Multi-Token Prediction

A training objective that asks each position to predict several future tokens through independent output heads on a shared model trunk.

Multi-token prediction is a language-model training objective where each position predicts the next N future tokens at once, using independent output heads on a shared trunk, instead of predicting only the single next token.

At a glance

Released

April 2024

Authors

Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, et al.

Optimizes

  • Sample Efficiency
  • Code Generation
  • Inference Speed
  • Next Token Prediction

What It Is

Multi-token prediction is a training objective for decoder language models. At each token position, the model still reads the same prefix of earlier tokens, but it is trained to predict several future tokens instead of only the next one. Each future offset gets its own output head on top of the same shared trunk, so the heads stay independent even though they read the same hidden state.

Why It Exists

Standard next-token prediction gives the model one supervised target per position. Multi-token prediction adds more future targets from the same context, which the paper argues can improve sample efficiency during pretraining and can optionally speed up inference when the extra heads are reused for self-speculative decoding.

How It Works

The baseline next-token objective trains one output head to predict token t+1 from the hidden state at position t. Multi-token prediction keeps that shared trunk but adds N independent heads that predict offsets t+1 through t+N from the same hidden state. During training, each head is scored against its matching future token, and the losses are combined. At inference time, the model can still generate with the primary next-token head, or use the auxiliary heads when a serving stack wants draft-and-verify style speedups.
Multi-token prediction keeps one shared trunk but adds independent heads for future offsets t+1 through t+N, while standard next-token prediction supervises only x_{t+1}.

Math Or Compute Schema

Both objectives read the same prefix and shared hidden state h_t at position t. The next-token objective supervises one future token. The multi-token objective averages the same style of loss across offsets 1 through N, with each offset using its own output head on h_t.
Next-token prediction objective
LNTP=−log⁡pθ(xt+1∣ht)\mathcal{L}_{\text{NTP}} = -\log p_\theta(x_{t+1} \mid h_t)
xtx_t
Token at position t.
hth_t
Hidden state at position t from the shared trunk.
xt+1x_{t+1}
Observed next token at offset 1.
x≤tx_{\leq t}
Prefix tokens up to position t.
θ\theta
Model parameters.
pθ(⋅)p_\theta(\cdot)
Conditional distribution from one output head on h_t.
Multi-token prediction objective
LMTP=1N∑k=1N−log⁡pθ(k)(xt+k∣ht)\mathcal{L}_{\text{MTP}} = \frac{1}{N}\sum_{k=1}^{N} -\log p_\theta^{(k)}(x_{t+k} \mid h_t)
xtx_t
Token at position t.
hth_t
Hidden state at position t from the shared trunk.
NN
Number of future offsets supervised at each position.
kk
Future offset index from 1 to N.
xt+kx_{t+k}
Observed token k steps ahead of position t.
x≤tx_{\leq t}
Prefix tokens up to position t.
θ\theta
Shared trunk parameters plus independent head parameters.
pθ(k)(⋅)p_\theta^{(k)}(\cdot)
Conditional distribution from output head k on h_t.

Compared To Nearby Modules

Next-token pretraining is the usual baseline objective: each position supervises one future token at offset t+1. Multi-token prediction changes training by adding independent heads for offsets t+2 through t+N from the same hidden state. Ordinary multi-step generation is different: at inference the model still emits one token per step and re-runs the trunk as the prefix grows, unless a serving stack uses auxiliary heads for drafting. Speculative decoding is an inference-time serving technique that can reuse MTP auxiliary heads, but ordinary generation can still rely on the primary next-token head alone.
Comparison dimensionNext-token pretrainingMulti-token predictionSpeculative decoding
Prediction targets per positionOne future token at offset t+1N future tokens at offsets t+1 through t+NDraft tokens proposed then verified by the main model
Training-time roleDefault pretraining objective for decoder language modelsAuxiliary objective with independent heads on a shared trunkNot a training objective by itself
Inference-time roleStandard autoregressive next-token generationCan keep the primary next-token head or reuse auxiliary heads for self-speculative decodingServing technique that can accept drafts from MTP heads or a separate draft model

Example Architectures

The source paper studies decoder-only transformer language models trained with multi-token prediction as an auxiliary objective on top of the standard trunk. The method is presented as a training recipe that can be added to large-scale pretraining runs rather than as one fixed public checkpoint family.

Limitations And Tradeoffs

Multi-token prediction adds output-head complexity and depends on the chosen horizon N. Larger N increases supervision signal but also increases head count and training overhead. The paper reports stronger gains at larger model sizes and on code generation, but the method is not a universal replacement for every serving or decoding optimization.

Why It Still Matters

Multi-token prediction matters because it reframes what each training position supervises without changing the core autoregressive trunk. That makes it a useful bridge between ordinary next-token pretraining and faster inference ideas that draft multiple tokens before verification.

Tags

References

  1. Gloeckle, Fabian, et al. "Better & Faster Large Language Models via Multi-token Prediction." arXiv, 2024, https://arxiv.org/abs/2404.19737.

Mamba Selective State-Space Module

Previous Page

U-Net

Next Page

On this page

What It IsWhy It ExistsHow It WorksMath Or Compute SchemaCompared To Nearby ModulesExample ArchitecturesLimitations And TradeoffsWhy It Still MattersRelatedTagsReferences