Технический анализ EAGLE-4

В сотрудничестве с Fireworks и fal мы публикуем технический анализ EAGLE-4 вместе с публичным набором данных для бенчмаркинга спекулятивного декодирования в творческих задачах.

Зачем нужно спекулятивное декодирование?

LLM генерируют текст авторегрессионно, предсказывая по одному токену за раз. При малых размерах батча декодирование опирается на «тонкие» матричные умножения, которые недоиспользуют вычислительную мощность GPU из-за ограничений пропускной способности памяти. Батчинг запросов от нескольких пользователей превращает эти операции в более крупные матричные умножения и повышает общую пропускную способность, но он может снизить скорость генерации токенов, ощущаемую каждым пользователем.

Спекулятивное декодирование решает эту задачу обеспечения более быстрого ответа за счет использования легковесной модели для создания черновика из нескольких токенов, которые целевая модель затем проверяет одновременно. Проверка нескольких токенов увеличивает объем вычислений, выполняемых при каждом чтении весов модели, что позволяет эффективнее задействовать простаивающие в противном случае вычислительные мощности GPU. При правильной процедуре проверки спекулятивное декодирование сохраняет выходное распределение целевой модели.

Семейство EAGLE — это один из подходов к спекулятивному декодированию, используемый в продакшене. Его легковесный трансформер-генератор черновиков принимает скрытые состояния и эмбеддинги токенов от целевой модели и обучен предсказывать выходные токены целевой модели. EAGLE-3 использует скрытые состояния из трех слоев на разной глубине в сочетании с методами обучения, призванными увеличить длину принятия (acceptance length).

За год с момента выпуска EAGLE-3 появились и неавторегрессионные подходы, такие как DFlash и DSpark. Эти методы используют маскированную диффузию для параллельного создания черновика из нескольких токенов. Это может сделать более крупные генераторы черновиков практичными, поскольку составление черновика не требует отдельного последовательного шага для каждого предсказанного токена.

Почему моделям изображений и видео требуется быстрое декодирование

Многие современные системы генерации изображений и видео обрабатывают пользовательский промпт перед тем, как передать его в генеративную модель. Этот шаг, называемый «улучшением промпта» (prompt enhancement) или «расширением промпта» (prompt expansion), превращает краткий запрос в более подробное описание сцены. Он также адаптирует промпт под ожидаемый формат целевой модели, помогая получать более привлекательные результаты.

По мере того как модели изображений и видео становятся быстрее, улучшение промпта может составлять все большую долю сквозной задержки (end-to-end latency). Например, видеомодели могут генерировать видео в реальном времени, но этап расширения без каких-либо оптимизаций занял бы больше времени, чем сама генерация видео. Поэтому таким системам необходимы быстрые ответы при малых размерах батча, чтобы взаимодействие оставалось интерактивным, что делает улучшение промпта отличным сценарием применения для спекулятивного декодирования.

Чтобы привнести эти творческие рабочие нагрузки в оценку спекулятивного декодирования, мы выпускаем Prompt Enhancement SpecBench. Он содержит 400 примеров из 18 уникальных наборов данных в четырех категориях: текст-видео, изображение-видео, текст-изображение и изображение-изображение. Каждая категория включает 100 примеров.

Архитектура EAGLE-4

EAGLE-4 опирается на подход EAGLE-3: легковесный генератор черновиков объединяет скрытые состояния целевой модели с эмбеддингами токенов для предсказания возможных продолжений. Его эффективность зависит как от качества этих предсказаний, так и от того, насколько дешево он может их генерировать.

EAGLE-4 вносит два изменения в этот путь генерации черновиков. Внимание на основе скользящего окна (sliding-window attention) ограничивает контекст, используемый слоем внимания генератора. Голова кандидатов с низким рангом (low-rank candidate head) сначала оценивает словарь в меньшем представлении, а затем вычисляет полные логиты только для выбранных кандидатов. Эти изменения направлены на снижение стоимости внимания и прогнозирования следующего токена соответственно.

На диаграмме ниже показана более общая архитектура, включая опциональные компоненты. Несколько слоев черновика, поддержка смешанных экспертов (MoE) и Latent ReLU еще не включены в выпущенную модель. Мы ожидаем, что будущие модели будут использовать эти функции.

Деревья проверки EAGLE-4 и утилита управления (harness)

Архитектура влияет на предсказание генератора; дерево проверки расширяет список предсказаний, которые проверяет целевая модель. Единая цепочка фиксирует одно продолжение в каждой позиции, поэтому раннее несовпадение препятствует принятию последующих токенов в этой цепочке. Ветвление предоставляет верификатору альтернативные варианты продолжения для рассмотрения, что полезно, когда правдоподобно выглядят несколько следующих токенов.

EAGLE представляет эти альтернативы в виде дерева и оценивает выбранные узлы за один прямой проход целевой модели. Пример ниже показывает, как проверка следует по принятому пути, а затем продолжает работу с токеном от целевой модели, когда черновик больше не совпадает.

Чтобы оценивать эти ветви вместе, целевой модели необходимо сохранять их контексты раздельными. Пользовательская маска внимания позволяет каждому токену обращаться к общему префиксу, своим предкам и к самому себе, скрывая при этом токены-братья и другие ветви. Это сохраняет контекст каждого кандидатного продолжения в рамках комбинированного прямого прохода.

Более крупное дерево предлагает больше альтернатив, но его создание и проверка также стоят дороже. Утилита управления EAGLE-4 управляет этим компромиссом, используя измерения задержки на GPU. Она расширяет дерево уровень за уровнем, решает, когда прекратить создание черновика, и выбирает кандидатные узлы для проверки на основе стоимости проверки. После проверки она сохраняет принятое продолжение, освобождает отклоненные записи KV-кэша и начинает следующий раунд.

Перед обслуживанием единовременная автономная калибровка на целевом типе GPU и параллелизме создает таблицу стоимости: задержка проверки в зависимости от размера батча, длины контекста и количества кандидатов. На каждом уровне черновика утилита оценивает прирост кумулятивной вероятности, деленный на измеренную предельную стоимость на токен. Она расширяет или углубляет дерево, пока это соотношение остается выше порогового значения, после чего останавливается. Топовые кандидаты проверяются за один прямой проход целевой модели с маской внимания дерева, а выборка с отклонением (rejection sampling) принимает самый длинный совпадающий префикс, поэтому выходное распределение остается неизменным. В результате получается широкое глубокое дерево при размере батча 1 или длинном контексте, либо узкое дерево или простая цепочка при большом батче, используя один генератор без фиксированной глубины или настроек top-k.

Бенчмарки

Все эксперименты запускались на GPU NVIDIA B300 с параллелизмом тензоров, равным 4 (TP=4).

Наборы данных для оценки

Мы оценили EAGLE-4 на двух тщательно отобранных наборах данных:

• Набор данных расширения промптов fal: 400 промптов в категориях текст-видео, текст-изображение, изображение-видео и изображение-изображение, по 100 промптов на категорию.
• Набор данных спекулятивного декодирования Fireworks: chat-30, code-30, stem-30, math-30 и coding-long-20. Из 20 длинных промптов по программированию 18 оказались пригодными к использованию, что в сумме составило 138 промптов.

Приведенные ниже сравнения цепочки и дерева используют все 138 промптов Fireworks и 200 текстовых промптов fal. Мы приводим длину принятия (AL) для анализа поведения генератора, за которой следует пропускная способность декодирования для демонстрации влияния на время выполнения.

Длина принятия и сложность рабочей нагрузки

Творческое расширение промптов включает в себя продолжения с открытым концом, которые генератору может быть сложно предсказать. Анализ энтропии ниже показывает меньшую длину принятия для промптов в группах с более высокой энтропией: 2,66 в квартиле с самой высокой энтропией по сравнению с 2,97 в квартиле с самой низкой энтропией.

Прирост дерева по задачам

При конфигурациях, показанных ниже, дерево увеличивает длину принятия для каждого разделения задач. Совокупный относительный прирост составляет 9,6% на датасете Fireworks и 12,4% на текстовом подмножестве fal.

Масштабирование бюджета верификации

При каждом протестированном бюджете верификации оба метода имеют меньшую длину принятия на творческих промптах fal, чем на датасете Fireworks, как и ожидалось. Эти рабочие нагрузки сложнее для генераторов в данной оценке, а больший бюджет верификации улучшает принятие на обоих датасетах.

Преимущество дерева растет по мере увеличения бюджета верификации. Оба метода улучшаются, но у цепочки наблюдается убывающая отдача, особенно на датасете fal, в то время как дерево продолжает наращивать длину принятия.

При бюджете верификации в 16 токенов, самом большом из протестированных, дерево увеличивает длину принятия на 21,1% на датасете Fireworks и на 26,1% на текстовом подмножестве fal по сравнению с цепочкой.

Где дерево получает прирост принятия

Частота попаданий по позициям помогает объяснить эти приросты. На следующем графике сравнивается дерево при L=8 с цепочкой при N=7 на датасете Fireworks. Каждое значение представляет собой вероятность того, что сгенерированный токен на этой позиции будет принят вдоль пути.

Преимущество дерева сосредоточено в ранних позициях. На датасете Fireworks частота попаданий для первого токена возрастает с 72,5% до 85,3%, что составляет увеличение на 12,8 процентных пункта. Исследование альтернативных продолжений ближе к началу дерева является вероятным объяснением.

При фиксированном бюджете верификации ветвление ближе к началу может улучшить общую длину принятия, даже если принимается меньше глубоких токенов. Длинная цепочка тратит большую часть своего бюджета на более поздние позиции, которые с меньшей вероятностью будут достигнуты.

Скорость декодирования

Длина принятия помогает объяснить поведение генератора, в то время как скорость декодирования отражает совокупную стоимость генерации и верификации. Мы сравниваем DSpark, DFlash2, статическую цепочку EAGLE-4 и EAGLE-4 Harness при L=8 на Fireworks Bench и Prompt Enhancement SpecBench, при C=1 и C=8.

EAGLE-4 Harness имеет самую высокую скорость во всех четырех сравнениях. При C=1 она достигает 660,6 токенов в секунду на Fireworks Bench и 622,9 на Specbench, что в 1,63 и 1,85 раза превышает скорость самой быстрой альтернативы в каждом бенчмарке соответственно. При C=8 скорость достигает 3566,4 и 3551,9 токенов в секунду с соответствующим преимуществом в 1,36 и 1,62 раза.

Преимущество больше на Specbench при обеих настройках. При C=8 Harness достигает практически такой же скорости на Specbench, как и на Fireworks Bench, в то время как другие методы показывают больший разрыв между этими двумя рабочими нагрузками.

Увеличение бюджета верификации дерева

В предыдущем эксперименте со скоростью также тестировалось значение L=16. В этом сравнении увеличение бюджета дерева с L=8 до L=16 увеличило скорость с 660,1 до 752 токенов в секунду с расчетным ускорением в 3,77 раза по сравнению с декодированием без спекулятивной генерации.

Заключение

EAGLE-4 показывает, как эффективный генератор, верификация с ветвлением и обвязка (harness), учитывающая затраты на аппаратное обеспечение, могут работать вместе для повышения скорости декодирования. В нашей оценке Harness лидирует в обоих бенчмарках при C=1 и C=8. Его преимущество особенно сильно при улучшении творческих промптов: на SpecBench он обеспечивает в 1,85 раза большую скорость по сравнению с самой быстрой альтернативой при C=1 и в 1,62 раза при C=8.

Творческие рабочие нагрузки по-прежнему сложнее предсказывать для генераторов, но дерево восстанавливает больше полезных продолжений из этой неопределенности. По мере того как генерация изображений и видео становится быстрее, сокращение времени, затрачиваемого на улучшение промптов, приобретает все большую ценность. Мы выпускаем Prompt Enhancement SpecBench, чтобы сделать эти рабочие нагрузки регулярной частью оценки спекулятивного декодирования.

Спасибо нашим партнерам из Fireworks за совместную работу над этим анализом и оценкой.