Speculative Decoding: от идеи до SOTA

1. Почему вообще декодинг медленный

Autoregressive-инференс LLM упирается не в compute, а в memory bandwidth. Чтобы сгенерировать один токен, нужно прочитать все веса модели из HBM в SRAM GPU. Арифметическая интенсивность при batch=1 крайне низкая — тензорные ядра простаивают.

Ключевое наблюдение: forward pass на $K$ токенах стоит почти столько же по времени, сколько на $1$ токене (пока $K$ невелико), потому что bottleneck — загрузка весов, а не вычисления. Это и эксплуатирует speculative decoding.

2. Базовая идея (Leviathan et al. 2023, Chen et al. 2023)

Берём две модели:

• Target модель $M_p$ с распределением $p(x)$ — большая, медленная, качественная.
• Draft модель $M_q$ с распределением $q(x)$ — маленькая, быстрая, чуть менее точная.

Цикл генерации:

1. Draft модель авторегрессивно генерирует $\gamma$ токенов: $x_1,x_2,\dots ,x_{\gamma }$.
2. Target модель за один forward pass считает $p(\cdot \mid x_{<i})$ для всех $i=1,\dots ,\gamma +1$.
3. Слева направо принимаем/отклоняем токены по правилу rejection sampling.
4. Если все $\gamma$ приняты — бонусом сэмплируем $(\gamma +1)$-й токен из $p$.

3. Математика: почему это корректно

Хотим, чтобы итоговое распределение принятого токена было ровно $p(x)$.

Acceptance rule. Для $x\sim q(\cdot )$ принимаем с вероятностью

$$\alpha (x)=\min \left(1,\frac{p(x)}{q(x)}\right)$$

Resampling rule. Если отклонили, сэмплируем из скорректированного распределения

$$p^{\prime }(x)=\frac{\max (0,p(x)-q(x))}{{\sum }_{x^{\prime }}\max (0,p(x^{\prime })-q(x^{\prime }))}$$

Теорема. Итоговое распределение равно $p(x)$.

Доказательство.

$$P(\text{output}=x)=q(x)\alpha (x)+(1-\beta )p^{\prime }(x),\beta =\sum _{x}q(x)\alpha (x)$$

Подставляем:

$$q(x)\cdot \min \left(1,\frac{p(x)}{q(x)}\right)=\min (p(x),q(x))$$

$$(1-\beta )p^{\prime }(x)=\max (0,p(x)-q(x))$$

Складываем:

$$\min (p(x),q(x))+\max (0,p(x)-q(x))=p(x).■$$

Главный вывод: декодинг бесплатный по качеству — никакой деградации распределения.

4. Ожидаемое ускорение

Пусть:

• $\alpha$ — средняя вероятность принятия токена,
• $c=T_q/T_p$ — отношение латентности draft/target,
• $\gamma$ — число draft-токенов на цикл.

Ожидаемое число сгенерированных токенов за цикл (геометрическая сумма):

$$\mathbb{E}[\#\text{tokens}]=\frac{1-{\alpha }^{\gamma +1}}{1-\alpha }$$

Время цикла: $\gamma T_q+T_p$. Ускорение относительно ванильного декодинга:

$$\text{Speedup}(\alpha ,c,\gamma )=\frac{1-{\alpha }^{\gamma +1}}{(1-\alpha )(c\gamma +1)}$$

Оптимальное $\gamma$ обычно лежит в диапазоне 4–8; при $\alpha \approx 0.7$ и $c\approx 0.05$ получаем 2.5–3x.

5. Откуда брать draft-модель

Подход	Суть
Готовая малая модель	Llama-7B для Llama-70B, Qwen-0.5B для Qwen-72B и т.п.
Дистилляция	Тренируем draft под распределение target
Self-speculative	Target сама себе draft через layer-skip / early exit (LayerSkip, Draft&Verify)
N-gram / Prompt Lookup	Достаём n-грамм-кандидатов прямо из контекста — отлично для кода и шаблонов
Retrieval (REST)	Достаём кандидатов из внешнего корпуса по префиксу

6. Tree-based верификация (SpecInfer, Medusa)

Вместо одной цепочки draft-токенов строим дерево кандидатов и верифицируем все ветви одновременно через специальную tree attention mask. За один target forward pass проверяем сразу много гипотез — ожидаемое число принятых токенов растёт.

Medusa (Cai et al. 2024) — приклеивает к target модели $K$ дополнительных голов, каждая предсказывает токен на позиции $t+k$ напрямую из текущего hidden state $h_t$:

$${\hat{p}}_k(x_{t+k})=\mathrm{softmax}(W_k\cdot (h_t+\mathrm{SiLU}(W_k^{\prime }{h}_t)))$$

Топ-кандидаты от каждой головы образуют дерево, проверяемое одним forward pass. Отдельная draft модель не нужна — только дообучить головы.

7. EAGLE — SOTA-семейство (2024)

Главная идея EAGLE: предсказывать не токены, а признаки (hidden states предпоследнего слоя). Feature space гораздо регулярнее token space и лучше моделируется маленькой сетью.

Архитектура: 1-слойный autoregressive transformer-head, который на вход получает $[f_t,e_{t+1}]$ — feature на шаге $t$ и embedding следующего токена — и предсказывает $f_{t+1}$.

Loss:

$$\mathcal{L}={\mathcal{L}}_{\text{feat}}+w_{\text{tok}}\cdot {\mathcal{L}}_{\text{tok}}$$

$${\mathcal{L}}_{\text{feat}}=SmoothL1(f_{t+1},{\hat{f}}_{t+1}),{\mathcal{L}}_{\text{tok}}=\mathrm{CE}(p_{t+1},LM\_head({\hat{f}}_{t+1}))$$

• EAGLE-1: статическое дерево драфтов, 3x speedup.
• EAGLE-2: динамическое дерево по предсказанной confidence — больше принятых токенов при том же бюджете.
• EAGLE-3: убрана feature-regression регуляризация (она ограничивала draft), добавлен training-time test — моделирование ошибок драфта на обучении. До 5x на современных бенчмарках.

8. Lookahead Decoding

Другой подход — без draft-модели вообще. Использует Jacobi iteration для параллельного решения системы уравнений авторегрессии и поддерживает n-gram pool из истории — совпадения подставляются как draft-кандидаты. Удобно, когда тренировать draft неоткуда.

9. Что выбрать на практике

Сценарий	Метод
Есть малая модель того же семейства	Vanilla speculative decoding
Plug-and-play поверх готовой модели	Medusa / EAGLE (тренируем только головы)
Код, JSON, повторяющиеся шаблоны	Prompt Lookup Decoding
Production-serving (vLLM, SGLang)	EAGLE-2 / EAGLE-3
Очень длинный контекст	Self-speculative с layer-skip

10. Подводные камни

• Температура. При высокой $T$ распределение размывается, $\alpha$ падает — speedup тает.
• Batched serving. При больших batch size target уже compute-bound, выигрыш существенно меньше.
• Tree attention. Требует поддержки custom масок и position ids — не везде из коробки.
• Потолок acceptance. Даже идеальный драфт ограничен энтропией $p$ — для творческих задач $\alpha$ редко выше 0.8.
• KV-cache. Отклонённые токены требуют отката KV-кэша target модели — реализация неаккуратно может съесть весь выигрыш.

11. TL;DR

Speculative decoding — это rejection sampling над авторегрессивной генерацией, эксплуатирующее тот факт, что target forward pass на K токенах почти так же дёшев, как на 1. Корректность доказывается одной строчкой ($\min +\max =p$), ускорение задаётся формулой $(1-{\alpha }^{\gamma +1})/((1-\alpha )(c\gamma +1))$. Современная линия развития — feature-level prediction + динамические деревья (EAGLE-2/3), дающая 3–5x без потери качества.