Как улучшили EAGLE-3

Сегодняшняя статья — о EAGLE-3. Это новая версия популярного метода спекулятивного декодинга. Расскажем, как её улучшили по сравнению с прошлыми итерациями.

Спекулятивный декодинг — это способ ускорения инференса, предполагающий использование черновой (draft) модели, которая предлагает варианты продолжения генераций. Основная модель проверяет эти варианты, выбирая один с помощью процедуры верификации. Качество генераций при этом не страдает, ведь окончательное решение о принятии тех или иных токенов лежит на основной модели.

Один из самых известных методов спекулятивного декодинга — Extrapolation Algorithm for Greater Language-model Efficiency (EAGLE). В его рамках модель принимает не только прошлые токены, но и их feature-вектора. Это позволяет увеличить точность угадывания токенов.

Обновлённая версия EAGLE — EAGLE-3 — призвана сделать угадывание ещё более точным. Для этого можно просто налить больше данных в обучение EAGLE-модели. Однако, как показала практика, такой подход работает не слишком хорошо. Авторы метода посчитали, что здесь мешает feature loss, на который учится EAGLE. Выход — избавиться от feature loss и учить только на KL-лосс между предсказаниями EAGLE-головы и основной модели.

Проверка этой гипотезы показала, что без feature loss точность угадывания первого токена действительно увеличивается при добавлении новых данных. Однако она падает для следующих токенов. Всё из-за того, что теряется способность предсказывать в глубину. Решение: во время обучения делать не одну, а сразу несколько итераций EAGLE-головы, осуществляя предсказание в глубину.

Авторы сделали ещё одно улучшение. В прошлых версиях метода в EAGLE-модель подавали хиддены с последнего слоя таргет-модели, а также эмбеддинги токенов, отсэмплированных из них. Исследователи предположили, что в хидденах недостаточно информации, чтобы эффективно предсказывать токены. Вероятно, больше данных содержится в хидденах с промежуточных трансформерных слоёв. В EAGLE-3 авторы конкатенируют хиддены с трёх decoder-слоёв — с третего от начала, третьего от конца и слоя в середине между ними — и уже их передают на вход EAGLE-модели.

Суммируя, авторы EAGLE-3:

— убрали feature loss;
— добавили несколько шагов на обучении и увеличили объём данных;
— решили отправлять в EAGLE-модель хиддены с нескольких слоёв.

По сравнению с инференсом без использования EAGLE, всё это позволило получить прирост в скорости в 6,5 раза — и без потери качества. Число токенов за одну итерацию увеличилось на 50% по сравнению с EAGLE-2: с 4,05 до 6,13.

Разбор подготовил ❣ Алексей Гликин

Душный NLP