Что такое EAGLE-3 и чем он отличается от стандартного speculative decoding?

Краткий тезис

EAGLE-3 — это третий этап эволюции семейства методов EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency), которые относятся к feature-aware speculative decoding. В отличие от стандартного speculative decoding, где draft model (model|маленькая модель-черновик) обучается независимо от target model (большой целевой модели), EAGLE-3 использует скрытые представления (hidden states) target модели для предсказания следующих токенов. Это позволяет достичь acceptance rate 78–82% против 50–65% у классических подходов. Ключевые улучшения: EAGLE-2 добавил динамическое построение дерева кандидатов, а EAGLE-3 — penalty|KL divergence penalty|KL divergence top-k penalty|KL divergence loss для более точного обучения draft модели.

1. Термин: Speculative Decoding (спекулятивная декодировка)

Speculative decoding — это техника ускорения инференса LLM без потери качества. Идея: маленькая и быстрая draft model генерирует несколько токенов (кандидатов), а большая target model проверяет их за один forward pass. Если кандидаты принимаются — мы сэкономили время, так как target model не пришлось генерировать каждый токен последовательно.

Почему это работает

• Draft model намного меньше (например, 100M параметров против 7B), поэтому её forward pass быстрее.
• Target model может проверить несколько токенов параллельно (благодаря masked attention).
• Acceptance rate — доля токенов, которые target model принимает. Чем выше acceptance rate, тем больше ускорение.

Ограничение стандартного подхода
Draft model обучается на текстовых данных независимо от target model. Она не знает, как target model «думает» (её скрытые состояния). Из-за этого acceptance rate ограничен (обычно 50–65%).

2. Термин: Feature-Aware Draft Model (модель-черновик, учитывающая признаки)

Feature-aware draft model — это draft модель, которая использует не только предыдущие токены, но и скрытые представления (hidden states) target модели для предсказания следующего токена. Это ключевое отличие EAGLE от классического speculative decoding.

Как это работает в EAGLE

1. На каждом шаге target модель вычисляет скрытое состояние для текущего токена.
2. Draft модель получает это скрытое состояние как дополнительный вход.
3. Draft модель предсказывает следующий токен, «зная», как target модель «воспринимает» контекст.
4. Target модель проверяет предсказание — если оно совпадает с её собственным распределением, токен принимается.

Преимущество
Draft модель учится имитировать распределение target модели, а не просто языковое распределение. Это даёт более высокий acceptance rate.

3. Эволюция EAGLE: от 1 до 3

Термин «Tree attention» (древовидное внимание):
Вместо генерации одной последовательности токенов draft модель генерирует несколько ветвей (дерево). Target модель проверяет все ветви параллельно, выбирая самую длинную принятую последовательность. Это увеличивает acceptance rate, так как даже если одна ветвь не принята, другая может быть принята.

Термин «Top-k KL divergence loss»:
Функция потерь для обучения draft модели. Вместо минимизации KL-дивергенции между полными распределениями draft и target моделей, минимизируется расхождение только для top-k наиболее вероятных токенов (по версии target модели). Это заставляет draft модель фокусироваться на важных токенах, игнорируя маловероятные.

4. Как работает EAGLE-3 (пошагово)

1. Инициализация
   Target модель (например, LLaMA-7B) загружена в память. Draft модель (например, 100M параметров) обучена с top-k KL divergence loss на скрытых состояниях target модели.

2. Генерация кандидатов

   • Target модель вычисляет скрытое состояние h_t для текущего токена x_t.
   • Draft модель получает h_t и предсказывает следующий токен x_{t+1}.
   • Draft модель также генерирует несколько альтернативных токенов (ветви дерева) с помощью tree attention.

3. Проверка target моделью

   • Target модель получает все кандидаты (дерево) и вычисляет для каждого токена вероятность p_target(x | context).
   • Если p_target(x_draft) >= p_target(x_target), токен принимается.
   • Иначе токен отклоняется, и target модель генерирует свой токен (сэмплирование из скорректированного распределения).

4. Принятие последовательности

   • Target модель принимает самую длинную префиксную последовательность из дерева, где все токены удовлетворяют условию принятия.
   • Оставшиеся токены генерируются target моделью последовательно (но их обычно мало, так как acceptance rate высок).

Пример кода (упрощённый псевдокод):

def eagle3_generate(target_model, draft_model, prompt, max_tokens=256):
    tokens = tokenize(prompt)
    hidden_state = target_model.get_hidden_state(tokens[-1])
    
    for _ in range(max_tokens):
        # Draft model генерирует дерево кандидатов
        candidates = draft_model.generate_tree(hidden_state, top_k=5)
        
        # Target model проверяет все кандидаты
        accepted, new_hidden = target_model.verify(candidates, tokens)
        
        # Добавляем принятые токены
        tokens.extend(accepted)
        
        # Обновляем скрытое состояние для следующего шага
        hidden_state = new_hidden
        
        if tokens[-1] == EOS:
            break
    
    return decode(tokens)

5. Сравнение EAGLE-3 со стандартным speculative decoding

Почему EAGLE-3 быстрее

• Выше acceptance rate → меньше шагов target model.
• Tree attention позволяет проверять больше кандидатов за один forward pass.
• Top-k KL divergence loss делает draft модель более точной в важных токенах.

6. Термин: Acceptance Rate (коэффициент принятия)

Acceptance rate — доля токенов, сгенерированных draft моделью, которые были приняты target моделью без изменений. Формально:

acceptance_rate = (количество принятых токенов draft) / (общее количество сгенерированных токенов draft)

Влияние на ускорение
Ускорение ≈ 1 / (1 - acceptance_rate) при условии, что draft модель значительно быстрее target. Например, при acceptance rate 80% ускорение ~5x, при 60% — ~2.5x.

Почему EAGLE-3 достигает 78–82%

• Draft модель «видит» скрытые состояния target, поэтому её предсказания лучше согласованы.
• Top-k KL divergence loss заставляет draft модель фокусироваться на токенах, которые target модель считает вероятными.
• Tree attention даёт несколько попыток: даже если одна ветвь не принята, другая может быть принята.

7. Термин: Top-k KL Divergence Loss

KL divergence (дивергенция Кульбака-Лейблера) — мера различия между двумя распределениями вероятностей. В контексте speculative decoding, KL(P_target || P_draft) показывает, насколько распределение draft модели отличается от распределения target модели.

Стандартный подход
Минимизировать KL(P_target || P_draft) по всем токенам словаря (обычно 50k+ токенов). Это дорого и неэффективно, так как большинство токенов имеют очень низкую вероятность.

Top-k KL divergence loss (EAGLE-3):

• Выбираются top-k токенов по вероятности target модели (например, k=10).
• KL-divergence считается только по этим токенам.
• Остальные токены игнорируются (или штрафуются минимально).

Преимущества

• Обучение быстрее (меньше вычислений).
• Draft модель учится точно предсказывать только важные токены.
• Acceptance rate растёт, так как draft модель не тратит ресурсы на маловероятные токены.

Формула (упрощённо):

L = Σ_{i in top-k} P_target(x_i) * log(P_target(x_i) / P_draft(x_i))

8. Когда использовать EAGLE-3

EAGLE-3 эффективен, когда:

• Target модель очень большая (7B+ параметров) и каждый forward pass дорог.
• Нужно ускорение в 3–5 раз без потери качества.
• Есть доступ к hidden states target модели (не все API предоставляют).
• Draft модель можно обучить на скрытых состояниях (требуется вычислительные ресурсы).

Когда стандартный speculative decoding может быть достаточен:

• Target модель небольшая (1–3B параметров) — ускорение от EAGLE-3 невелико.
• Нет доступа к hidden states (например, через API).
• Acceptance rate 50–65% даёт приемлемое ускорение (2–3x).

Ограничения EAGLE-3

• Требуется модификация инференс-пайплайна (доступ к hidden states).
• Обучение draft model сложнее (нужны данные с hidden states).
• Не все архитектуры поддерживают tree attention (требуется реализация).

9. Пет-проект для закрепления

Задача Реализовать упрощённую версию EAGLE-3 для небольшой target модели (например, GPT-2) и сравнить со стандартным speculative decoding.

Инструменты

• Python, PyTorch, Transformers (Hugging Face).
• GPT-2 (124M) как target model, маленькая GPT-2 (12M) как draft model.
• Датасет: WikiText-2 или любой текстовый корпус.

Шаги:

1. Подготовка данных

   • Взять 1000 предложений из датасета.
   • Для каждого предложения получить hidden states target модели (последний слой).

2. Обучение draft model (EAGLE-3):

   • Draft model принимает на вход: предыдущий токен + hidden state target.
   • Функция потерь: top-k KL divergence (k=5).
   • Обучить 10 эпох.

3. Обучение draft model (стандартный подход):

   • Draft model принимает только предыдущий токен.
   • Функция потерь: cross-entropy на тексте.
   • Обучить 10 эпох.

4. Реализация speculative decoding

   • Написать функцию, которая генерирует текст с помощью draft + target.
   • Для EAGLE-3: передавать hidden states target.
   • Для стандартного: только токены.

5. Сравнение

   • Замерить acceptance rate на 100 запросах.
   • Замерить время генерации 100 токенов.
   • Проверить, что финальный текст идентичен (или почти идентичен) тексту, сгенерированному только target моделью.

Ожидаемый результат

• EAGLE-3 покажет acceptance rate ~70–75% (против 50–60% у стандартного).
• Ускорение: 2–3x для EAGLE-3, 1.5–2x для стандартного.
• Качество текста не ухудшится (perplexity совпадает).

10. Связь с другими вопросами

11. Навигация

• Предыдущий: 157
• Следующий: 159
• Индекс: 00. Индекс разборов

Навигация

• Предыдущий: 157
• Следующий: 159
• Индекс: 00. Индекс разборов