Multi-Token Prediction

Зачем это нужно

Языковая модель учится на одной задаче: по префиксу угадать следующий токен. Из текста длиной $T$ токенов получается $T$ обучающих примеров, по одному на позицию. Это и есть весь сигнал, за который заплачено триллионами токенов и миллионами GPU-часов.

предсказание нескольких токеновпредсказание нескольких токенов (MTP)EN multi-token prediction (MTP) задаёт вопрос: а что если на каждой позиции требовать ещё и токен через один, через два? Модель получает больше сигнала с того же текста, и ей приходится «заглядывать вперёд», а не только реагировать на последние символы. Есть и второй бонус. Голова, которая предсказывает следующий-за-следующим токен, — готовый черновик для спекулятивное декодированиеспекулятивное декодированиеEN speculative decoding, и генерация ускоряется почти вдвое.

DeepSeek-V3 обучалась с MTP. Затем его взяли Qwen3-Next, Xiaomi MiMo и другие. В этом уроке разберём две конструкции MTP, соберём одну из них на PyTorch и своими руками ускорим генерацию.

Идея: учиться предсказывать дальше

Обычный loss из урока 1.3.2 — средняя cross-entropy следующего токена:

$$\mathcal{L}_{\text{main}} = -\frac{1}{T}\sum_{i=1}^{T} \log p_\theta\left(t_{i+1} \mid t_{1..i}\right).$$

MTP глубины $D$ добавляет к нему слагаемые для токенов $t_{i+2}, \dots, t_{i+D+1}$. Каждая позиция теперь даёт не один обучающий пример, а $D + 1$. Текст тот же (у V3 — те же 14.8T токенов), а сигнала больше.

Почему это помогает? Возьмём префикс «Москва — столица». Следующий токен «России» угадывается почти по шаблону. А вот чтобы предсказать токен через два-три шага, модели нужно «держать в голове», куда идёт фраза. Предсказание следующего токена можно частично решить локальными подсказками. Предсказание дальше требует представлений, в которых закодирован план.

Авторы работы Meta (Gloeckle и соавторы, 2024) показали это на коде: модель на 13B с MTP решала на 12% больше задач HumanEval и на 17% больше задач MBPP (два набора задач «напиши функцию на Python», проверяемых тестами), чем такая же модель без MTP. Важная оговорка из той же работы: польза растёт с размером модели, а у маленьких моделей MTP скорее мешает. Им не хватает ёмкости на дополнительную задачу.

Вариант 1: параллельные головы (Meta, 2024)

У модели общий «ствол» (trunk) — стопка блоков трансформера. Над стволом стоят $n$ независимых голов, каждая — ещё один блок трансформера: голова $k$ из того же скрытого состояния ствола $h_i$ предсказывает токен $t_{i+k}$. Выходной слой (unembedding) у всех голов общий. Лучшим для кода оказалось $n = 4$.

flowchart LR
T["токены t1 … ti"] --> S["общий ствол"]
S --> H1["голова 1"] --> P1["t(i+1)"]
S --> H2["голова 2"] --> P2["t(i+2)"]
S --> H3["голова 3"] --> P3["t(i+3)"]
S --> H4["голова 4"] --> P4["t(i+4)"]

Параллельные головы Meta: каждая предсказывает свой токен вперёд из одного и того же h_i

Инженерная деталь: логиты $n$ голов по словарю $V$ — это $n$ огромных тензоров. Чтобы не держать их одновременно, авторы делают прямой и обратный проход головы по очереди и накапливают градиент на выходе ствола. Пиковая память почти та же, что у обычной модели.

Слабость конструкции видна на примере. После «Я люблю» следующим может быть «кошек» или «собак», и голова 2 должна предсказать токен после этого слова, не зная, какое из них выбрано. Ей приходится «усреднять» по обоим вариантам. Головы не согласованы друг с другом.

Вариант 2: последовательные модули (DeepSeek-V3)

DeepSeek-V3 строит MTP-модули цепочкой. Модуль глубины $k$ берёт представление с предыдущей глубины ($h^0$ — выход основной модели) и эмбеддинг токена $t_{i+k}$, который уже известен из текста:

$$\begin{gathered} h_i'^{\,k} = M_k\left[\mathrm{RMSNorm}\left(h_i^{k-1}\right);\ \mathrm{RMSNorm}\left(\mathrm{Emb}(t_{i+k})\right)\right], \\ h^{k} = \mathrm{TRM}_k\left(h'^{\,k}\right), \qquad P^{k}_{i+k+1} = \mathrm{OutHead}\left(h_i^{k}\right). \end{gathered}$$

Что здесь написано:

1. Два вектора размерности $d$ нормируются и склеиваются в вектор длины $2d$ (квадратные скобки с точкой с запятой — склейка).
2. Матрица $M_k$ размера $d \times 2d$ сжимает его обратно до $d$.
3. Свой блок трансформера $\mathrm{TRM}_k$ (с причинным вниманием по позициям).
4. Общий с основной моделью выходной слой выдаёт распределение для токена $t_{i+k+1}$.

Эмбеддинги и выходной слой общие с основной моделью — экономия памяти. Итоговый loss:

$$\mathcal{L} = \mathcal{L}_{\text{main}} + \frac{\lambda}{D}\sum_{k=1}^{D} \mathcal{L}^{k}_{\text{MTP}}.$$

В V3 глубина $D = 1$: один модуль, который предсказывает токен через один. Вес $\lambda = 0.3$ на первых 10T токенов и 0.1 на оставшихся 4.8T.

Разберём на примере. Пусть текст разбит на токены У · ▁лу · ко · морья · ▁дуб:

Позиция $i$	Основная модель видит	и предсказывает	MTP-модуль видит	и предсказывает
1	У	▁лу	$h_1$ + эмбеддинг ▁лу	ко
2	У▁лу	ко	$h_2$ + эмбеддинг ко	морья
3	У▁луко	морья	$h_3$ + эмбеддинг морья	▁дуб

MTP-модуль знает, каким оказался $t_{i+1}$, поэтому причинная цепочка сохраняется: прогноз на два шага вперёд опирается на настоящий следующий токен, а не на догадку. Это главное отличие от параллельных голов.

flowchart LR
X["t1 … ti"] --> MM["основная модель"]
MM --> HM["h_i"] --> O1["OutHead → t(i+1)"]
HM --> N1["RMSNorm"]
E["Emb(t(i+1))"] --> N2["RMSNorm"]
N1 --> C["склейка → M1"]
N2 --> C
C --> TR["блок TRM1"] --> O2["OutHead → t(i+2)"]

MTP-модуль глубины 1 в DeepSeek-V3: эмбеддинги и выходной слой общие с основной моделью

	Параллельные головы (Meta)	Последовательные модули (DeepSeek-V3)
Вход головы $k$	только $h_i$	$h_i^{k-1}$ и эмбеддинг $t_{i+k}$
Причинная цепочка	нет: $t_{i+2}$ без знания $t_{i+1}$	да
Типичная глубина	4	1
На инференсе	можно выбросить или взять для черновиков	то же

MTP на инференсе

У MTP-модулей два пути после обучения.

Выбросить. Основная модель работает сама по себе, модули ей не нужны. Так DeepSeek проводила абляции (§4.5 отчёта): модели с MTP и без него сравнивались при одинаковой стоимости инференса. Выигрыш на бенчмарках получен целиком за счёт лучшего обучения.

Оставить как черновик. Идея спекулятивного декодированияспекулятивное декодированиеEN speculative decoding: дешёвая модель предлагает следующие токены, а большая проверяет их за один проход. Проверить несколько токенов стоит почти столько же, сколько сгенерировать один. При генерации GPU больше всего времени тратит на чтение весов модели из памяти, а прочитанные веса можно сразу применить к нескольким позициям (подробно — в уроке 7.1.1). С MTP черновик пишет сама модель — это самоспекуляция, отдельная маленькая модель не нужна.

Один шаг генерации с MTP глубины 1 выглядит так:

1. Основная модель выдала следующий токен $t_{n+1}$ и скрытое состояние $h_n$.
2. MTP-модуль по $h_n$ и эмбеддингу $t_{n+1}$ предлагает черновик $\tilde t_{n+2}$.
3. Основная модель прогоняет последовательность с $t_{n+1}$ и $\tilde t_{n+2}$. На позиции $t_{n+1}$ она видит, какой токен выбрала бы сама. Совпало с черновиком — токен принят бесплатно, и тем же проходом получен ещё и $t_{n+3}$. Не совпало — берём токен основной модели.

При жадной генерации результат в точности совпадает с обычной генерацией: каждый токен в итоге выбран основной моделью. Для сэмплирования с температурой есть правило приёма с вероятностью $\min(1, p/q)$, где $p$ и $q$ — вероятности чернового токена у основной модели и у черновика (Leviathan и соавторы). Оно тоже не меняет распределение ответов. Подробно — в уроке 7.1.2.

Сколько это даёт? Пусть $a$ — доля принятых черновых токеновдоля принятых черновых токенов (acceptance rate)EN acceptance rateВ спекулятивном декодировании — какая доля токенов-черновиков совпала с выбором основной модели. У MTP-модуля DeepSeek-V3 — 85–90%., доля принятых черновиков. При глубине 1 один проход основной модели даёт в среднем $1 + a$ токенов. При глубине $D$ цепочка черновиков обрывается на первом отказе, и выходит $1 + a + a^2 + \dots + a^D$:

$a$	$D = 1$	$D = 2$	$D = 3$
0.6	1.60	1.96	2.18
0.8	1.80	2.44	2.95
0.9	1.90	2.71	3.44

DeepSeek сообщает для V3: второй токен принимается в 85–90% случаев, скорость генерации растёт примерно в 1.8 раза. По таблице выходит 1.85–1.9 токена за проход. Разница — накладные расходы: сам MTP-модуль и проверка тоже стоят времени.

Практика: GPT с MTP-модулем на PyTorch

Соберём маленькую GPT по образцу урока 4.6.2, добавим MTP-модуль глубины 1 в двух вариантах и сравним их. Корпус — начало пролога к «Руслану и Людмиле» Пушкина, токены — символы. Обучение двух моделей занимает около минуты на CPU.

import torch
import torch.nn as nn
import torch.nn.functional as F

TEXT = """У лукоморья дуб зелёный;
Златая цепь на дубе том:
И днём и ночью кот учёный
Всё ходит по цепи кругом;
Идёт направо — песнь заводит,
Налево — сказку говорит.
Там чудеса: там леший бродит,
Русалка на ветвях сидит;
Там на неведомых дорожках
Следы невиданных зверей;
Избушка там на курьих ножках
Стоит без окон, без дверей;
Там лес и дол видений полны;
Там о заре прихлынут волны
На брег песчаный и пустой,
И тридцать витязей прекрасных
Чредой из вод выходят ясных,
И с ними дядька их морской;
"""
# nn.RMSNorm в коде ниже есть в PyTorch 2.4 и новее
chars = sorted(set(TEXT))                       # посимвольный «токенизатор»
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in TEXT])
V, D, T = len(chars), 64, 64                    # словарь, ширина модели, длина контекста


class Block(nn.Module):
    """Pre-Norm блок трансформера: причинное внимание + MLP (урок 4.5.1)."""

    def __init__(self, d, n_heads=4):
        super().__init__()
        self.n1, self.n2 = nn.RMSNorm(d), nn.RMSNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d, bias=False), nn.Linear(d, d, bias=False)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d, bias=False), nn.GELU(), nn.Linear(4 * d, d, bias=False))
        self.h = n_heads

    def forward(self, x):
        B, L, d = x.shape
        # (B, L, 3d) → (B, L, 3, H, d_h) → permute → (3, B, H, L, d_h): первая ось распаковывается в q, k, v
        q, k, v = self.qkv(self.n1(x)).view(B, L, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        x = x + self.proj(a.transpose(1, 2).reshape(B, L, d))
        return x + self.mlp(self.n2(x))


class GPTWithMTP(nn.Module):
    """Маленькая GPT + MTP-модуль глубины 1.

    kind="seq" — как DeepSeek-V3: модуль видит h_i И эмбеддинг t_(i+1), предсказывает t_(i+2).
    kind="par" — как параллельные головы Meta: модуль видит только h_i.
    """

    def __init__(self, kind="seq"):
        super().__init__()
        self.kind = kind
        self.emb, self.pos = nn.Embedding(V, D), nn.Embedding(T + 2, D)
        nn.init.normal_(self.emb.weight, std=0.02)   # маленькие эмбеддинги → стартовый loss ≈ ln V
        nn.init.normal_(self.pos.weight, std=0.02)
        self.blocks = nn.Sequential(Block(D), Block(D))
        self.norm = nn.RMSNorm(D)
        self.mtp_norm_h, self.mtp_norm_e = nn.RMSNorm(D), nn.RMSNorm(D)
        self.mtp_proj = nn.Linear(2 * D if kind == "seq" else D, D, bias=False)   # матрица M_1
        self.mtp_block = Block(D)                                                  # свой блок TRM_1
        self.mtp_norm = nn.RMSNorm(D)

    def head(self, h):
        return h @ self.emb.weight.T                 # выходной слой общий с эмбеддингами

    def main(self, idx):
        """Основная модель: скрытые состояния h_i и логиты для t_(i+1)."""
        h = self.blocks(self.emb(idx) + self.pos(torch.arange(idx.shape[1])))
        return h, self.head(self.norm(h))

    def mtp(self, h, next_tokens):
        """Логиты для t_(i+2) в каждой позиции i."""
        z = self.mtp_norm_h(h)
        if self.kind == "seq":                       # склеиваем с эмбеддингом уже известного t_(i+1)
            z = torch.cat([z, self.mtp_norm_e(self.emb(next_tokens))], dim=-1)
        return self.head(self.mtp_norm(self.mtp_block(self.mtp_proj(z))))


def batch(bs=32):
    ix = torch.randint(len(data) - T - 2, (bs,))
    c = torch.stack([data[i : i + T + 2] for i in ix])
    return c[:, :T], c[:, 1 : T + 1], c[:, 2 : T + 2]    # вход, цель через 1 токен, цель через 2


def train(kind, steps, lam=0.3):
    torch.manual_seed(0)
    model = GPTWithMTP(kind)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    for step in range(steps + 1):
        x, y1, y2 = batch()
        h, logits = model.main(x)
        loss_main = F.cross_entropy(logits.reshape(-1, V), y1.reshape(-1))
        loss_mtp = F.cross_entropy(model.mtp(h, y1).reshape(-1, V), y2.reshape(-1))
        loss = loss_main + lam * loss_mtp            # L = L_main + λ · L_MTP
        opt.zero_grad()
        loss.backward()
        opt.step()
        if step % 50 == 0:
            print(f"{kind} шаг {step:3d}: main {loss_main.item():.3f}  mtp {loss_mtp.item():.3f}")
    return model


@torch.no_grad()
def generate_plain(model, prompt, n):
    """Обычная жадная генерация: один проход основной модели на токен."""
    ctx, passes = [stoi[c] for c in prompt], 0
    while len(ctx) < len(prompt) + n:
        _, logits = model.main(torch.tensor([ctx]))
        passes += 1
        ctx.append(int(logits[0, -1].argmax()))
    return ctx, passes


@torch.no_grad()
def generate_speculative(model, prompt, n):
    """Самоспекуляция: MTP предлагает черновик, основная модель проверяет его тем же проходом."""
    ctx = [stoi[c] for c in prompt]
    h, logits = model.main(torch.tensor([ctx]))
    passes, drafts, accepted = 1, 0, 0
    nxt = int(logits[0, -1].argmax())               # следующий токен — уже от основной модели
    while len(ctx) < len(prompt) + n:
        # 1) черновик токена ПОСЛЕ nxt (дёшево: один блок вместо всей модели).
        #    MTP на позиции i ждёт h_i и эмбеддинг t_(i+1), поэтому «следующие токены» — это ctx со сдвигом на 1 плюс nxt
        draft = int(model.mtp(h, torch.tensor([ctx[1:] + [nxt]]))[0, -1].argmax())
        drafts += 1
        # 2) один проход основной модели и проверяет черновик, и даёт следующий токен
        h2, logits2 = model.main(torch.tensor([ctx + [nxt, draft]]))
        passes += 1
        ctx.append(nxt)
        check = int(logits2[0, len(ctx) - 1].argmax())   # что основная модель хотела бы после nxt
        if check == draft:                               # совпало: черновик принят «бесплатно»
            accepted += 1
            ctx.append(draft)
            h, nxt = h2, int(logits2[0, -1].argmax())
        else:                                            # не совпало: берём токен основной модели
            h, nxt = h2[:, : len(ctx)], check
    return ctx[: len(prompt) + n], passes, accepted / drafts


prompt, n = "Там ", 50
for kind in ["seq", "par"]:
    model = train(kind, steps=150)
    plain, p1 = generate_plain(model, prompt, n)
    spec, p2, acc = generate_speculative(model, prompt, n)
    print(repr("".join(chars[i] for i in spec)))
    print(f"{kind}: тексты совпадают: {plain == spec}; проходов: {p1} → {p2}; принято черновиков: {acc:.0%}\n")

seq шаг   0: main 3.850  mtp 3.807
seq шаг  50: main 2.188  mtp 1.853
seq шаг 100: main 0.910  mtp 0.548
seq шаг 150: main 0.168  mtp 0.124
'Там лес и ви видений полны;\nТам о заре прихлынут волны'
seq: тексты совпадают: True; проходов: 50 → 27; принято черновиков: 96%

par шаг   0: main 3.860  mtp 3.776
par шаг  50: main 2.408  mtp 2.997
par шаг 100: main 1.327  mtp 2.158
par шаг 150: main 0.281  mtp 0.603
'Там лет — с при — придцаскрускрамых\nЧредой из водод ят'
par: тексты совпадают: True; проходов: 50 → 40; принято черновиков: 28%

Точные числа на вашей машине могут немного отличаться (другое число потоков CPU — другие округления). Что здесь видно:

• Стартовый loss 3.85 — почти ровно $\ln V$ для словаря из 45 символов ($\ln 45 \approx 3.81$). Проверка здравого смысла из урока 1.3.2 пройдена. Без std=0.02 у эмбеддингов он был бы больше 40: огромные стартовые логиты.
• Loss последовательного MTP ниже, чем параллельного (0.12 против 0.60). Модулю DeepSeek-типа известен $t_{i+1}$, и предсказать $t_{i+2}$ ему не сложнее, чем основной модели предсказать следующий токен. Параллельной голове приходится гадать.
• Тексты совпадают в обоих вариантах: при жадной генерации спекуляция не меняет ответ, только ускоряет его.
• Доля принятия решает всё. Последовательный модуль: 96% черновиков принято, 27 проходов вместо 50. Параллельный: 28%, 40 проходов. Хороший черновик — это черновик, который согласован с основной моделью.

Не переоценивайте 96%

Модель почти выучила стихотворение наизусть, поэтому её поведение легко предсказать. На реальном тексте доля принятия ниже: у DeepSeek-V3 — 85–90%. Попробуйте обучить модели дольше (steps=400) или короче (steps=60) и посмотрите, как меняется доля принятия. Ещё опыт: поставьте lam=0 — MTP-модуль перестанет учиться, и черновики станут случайными.

Обратите внимание на две детали кода. Генерация без KV-кэша прогоняет весь контекст заново: это просто, но медленно. С KV-кэшем (урок 4.6.4) экономия от спекуляции сохраняется. И проверка черновика возвращает два токена за проход, потому что основная модель уже посчитала логиты и на позиции черновика.

🖥 На вашей машине (4×GPU)

1. Добавьте MTP-модуль в свою GPT из урока 4.6.2 (или в MiniMind из урока 5.6.1) и обучите на настоящем русском корпусе на одной RTX 4080. Сравните val loss основной модели с MTP и без него при равном числе шагов. На моделях в десятки миллионов параметров не удивляйтесь, если MTP не поможет: по данным Meta, выигрыш появляется с ростом модели.
2. Измерьте долю принятия черновиков на валидационном тексте и пересчитайте ожидаемое ускорение по таблице выше.
3. В серверах инференса vLLM и SGLang есть спекулятивное декодирование через MTP-модули для моделей DeepSeek и Qwen3-Next. Запуск и замеры токенов в секунду — в модуле 7.1.

Три остальные карты в это время свободны: запустите на них тот же опыт с другими seed (CUDA_VISIBLE_DEVICES=1 ...), чтобы отличить эффект MTP от случайного разброса. Настройка карт — на странице «Ваша лаборатория (4×GPU)».

Китайский угол: от идеи Meta к стандарту китайских моделей

DeepSeek-V3 (декабрь 2024) взяла идею MTP у Meta, но перестроила её: вместо независимых голов — последовательные модули, сохраняющие причинную цепочку. В отчёте MTP проверен абляциями на моделях меньшего размера: качество на большинстве бенчмарков растёт при той же стоимости инференса. Модуль глубины 1 опубликован вместе с весами — поэтому в репозитории модели на Hugging Face 685B параметров (671B основной модели + 14B MTP-модуля). Пересчитать эти числа можно в уроке china/06.

Qwen3-Next (сентябрь 2025) встроил MTP как «родной» механизм. Команда Qwen подчёркивает две вещи: MTP-модуль даёт высокую долю принятых черновиков, и его обучают в многошаговом режиме, согласованном с тем, как он работает на инференсе.

Xiaomi MiMo тоже использует MTP — и в MiMo-7B, и в большой MiMo-V2-Flash. За полтора года MTP превратился из статьи Meta в типовую деталь китайских открытых моделей.

Где почитать по-китайски. В курсе ZOMI AIInfra есть глава «Генерация нескольких токенов»: параллельные и последовательные MTP, а затем спекулятивное декодирование вплоть до EAGLE-3 — хороший мостик к модулю 7.

Термины. MTP по-китайски — «предсказание нескольких токенов»; сам токен пишут либо по-английски, либо называют «элементом слова» (урок 4.2.1). Спекулятивное декодирование переводят дословно, и «спекулятивный» здесь то же слово, что в «биржевой спекуляции». Черновик и доля принятия тоже называются дословно.

Проверь себя

Проверь себя5 вопросов·

1 Что получает на вход MTP-модуль глубины 1 в DeepSeek-V3 на позиции $i$?

Только скрытое состояние основной модели $h_i$Скрытое состояние $h_i$ и эмбеддинг настоящего следующего токена $t_{i+1}$Логиты основной модели на позиции $i$Эмбеддинги всех будущих токенов

2 Почему параллельной голове Meta предсказывать $t_{i+2}$ труднее, чем MTP-модулю DeepSeek?

У неё меньше параметровОна не знает, каким был $t_{i+1}$, и должна учитывать все его варианты сразуОна не использует общий выходной слойЕй запрещено внимание на прошлые токены

3 Самоспекуляция с MTP глубины 1, доля принятия черновиков 0.9. Сколько токенов в среднем даёт один проход основной модели?

4 MTP глубины 2, каждый черновой токен принимается с вероятностью 0.8, цепочка обрывается на первом отказе. Сколько токенов за проход в среднем?

5 Какие утверждения о MTP в DeepSeek-V3 верны? Выберите все.Можно выбрать несколько ответов

Эмбеддинги и выходной слой MTP-модуля общие с основной модельюНа инференсе MTP-модуль можно выбросить, основная модель работает без негоВес $\lambda$ MTP-лосса снижали с 0.3 до 0.1 во второй части обученияПри жадной самоспекуляции текст получается немного другим, чем без неё

Проверить ответы

Итоги

• MTP требует на каждой позиции предсказывать не только следующий токен, но и дальнейшие: больше сигнала с того же текста. Польза растёт с размером модели.
• Meta: параллельные головы над общим стволом, $n = 4$. DeepSeek-V3: последовательный модуль, который видит эмбеддинг $t_{i+1}$, $D = 1$, $\mathcal{L} = \mathcal{L}_{\text{main}} + \lambda \mathcal{L}_{\text{MTP}}$.
• После обучения модуль можно выбросить или превратить в черновик для самоспекуляции: $1 + a$ токена за проход, у V3 — около 1.8× к скорости генерации.
• Чем согласованнее черновик с основной моделью, тем выше доля принятия. Последовательный модуль выигрывает именно здесь.

Дальше: распределённое обучение — как разложить модель на 2048 GPU и почему DeepSeek пришлось изобрести собственное расписание конвейера.