Многотокенное предсказание

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

Идея многотокенного предсказания (multi-token prediction, MTP) состоит в том, чтобы заставить модель или связанный с ней легкий модуль предсказывать несколько будущих токенов за один цикл. Но финальное право выбора остается за исходной авторегрессионной моделью. Поэтому MTP обычно рассматривают не как замену авторегрессии, а как частный случай спекулятивного декодирования: быстрый черновик предлагает продолжение, а целевая модель проверяет его параллельно.

Спекулятивное декодирование

Пусть есть целевая модель $p_{\theta }$, которую мы хотим ускорить, и черновая модель $q_{\varphi }$, которая дешевле генерирует кандидатов. Один цикл спекулятивного декодирования выглядит так:

1. Черновая модель авторегрессионно или параллельно предлагает блок из $\gamma$ токенов:

$${\hat{x}}_{t+1},{\hat{x}}_{t+2},\dots ,{\hat{x}}_{t+\gamma }.$$

1. Целевая модель одним forward pass считает распределения

$$p_{\theta }(\cdot \mid x_{\le t}),p_{\theta }(\cdot \mid x_{\le t},{\hat{x}}_{t+1}),\dots ,p_{\theta }(\cdot \mid x_{\le t},{\hat{x}}_{t+1:t+\gamma })$$

для всех проверяемых позиций.

1. Токены проверяются слева направо. В greedy-режиме токен принимается, если он совпал с токеном, который выбрала бы целевая модель. В sampling-режиме используется rejection sampling:

$${\alpha }_i=\min \left(1,\frac{p_{\theta }({\hat{x}}_{t+i}\mid x_{\le t},{\hat{x}}_{t+1:t+i-1})}{q_{\varphi }({\hat{x}}_{t+i}\mid x_{\le t},{\hat{x}}_{t+1:t+i-1})}\right).$$

1. Если токен отклонен, целевая модель исправляет продолжение: новый токен сэмплируется из остаточного распределения, пропорционального положительной части $p_{\theta }-q_{\varphi }$. После этого текущий спекулятивный блок заканчивается, потому что все следующие черновые токены были построены из уже неправильного префикса.

2. Если все $\gamma$ токенов приняты, целевая модель дополнительно выдает бонусный токен на позиции $t+\gamma +1$.

Главное свойство: при корректной процедуре принятия/отклонения итоговое распределение совпадает с распределением целевой модели. Черновая модель влияет на скорость, но не должна менять качество. Она может ошибаться; целевая модель все равно имеет последнее слово.

Более подробно базовую схему спекулятивного декодирования можно разобрать в отдельной заметке.

Почему ускорение вообще возможно

Обозначим:

• $T_{\text{base}}$ — среднюю задержку одного обычного авторегрессионного шага целевой модели;
• $T_{\text{draft}}$ — время генерации чернового блока;
• $T_{\text{verify}}$ — время одного параллельного прогона целевой модели для проверки блока;
• $A$ — число принятых черновых токенов в цикле;
• $\tau =\mathbb{E}[A+1]$ — среднюю acceptance length с учетом бонусного токена целевой модели.

Тогда средняя задержка на один итоговый токен примерно равна

$$T_{\text{token}}\approx \frac{T_{\text{draft}}+T_{\text{verify}}}{\tau },$$

а ускорение относительно обычной авторегрессии:

$$S\approx \frac{T_{\text{base}}}{T_{\text{token}}}=\frac{T_{\text{base}}\tau }{T_{\text{draft}}+T_{\text{verify}}}.$$

Из этой формулы видно, что спекулятивное декодирование выигрывает только при двух условиях:

• черновик достаточно дешевый;
• целевая модель часто принимает черновые токены, то есть $\tau$ заметно больше 1.

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

MTP как черновик для самой модели

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

Если базовая модель на позиции $t$ строит скрытое состояние

$$h_t=f_{\theta }(x_{\le t}),$$

то стандартная LM-голова предсказывает следующий токен:

$$p_{\theta }(x_{t+1}\mid x_{\le t})=\mathrm{softmax}(W{h}_t).$$

MTP добавляет $K$ дополнительных голов:

$$q_{\varphi ,k}(x_{t+k}\mid x_{\le t})=\mathrm{softmax}(g_{\varphi ,k}(h_t)),k=1,\dots ,K.$$

Интуитивно каждая голова смотрит на внутреннее состояние модели и пытается угадать не только ближайший токен, но и более далекое продолжение. Одна голова предсказывает токен на один шаг вперед, другая — на два, третья — на три и так далее. Модель как будто учится думать немного вперед, но проверка все равно остается авторегрессионной.

Практический пример такого подхода — MTP-drafters для Gemma 4. Google описывает их как специализированные черновые модули для семейства Gemma 4: drafter предлагает несколько будущих токенов, а основная Gemma 4 проверяет их параллельно. По официальному посту Google от 5 мая 2026 года, такие drafters дают ускорение до 3x без деградации качества, потому что итоговую верификацию сохраняет основная модель: Accelerating Gemma 4: faster inference with multi-token prediction drafters.

Функция потерь для MTP

Есть два близких режима обучения.

Первый режим — обучать MTP вместе с основной языковой моделью как дополнительную задачу. Для последовательности $x_1,\dots ,x_T$ стандартная функция потерь предсказания следующего токена:

$${\mathcal{L}}_{\text{NTP}}=-\sum _{t=1}^{T-1}\log p_{\theta }(x_{t+1}\mid x_{\le t}).$$

MTP добавляет к функции потерь предсказание нескольких будущих токенов:

$${\mathcal{L}}_{\text{MTP}}=-\sum _{t=1}^T\sum _{k=1}^K{\lambda }_k{1}_{t+k\le T}\log q_{\varphi ,k}(x_{t+k}\mid x_{\le t}).$$

Итоговая функция потерь:

$$\mathcal{L}={\mathcal{L}}_{\text{NTP}}+\beta {\mathcal{L}}_{\text{MTP}}.$$

Такой подход использовался как обучающая цель в работе Better & Faster Large Language Models via Multi-token Prediction: модель предсказывает несколько будущих токенов через независимые головы поверх общего ствола сети.

Второй режим — заморозить целевую модель и дообучить только MTP-drafter. Это ближе к ускорению инференса уже готовой модели:

$$\theta \text{заморожены},\varphi \text{обучаются}.$$

Тогда оптимизируется только

$${\mathcal{L}}_{\text{draft}}=-\sum _{t=1}^T\sum _{k=1}^K{\lambda }_k{1}_{t+k\le T}\log q_{\varphi ,k}(x_{t+k}\mid h_t).$$

Веса ${\lambda }_k$ часто делают убывающими по $k$, например ${\lambda }_k={\rho }^{k-1}$ при $0<\rho \le 1$. Причина простая: ошибка на первом черновом токене обнуляет пользу всех следующих токенов в блоке. Поэтому ближайшие позиции важнее дальних.

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

Почему hidden states помогают предсказывать будущее

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

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

Формально скрытое состояние можно воспринимать как сжатую статистику контекста:

$$h_t\approx s(x_{\le t}),$$

где $s$ содержит не только информацию для ближайшего токена, но и признаки, полезные для распределений

$$p(x_{t+1}\mid x_{\le t}),p(x_{t+2}\mid x_{\le t}),\dots ,p(x_{t+K}\mid x_{\le t}).$$

MTP-голова не изобретает новое рассуждение с нуля. Она читает уже посчитанное состояние целевой модели и проецирует его в пространство будущих токенов. Поэтому MTP особенно естественно работает как самоспекуляция: drafter и target модель имеют общую внутреннюю геометрию, общий tokenizer и похожие предпочтения.

Простая MTP-голова

Схематически одну residual MTP-голову можно записать так:

import torch
from torch import nn


class ResidualMTPHead(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        vocab_size: int,
        expansion: int = 4,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        inner_size = expansion * hidden_size

        self.net = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, inner_size),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Linear(inner_size, hidden_size),
            nn.LayerNorm(hidden_size),
        )
        self.proj = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        future_h = h + self.net(h)
        return self.proj(future_h)

Смысл блоков:

• вход $h_t$ — скрытое состояние, которое базовая модель уже посчитала на позиции $t$
• residual MLP слегка поворачивает $h_t$ в подпространство будущего токена
• проекционный слой выдает logits по словарю
• отдельная голова $g_{\varphi ,k}$ отвечает за свой горизонт $k$

В production-реализациях такая голова может быть сложнее: с привязкой к embedding/LM-head целевой модели, с conditioning на уже предложенные черновые токены, с древовидным декодированием или с общей параметризацией нескольких горизонтов. Но базовая идея остается той же: использовать уже вычисленные скрытые состояния как дешевый источник информации о будущем.

Как MTP используется при декодировании

Для greedy-декодирования минимальная схема такая:

1. Целевая модель получает текущий префикс $x_{\le t}$ и считает $h_t$
2. MTP-головы предлагают черновой блок:

$${\hat{x}}_{t+1},\dots ,{\hat{x}}_{t+K}.$$

1. Целевая модель проверяет весь блок одним forward pass
2. Слева направо принимаются совпавшие токены
3. При первом несовпадении целевая модель выдает свой токен, а следующий цикл начинается уже от исправленного префикса

Пример:

префикс:
  "Actions speak louder than"

черновые предположения:
  ["words", ".", " The"]

проверка целевой модели:
  ["words", ".", " It"]

принято:
  ["words", "."]

исправление целевой модели:
  " It"

новый префикс:
  "Actions speak louder than words. It"

Черновик предложил три токена, но целевая модель приняла только первые два. Третий токен заменен, потому что после этого места все последующие черновые предположения уже построены из неправильного префикса.

Метрики качества черновой модели

Главная метрика — acceptance length:

$$\tau =\mathbb{E}[A+1],$$

где $A$ — число принятых черновых токенов, а $+1$ — бонусный токен целевой модели, если блок принят до конца или после исправления. Часто также считают:

$$acceptance\; rate=\frac{\text{число принятых черновых токенов}}{\text{число предложенных черновых токенов}}.$$

Фактический потолок достигается, когда все предложенные черновые токены принимаются целевой моделью. Тогда один дорогой verification pass дает сразу $\gamma +1$ итоговых токенов. На практике вероятность полного принятия падает с длиной блока, поэтому возникает компромисс:

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

DFlash

MTP ускоряет генерацию, но многие подходы с черновыми предсказаниями все еще генерируют кандидатов последовательно или зависят от небольшого числа голов будущего предсказания. DFlash предлагает другой вариант: использовать легкую block diffusion model как черновую модель.

Идея DFlash описана в статье DFlash: Block Diffusion for Flash Speculative Decoding и на странице проекта. Вместо того чтобы генерировать черновые токены один за другим, DFlash генерирует целый блок токенов параллельно. Целевая модель затем проверяет этот блок через спекулятивное декодирование.

У DFlash три ключевых компонента:

• Признаки контекста целевой модели. После prefill или verification целевая модель отдает скрытые состояния с нескольких слоев. Эти признаки сжимаются легкой проекцией.
• Инъекция в KV-cache. Сжатые признаки целевой модели инжектируются в Key/Value cache каждого чернового слоя. Это делает conditioning устойчивым на глубине: черновая модель не теряет информацию целевой модели после первого слоя.
• Блочное диффузионное черновое декодирование. Черновая модель предсказывает следующий блок токенов как masked/denoising задачу, параллельно по позициям блока.

В обозначениях:

$$C_t=F(h_t^{({\ell }_1)},h_t^{({\ell }_2)},\dots ,h_t^{({\ell }_m)})$$

— объединенный контекст из нескольких слоев целевой модели. DFlash-drafter задает распределение блока:

$$q_{\varphi }(x_{t+1:t+B}\mid x_{\le t},C_t).$$

В отличие от авторегрессионного черновика, стоимость генерации блока не обязана линейно расти с $B$:

$$T_{\text{draft}}^{AR}(B)\approx B\cdot T_{\text{step}},$$

а для block diffusion:

$$T_{\text{draft}}^{diff}(B)\approx T_{\text{block}}.$$

Именно это меняет дизайн-пространство. Если черновой блок можно получить одним параллельным прогоном, drafter может быть глубже и выразительнее без такого же роста latency, как у авторегрессионного черновика.

По результатам авторов DFlash, метод дает ускорение без изменения итогового распределения свыше 6x на ряде моделей и задач и до 2.5x больший speedup по сравнению с EAGLE-3. Готовые модели собраны в коллекции z-lab/dflash на Hugging Face, а код доступен в репозитории z-lab/dflash.

Обучение DFlash

DFlash обучается как черновой адаптер к замороженной целевой модели. Для последовательности "запрос + ответ":

1. Целевая модель прогоняет чистую последовательность и отдает скрытые признаки.
2. Из ответа случайно выбираются якорные позиции.
3. Каждый якорь становится первым чистым токеном блока.
4. Остальные позиции внутри блока маскируются.
5. Диффузионный черновик учится восстановить замаскированные токены, используя якорь, внутриблочное внимание и инжектированный контекст целевой модели.

Функция потерь похожа на взвешенную кросс-энтропию по позициям блока:

$${\mathcal{L}}_{\text{DFlash}}=-\sum _b\sum _{j=1}^B{w}_j\log q_{\varphi }(x_{b,j}\mid \mathrm{mask}(x_{b,1:B}),C_b),$$

где $b$ — индекс блока, $j$ — позиция внутри блока, а веса убывают по мере удаления от anchor:

$$w_j=\exp (-\alpha (j-1)).$$

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

Чем DFlash похож на диффузию

Обычные диффузионные модели учатся восстанавливать чистый объект из зашумленного состояния. В текстовом block diffusion вместо непрерывного шума часто используется маскирование: часть токенов скрыта, модель должна восстановить их из контекста.

В DFlash "чистым объектом" является следующий блок токенов, а "зашумленным" состоянием -- блок с замаскированными позициями. Но задача DFlash уже, чем полноценная генерация текста диффузионной моделью. Она не обязана быть самостоятельной языковой моделью уровня целевой модели. Ей достаточно быть хорошим черновиком: быстро дать токены, которые целевая модель с высокой вероятностью примет.

Это важное разделение ролей:

• диффузионный черновик оптимизируется на скорость и близость к целевой модели;
• целевая модель гарантирует качество и корректное итоговое распределение;
• спекулятивная проверка делает ускорение без изменения результата относительно целевой модели.

Baseline vs MTP vs DFlash

Критерий	Базовая авторегрессия	MTP / самоспекуляция	DFlash
Генерация кандидатов	Нет кандидатов, один токен целевой модели за шаг	Несколько голов будущего предсказания или связанный MTP-drafter	Блок токенов через диффузионный черновик
Параллельность	Низкая на уровне токенов	Средняя: головы параллельны, но сложные варианты могут быть частично последовательными	Высокая: блок предсказывается параллельно
Память	Только целевая модель	Целевая модель + небольшие головы/drafter	Целевая модель + отдельный диффузионный черновик
Качество	Качество целевой модели	Без изменения результата при корректной проверке	Без изменения результата при корректной проверке
Скорость	Базовая	Хорошая, если acceptance rate высок и головы дешевые	Очень высокая при длинных принимаемых блоках
Стоимость обучения	Нет дополнительного обучения	Нужно обучить MTP-головы/drafter или брать готовые	Нужно обучить отдельный черновой адаптер
Совместимость с целевой моделью	Очень высокая	Обычно высокая: общий tokenizer, скрытые состояния, KV/cache	Зависит от conditioning; DFlash специально использует признаки целевой модели и общие embedding/LM head
Риск рассогласования	Нет	Низкий	Выше, если внешний drafter плохо выровнен с целевой моделью
Проверка целевой моделью	Не нужна	Нужна	Нужна

Смысл картинки: короткий MTP-блок может быть менее амбициозным, но часто лучше выровнен с целевой моделью. DFlash-блок длиннее и потенциально быстрее, но одна ранняя ошибка уничтожает пользу оставшейся части блока.

Почему выравнивание так важно?

Спекулятивное декодирование проверяет не смысл строки, а конкретную последовательность token id. Поэтому два продолжения могут быть семантически одинаковыми и даже декодироваться в похожую строку, но одно будет принято, а другое отклонено.

Пусть черновик предложил токены

$$u=(u_1,u_2,u_3)$$

а целевая модель на том же месте предпочитает

$$v=(v_1,v_2).$$

Даже если

$$\mathrm{decode}(u_1,u_2,u_3)=\mathrm{decode}(v_1,v_2),$$

проверка сравнивает токены слева направо. Если $u_1\ne v_1$, первый черновой токен отклоняется, и весь хвост $u_2,u_3$ становится бесполезным.

Пример:

черновик:
  ["gr", "avity"]

целевая модель:
  ["grav", "ity"]

Строка может выглядеть почти одинаково, но путь на уровне токенов другой. Для спекулятивного декодирования это разные события. Вероятность принять весь блок длины $B$ примерно раскладывается как произведение условных вероятностей совпадения:

$$\mathrm{Pr}(A=B)=\prod _{j=1}^B\mathrm{Pr}({\hat{x}}_{t+j}=x_{t+j}^{\star }\mid {\hat{x}}_{t+1:t+j-1}=x_{t+1:t+j-1}^{\star }).$$

Чем длиннее блок, тем сильнее накапливается риск. Поэтому важны:

• общий tokenizer
• общий или совместимый vocabulary
• совпадение локального стиля целевой модели
• обучение черновика на outputs самой целевой модели
• conditioning на скрытых состояниях целевой модели, а не только на текстовый префикс

MTP обычно имеет преимущество в выравнивании, потому что читает внутренние состояния целевой модели. DFlash снижает риск рассогласования через признаки целевой модели, инъекцию в KV-cache и общие embedding/LM head, но как внешний черновик он все равно сильнее зависит от качества выравнивания.

Когда выбирать MTP, а когда DFlash

MTP стоит выбирать, если:

• нужен простой самоспекулятивный слой поверх одной целевой модели;
• важна тесная совместимость с tokenizer, KV cache и скрытыми состояниями;
• достаточно коротких черновых блоков;
• есть готовые MTP-веса, как в Gemma 4;
• память ограничена и отдельный diffusion-drafter слишком дорог.

DFlash стоит выбирать, если:

• узкое место — именно последовательное черновое декодирование;
• нужен длинный блок кандидатов за один параллельный шаг;
• есть готовый DFlash-checkpoint для нужной целевой модели;
• GPU лучше загружается крупными параллельными операциями;
• дополнительная память под черновую модель приемлема.

TL;DR

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

MTP делает черновик максимально близким к целевой модели: дополнительные головы или связанный MTP-модуль читают скрытые состояния самой модели. Это дает хорошее выравнивание и относительно дешевую интеграцию.

DFlash делает другой ход: превращает черновое предсказание в block diffusion задачу, где несколько токенов генерируются параллельно. Это увеличивает потенциальный speedup, но требует отдельного хорошо выровненного чернового адаптера.

Главная инженерная метрика в обоих случаях — не perplexity черновика сама по себе, а acceptance length на реальной целевой модели. Чем больше токенов целевая модель принимает за один verification pass и чем дешевле черновик, тем ближе мы к идеальному ускорению без изменения качества.