Как работает speculative decoding с несколькими draft моделями?

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

Speculative decoding с несколькими draft моделями — это техника ускорения инференса LLM, при которой несколько «черновиков» (draft моделей) параллельно генерируют кандидатные последовательности токенов, а одна большая target модель проверяет их и принимает или отвергает. Это позволяет повысить acceptance rate (долю принятых токенов) за счёт разнообразия гипотез, особенно в задачах, где одна draft модель плохо предсказывает сложные паттерны. Реализации делятся на ensemble-based (например, Medusa-2) и tree-based (например, EAGLE-3).

1. Термины и базовые понятия

Speculative decoding — метод ускорения генерации текста в LLM, при котором маленькая быстрая модель (draft model) генерирует несколько токенов за один шаг, а большая модель (target model) проверяет их и принимает или отвергает, пересчитывая логиты. Цель — уменьшить количество дорогих forward-проходов target модели.

Draft model — лёгкая модель (часто на 1–2 порядка меньше target), обученная предсказывать распределение target модели. Может быть отдельной сетью или головой (head) на общем backbone.

Target model — основная большая LLM, которую мы хотим ускорить. Её forward-проходы дороги, поэтому мы стараемся минимизировать их количество.

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

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

Ensemble — комбинация нескольких независимых draft моделей (или голов), каждая из которых выдаёт свой набор кандидатных токенов. Target модель проверяет все кандидаты сразу.

2. Базовая идея speculative decoding (одна draft модель)

Для понимания множественных draft моделей нужно вспомнить классический speculative decoding:

1. Draft модель генерирует K токенов авторегрессивно (быстро, так как она маленькая).
2. Target модель выполняет один forward-проход на всей последовательности из K+1 токенов (включая последний реальный).
3. Для каждой позиции сравниваются логиты draft и target. Если распределения совпадают (по критерию rejection sampling), токен принимается. Если нет — токен отвергается, и генерация откатывается до последнего принятого.
4. Процесс повторяется.

Ускорение достигается, если acceptance rate высокий (например, 0.8–0.9). Но одна draft модель может плохо предсказывать в некоторых контекстах (редкие слова, сложные синтаксические конструкции), что снижает acceptance rate.

3. Мотивация для нескольких draft моделей

Одна draft модель — это компромисс между скоростью и качеством предсказания. Если сделать draft модель слишком маленькой, acceptance rate падает. Если слишком большой — выигрыш в скорости исчезает.

Использование нескольких draft моделей позволяет:

• Покрыть разные режимы распределения target модели (каждая draft модель специализируется на своих паттернах).
• Параллельно генерировать несколько кандидатных последовательностей, увеличивая шанс, что хотя бы одна из них будет принята.
• Улучшить acceptance rate без увеличения размера одной draft модели.

4. Ensemble-based подход (Medusa-2)

Medusa (и её версия Medusa-2) — это метод, при котором на финальные скрытые состояния (hidden states) target модели добавляются несколько лёгких голов (heads). Каждая голова — это маленькая нейросеть (например, MLP), которая предсказывает распределение следующего токена, но с разными начальными условиями или архитектурой. Все головы работают параллельно на одном и том же скрытом состоянии.

Как работает с несколькими головами

1. Target модель вычисляет скрытое состояние для текущего токена.
2. N голов (draft моделей) параллельно предсказывают распределения для следующего токена.
3. Из каждой головы семплируется один или несколько кандидатных токенов (обычно top-k).
4. Все кандидаты собираются в пул (например, до M уникальных токенов).
5. Target модель делает один forward-проход для всех кандидатов сразу (с помощью расширения батча или tree attention).
6. Для каждого кандидата проверяется acceptance: если распределение target модели согласуется с распределением головы, токен принимается.
7. Принимается самый длинный префикс, согласованный с target.

Преимущества простая реализация, не требует отдельного обучения draft моделей (головы дообучаются вместе с target). Недостаток: головы могут быть скоррелированы, что снижает разнообразие.

5. Tree-based подход (EAGLE-3)

EAGLE (Extrapolation of Autoregressive Generation with Less Energy) — семейство методов, где draft модель строит дерево возможных продолжений. В EAGLE-3 используется несколько draft моделей (или одна draft модель с разными стратегиями семплирования), которые генерируют не одну последовательность, а дерево.

Как работает

1. Draft модель (или несколько) генерирует несколько вариантов следующего токена (например, top-3).
2. Для каждого варианта draft модель продолжает генерацию ещё на несколько шагов, создавая ветви.
3. Получается дерево глубиной D и шириной W (каждая ветвь — своя последовательность).
4. Target модель обрабатывает всё дерево за один forward-проход с помощью tree attention — модифицированного механизма, где каждая ветвь видит только свой путь (маскировка), но все ветви обрабатываются параллельно.
5. Target модель проверяет каждую ветвь и выбирает самую длинную принятую последовательность.

Роль нескольких draft моделей в EAGLE-3 draft модель может быть одна, но она генерирует несколько вариантов (top-k). Однако можно использовать несколько разных draft моделей, каждая из которых предлагает свои ветви. Это увеличивает разнообразие дерева.

Преимущества высокий acceptance rate за счёт перебора множества вариантов. Недостаток: сложность tree attention и рост вычислительных затрат при увеличении дерева.

6. Сравнение ensemble-based и tree-based

7. Acceptance rate и ускорение

Формула ускорения для speculative decoding с одной draft моделью:

Speedup = (1 + K * acceptance_rate) / (1 + K * (cost_draft / cost_target))

где K — число токенов, генерируемых draft моделью за шаг.

С несколькими draft моделями acceptance_rate повышается, так как target модель может выбрать лучший вариант из нескольких. Например, если одна draft модель даёт acceptance_rate 0.7, а две независимые — 0.85 (вероятность, что хотя бы одна из двух угадает). Однако нужно учитывать, что стоимость forward-прохода target модели возрастает (нужно обработать больше кандидатов). На практике выигрыш в acceptance_rate перевешивает рост стоимости, если количество кандидатов невелико (2–4).

Пример расчёта

• cost_target = 10x cost_draft (одна draft модель).
• K = 5.
• Одна draft: acceptance_rate = 0.7 → Speedup ≈ (1+50.7)/(1+50.1) = (4.5)/(1.5) = 3.0.
• Две draft (ensemble): acceptance_rate = 0.85, но cost_target увеличивается на 20% (из-за обработки двух кандидатов) → cost_draft удваивается? На самом деле draft модели работают параллельно, их стоимость суммируется. Пусть две draft модели стоят 2x cost_draft. Тогда Speedup = (1+50.85)/(1+50.2) = (5.25)/(2.0) = 2.625. В этом примере ускорение даже меньше, потому что рост стоимости draft превысил рост acceptance. Поэтому важно, чтобы draft модели были очень лёгкими (например, головы на общем backbone, где дополнительная стоимость минимальна).

8. Практические аспекты

• Выбор количества draft моделей обычно 2–4. Больше — растёт overhead и падает ускорение.
• Обучение draft моделей в Medusa головы дообучаются на небольшом датасете (10–100k примеров) с дистилляцией из target. В EAGLE draft модель обучается предсказывать скрытые состояния target.
• Интеграция с Agentic RAG speculative decoding может ускорять вызовы LLM в агентах (например, при генерации планов, ответов на запросы). В контексте Agentic RAG multiple draft models могут быть полезны, когда агент должен быстро сгенерировать несколько вариантов действий.
• Ограничения speculative decoding не меняет качество ответов (теоретически идентично обычной генерации), но требует дополнительной памяти для хранения draft моделей и модификации инференса.

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

Задача Реализовать простой speculative decoding с двумя draft моделями (маленькая и средняя) для генерации текста на небольшой LLM (например, GPT-2).

Инструменты Python, PyTorch, transformers (Hugging Face).

Шаги:

1. Выбрать target модель (например, distilgpt2) и две draft модели: distilgpt2 (та же, но с меньшим числом слоёв) и gpt2 (средняя). На самом деле для демонстрации можно использовать одну и ту же модель с разными dropout или разными seed.
2. Реализовать функцию generate_draft(model, prompt, K), которая генерирует K токенов авторегрессивно.
3. Реализовать verify(target_model, prompt, draft_tokens), которая делает один forward-проход target модели и возвращает принятые токены по правилу rejection sampling.
4. Модифицировать для двух draft моделей: каждая генерирует свою последовательность, затем target проверяет обе и выбирает самую длинную принятую.
5. Сравнить скорость (токенов в секунду) и acceptance rate для одной и двух draft моделей на тестовых промптах.

Ожидаемый результат Вы увидите, что с двумя draft моделями acceptance rate выше, но общая скорость может быть ниже из-за дополнительных вычислений. Для реального ускорения нужно оптимизировать параллелизм (например, использовать tree attention).

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

Навигация

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