Multi-Token Prediction

Standard causal LM attaches one cross-entropy term to each position: \(-\log P_\theta(x_{t+1} \mid x_{1:t})\).

Multi-token prediction attaches \(n\):

The same shared trunk produces one latent \(z_t\), and \(n\) thin heads (typically one transformer layer each, over a shared unembedding matrix) read \(n\) different future positions off it. The trunk is therefore forced to encode information about tokens \(t+2 \dots t+n\) that a next-token objective never explicitly asks for.

Cause: each head materialises logits of shape (batch, sequence, vocab). With a 32k–128k vocabulary those tensors dwarf the trunk's own activations, and \(n\) heads held simultaneously multiply that dominant term by \(n\).

Fix: order the computation instead of parallelising it. Run the trunk forward once, then loop over heads — forward head \(i\), backpropagate its gradient into the trunk's accumulator, free its logits, move to head \(i+1\). Peak activation memory falls to that of a single head regardless of \(n\); the cost is a longer sequential backward pass.

It is free at inference only because it was paid for during training.

• Scale dependence. Gains grow with model size; Gloeckle et al. report benefit at 7B and 13B while small models can be actively hurt, as the extra heads compete for limited trunk capacity. Validating at 125M and assuming transfer is a mistake.
• Choosing \(n\). Around 4 is the reported sweet spot at 7B–13B. Beyond that, distant targets are genuinely unpredictable — close to noise — so their gradients dilute the useful signal instead of adding to it.
• Irreversible commitment. \(n\) (or \(\lambda\) and \(D\)) is fixed at the start of a multi-trillion-token run and cannot be cheaply re-tuned afterwards.
• Uneven gains. Largest on code, where continuations are more algorithmically determined; weaker on open-ended natural language.

Classical speculative decoding is purely an inference-time technique: attach a separate small draft model to an already-trained large one, verify its guesses in one forward pass, and recover the exact same output distribution. Training is untouched, the decision is reversible, but you must host and keep resident a second model.

Multi-token prediction changes the pretraining objective. Its heads 2…\(n\) happen to be draft heads, so you get self-speculative decoding — no second model, and the drafts reuse the trunk's existing KV cache. The additional payoff is that the objective may make the trunk itself better. The additional risk is that the decision must be made before pretraining and cannot be evaluated until the run finishes.

Historically the direction ran the other way: blockwise parallel decoding (Stern et al., 2018) added future-position heads purely for speed. MTP's claim is that the same heads, trained from scratch, also improve representation quality.

A deliberately conservative variant:

• Depth \(D = 1\) — one extra token beyond the next one, not four.
• Sequential, not parallel. Rather than one latent fanning out to independent futures, the MTP module combines the trunk's hidden state at position \(i\) with the embedding of the token at \(i+1\) before predicting \(i+2\), keeping the prediction chain causal.
• Shared embedding and output layers with the main model, so the extra parameter cost is small.
• Annealed weight: \(\lambda = 0.3\) for the first 10T training tokens, dropped to \(0.1\) for the remaining 4.8T.

Payoff: the module can be dropped at inference, leaving the main model unchanged, or kept for self-speculation. DeepSeek reports the second token being accepted 85–90% of the time, worth roughly 1.8× tokens per second.

Standard causal LM attaches one cross-entropy term to each position: \(-\log P_\theta(x_{t+1} \mid x_{1:t})\).

Multi-token prediction attaches \(n\):

The same shared trunk produces one latent \(z_t\), and \(n\) thin heads (typically one transformer layer each, over a shared unembedding matrix) read \(n\) different future positions off it. The trunk is therefore forced to encode information about tokens \(t+2 \dots t+n\) that a next-token objective never explicitly asks for.

Cause: each head materialises logits of shape (batch, sequence, vocab). With a 32k–128k vocabulary those tensors dwarf the trunk's own activations, and \(n\) heads held simultaneously multiply that dominant term by \(n\).

Fix: order the computation instead of parallelising it. Run the trunk forward once, then loop over heads — forward head \(i\), backpropagate its gradient into the trunk's accumulator, free its logits, move to head \(i+1\). Peak activation memory falls to that of a single head regardless of \(n\); the cost is a longer sequential backward pass.

It is free at inference only because it was paid for during training.

• Scale dependence. Gains grow with model size; Gloeckle et al. report benefit at 7B and 13B while small models can be actively hurt, as the extra heads compete for limited trunk capacity. Validating at 125M and assuming transfer is a mistake.
• Choosing \(n\). Around 4 is the reported sweet spot at 7B–13B. Beyond that, distant targets are genuinely unpredictable — close to noise — so their gradients dilute the useful signal instead of adding to it.
• Irreversible commitment. \(n\) (or \(\lambda\) and \(D\)) is fixed at the start of a multi-trillion-token run and cannot be cheaply re-tuned afterwards.
• Uneven gains. Largest on code, where continuations are more algorithmically determined; weaker on open-ended natural language.

Classical speculative decoding is purely an inference-time technique: attach a separate small draft model to an already-trained large one, verify its guesses in one forward pass, and recover the exact same output distribution. Training is untouched, the decision is reversible, but you must host and keep resident a second model.

Multi-token prediction changes the pretraining objective. Its heads 2…\(n\) happen to be draft heads, so you get self-speculative decoding — no second model, and the drafts reuse the trunk's existing KV cache. The additional payoff is that the objective may make the trunk itself better. The additional risk is that the decision must be made before pretraining and cannot be evaluated until the run finishes.

Historically the direction ran the other way: blockwise parallel decoding (Stern et al., 2018) added future-position heads purely for speed. MTP's claim is that the same heads, trained from scratch, also improve representation quality.

A deliberately conservative variant:

• Depth \(D = 1\) — one extra token beyond the next one, not four.
• Sequential, not parallel. Rather than one latent fanning out to independent futures, the MTP module combines the trunk's hidden state at position \(i\) with the embedding of the token at \(i+1\) before predicting \(i+2\), keeping the prediction chain causal.
• Shared embedding and output layers with the main model, so the extra parameter cost is small.
• Annealed weight: \(\lambda = 0.3\) for the first 10T training tokens, dropped to \(0.1\) for the remaining 4.8T.

Payoff: the module can be dropped at inference, leaving the main model unchanged, or kept for self-speculation. DeepSeek reports the second token being accepted 85–90% of the time, worth roughly 1.8× tokens per second.