← All writing
Paper Breakdown

Multi-token prediction: what the Meta FAIR paper actually says

Reading Gloeckle, Idir, Rozière, Lopez-Paz, and Synnaeve (Meta FAIR, ICML 2024) after hitting the operational overhead of running two models in production for speculative decoding.

The speculative decoding setup had worked perfectly in the benchmark: a 7B draft model generating candidate tokens, a 70B target model verifying in parallel, roughly 2.5× throughput improvement. Then we tried to actually run it in production.

The draft model needed its own replica set. When the target model got a new fine-tuned checkpoint, the draft model's acceptance rate dropped until we fine-tuned it too. The vocab mismatches between a quantized draft and a full-precision target created token boundary edge cases that took three weeks to debug. The memory footprint was 70B + 7B per serving instance, which changed our instance sizing math. And when the acceptance rate fell below about 60% under distribution shift, the "2.5×" throughput improvement became worse than just running the 70B model directly, because the overhead of running the draft and then rejecting its tokens added latency instead of removing it.

Speculative decoding (Leviathan et al., ICML 2023) is the right idea. The draft-model architecture is an operational mess.

Multi-Token Prediction — "Better & Faster Large Language Models via Multi-Token Prediction", Gloeckle, Idir, Rozière, Lopez-Paz, and Synnaeve, Meta FAIR, ICML 2024 — approaches the same problem differently: instead of building a separate draft model at inference time, bake the drafting capability into the main model at training time. You train the model to predict the next N tokens simultaneously using N independent output heads. The same model that generates text also supplies its own speculative drafts — no second deployment, no versioning mismatch, no acceptance rate cliff when you update the checkpoint.

The problem the paper is actually solving

To understand why MTP works, you need to understand what standard next-token prediction training actually teaches the model to do.

In standard autoregressive training, at each position t, the model takes tokens 1…t and predicts token t+1. The gradient signal at position t is: "given everything you've seen, how surprised were you by the actual next token?" The model learns to maintain internal representations sufficient to predict one step ahead.

The paper makes an argument that sounds simple but has real consequences: one-step-ahead prediction may not be the richest training signal you can extract from your data. At position t, you know not just token t+1 but tokens t+2, t+3, …, t+n. Why discard that information? If you train the model to predict token t+2 simultaneously with t+1, the gradient from the t+2 prediction forces the representation at t to encode information about the token after next — more global structure, more semantic planning.

For natural language, this makes a modest difference. For code and math — where adjacent tokens carry strong dependencies (a function opening ( strongly predicts a closing ) many tokens later, an if statement predicts the structure of an else) — it makes a larger one.

How training works

The architecture is straightforward. The transformer body stays unchanged. You add N independent output heads on top, where each head i predicts token t+i:

Head 1: standard LM head → predicts token t+1
Head 2: small residual block + LM projection → predicts token t+2
Head 3: small residual block + LM projection → predicts token t+3
Head 4: small residual block + LM projection → predicts token t+4

Each "small residual block" is a lightweight transformer layer that takes the previous head's intermediate representation and refines it toward the next-step prediction. The key detail is that these blocks are lightweight — the bulk of the compute and parameters remains in the shared trunk.

The training objective is the sum of per-head cross-entropy losses:

L_total = L_1 + L_2 + ... + L_N

Each L_i is the standard next-token prediction loss, but for the token i steps ahead. During training, all heads run in parallel; the gradient flows through all N heads and back through the shared trunk. The trunk learns representations that simultaneously support 1-step, 2-step, and 4-step prediction.

The parameter overhead is real but manageable. If your model has hidden dimension d and vocabulary size V, each additional head adds roughly one transformer block's worth of parameters plus a d×V projection. For a 7B model with d=4096 and V=32,000: each additional head is approximately 130M parameters (for the LM head projection alone, excluding the transformer block). With 3 additional heads, you're adding ~400–600M parameters on top of a 7B model — roughly 7–8% overhead.

At inference time, you discard the additional heads entirely if you're doing standard autoregressive generation. They cost nothing. They only matter if you're using MTP-based speculative decoding.

The inference-time payoff

Here's where the operational argument gets interesting.

In standard speculative decoding, you have two models: a small draft model that proposes K tokens, and a large target model that verifies all K tokens in a single parallel forward pass. The speedup comes from the fact that the target model's forward pass costs roughly the same whether it's processing 1 token or K tokens (you're memory-bandwidth bound, not compute bound), so verifying K tokens at once is almost free compared to generating them one at a time.

The catch: the draft model is a completely separate model with its own weights. Its representations are trained on different objectives (or the same objective but with far fewer parameters), so its token probability distributions differ from the target model's distributions. The acceptance rate — the fraction of draft tokens the target model "accepts" — depends entirely on how well the draft model's distribution approximates the target's. In practice, acceptance rates of 60–80% are typical; below 60%, the overhead outweighs the savings.

With MTP heads, the draft tokens come from output heads that share the target model's full trunk. At position t, after the target model's forward pass generates the representation h_t:

  1. Head 1 produces the prediction for t+1 (standard output)
  2. Head 2 takes a lightweight transform of h_t and produces the prediction for t+2
  3. Head 3 similarly produces the prediction for t+3
  4. Head 4 produces the prediction for t+4

These are draft tokens from the same representation space as the final prediction. The model trunk that produced h_t is the same trunk used to verify token t+1 in the next forward pass. There is no distribution mismatch of the kind that afflicts separate draft models.

In the paper's experiments, this translates to higher acceptance rates — and therefore better speedups — compared to equivalent separate draft models, particularly on structured tasks where the additional heads have useful information about future token structure.

The serving flow becomes:

  1. One forward pass of the 70B model at position t → generates predicted token t+1 AND draft tokens for t+2, t+3, t+4
  2. Next forward pass: process the draft tokens for t+2, t+3, t+4 in parallel, verifying all three at once
  3. Accept the draft tokens that match, reject the first mismatch, update position accordingly
  4. Repeat

Same hardware, same model, no separate process. The additional heads add a small amount of compute per forward pass (~5–10% overhead from the lightweight residual blocks), but you recover it through the higher acceptance rate.

What the results actually show

The headline number is code generation. A 7B model trained with 4-token MTP improves HumanEval pass@1 by approximately 11 percentage points compared to the same architecture trained with standard next-token prediction. On MBPP (another code benchmark), similar gains. The pattern is consistent across model sizes — the benefit doesn't disappear at larger scale.

Natural language benchmarks (common-sense reasoning, reading comprehension) show smaller but positive improvements. The explanation is entropy: in natural language, the distribution over token t+2 given token t is relatively flat — there are many plausible continuations. In code, that distribution is much sharper, so the gradient from predicting t+2 carries more information.

Math tasks fall in between code and language in terms of gain — structured, but less syntactically constrained than code.

The inference-time speculative decoding results show speedups in the 1.5–2.5× range on typical use cases, which is competitive with separate-draft-model approaches but without the operational complexity.

One important result the paper characterizes: MTP training helps most at the tasks that benefit most from planning. If your workload is dominated by free-form creative text generation, you'll see modest improvement. If your workload is code completion, structured output generation, or multi-step reasoning where the next several tokens are partly determined by structure, you'll see larger gains.

What breaks in production

The additional heads require a different checkpoint format. If you're building on a pretrained base (LLaMA, Mistral, etc.), you can't take the existing checkpoint and add MTP — you need to retrain or continue training with the MTP objective from a sufficiently early point that the trunk has learned to support multi-head prediction. Fine-tuning MTP heads onto a frozen trunk (a la Medusa) gets you some inference benefit but doesn't give you the training-quality improvements.

The acceptance rate is task-dependent in ways that are hard to predict before deployment. Your offline benchmarks might show 80% acceptance rates on your test set, and your production distribution might hit 55%. At 55%, you're burning compute on the additional heads and getting minimal speedup. You need monitoring on acceptance rate per request type in production.

The lightweight residual blocks are not actually trivial to implement efficiently in typical serving stacks. vLLM, TensorRT-LLM, and similar frameworks have optimized paths for standard speculative decoding with a separate draft model. MTP requires either custom CUDA kernels for the additional heads or using the framework's generic implementation, which may not achieve the same efficiency. Check your serving stack's MTP support before assuming you'll capture the paper's speedup numbers.

KV cache management gets more complex. In standard autoregressive decoding, the KV cache grows by exactly one token per step. With MTP-based speculative decoding, you speculatively write K tokens to the KV cache before verification, and if tokens are rejected, you need to roll back those entries. This is the same problem as in standard speculative decoding, but it's worth confirming your serving stack handles it correctly — cache corruption on rejection is a subtle bug.

When NOT to use it

Don't use MTP if you're serving a general-purpose assistant with diverse workloads. The training benefit is real but modest for open-ended tasks, and the operational complexity — new checkpoint format, serving stack changes, acceptance rate monitoring — may not pay off unless you're running at scale or code/math tasks dominate your traffic.

Don't use MTP if your serving stack doesn't have first-class speculative decoding support. The training-quality improvement alone (better code benchmarks) might be worth it for specialized use cases, but if you can't capture the inference speedup, you're paying the training overhead without the throughput payback.

Don't use MTP if you need to hot-swap checkpoints frequently. The whole-model retraining requirement means your iterative fine-tuning cycle becomes more expensive if you want to preserve MTP properties. Teams doing weekly checkpoint refreshes from new data may find the overhead unsustainable.

Don't assume MTP replaces a good draft model for all cases. A separate 7B draft model for a 70B target, properly fine-tuned on your task distribution, can still outperform MTP if your serving infrastructure handles the dual-model complexity well. MTP removes operational overhead at the cost of flexibility in choosing draft-model architecture and training.

The design decision that matters

The paper's most practical contribution isn't the benchmark numbers — it's the insight that speculative decoding's draft model can be part of the target model's training, not a separate engineering artifact.

The operational failure mode of separate-draft-model speculative decoding is version drift: your target model gets updated and your draft model's acceptance rate silently degrades. With MTP, the draft heads are trained jointly with the target and always version-locked. You can't get out of sync because there's only one model.

This is the same reason that SWE-bench performance on code models has been dominated by models trained with MTP-like objectives: the model learns to plan its output at the representation level, not just predict locally optimal next tokens. When you ask a model to write a complete function rather than just the next line, that planning signal during training is worth having.

Whether to use MTP comes down to your training setup and workload. If you're training from scratch or have capacity for an extended continued-training run, and if your production traffic is code-heavy or structured-output-heavy, MTP is worth the implementation effort. If you're fine-tuning a pretrained base for a general-purpose task and running it on heterogeneous traffic, the juice probably isn't worth the squeeze.

The paper itself is short (12 pages plus appendix) and unusually clear about what it's measuring and why. Worth reading directly if you're making this decision.


Previously: Speculative Decoding: What the Paper Actually Says covers the Leviathan et al. draft-model approach this post builds on. Continuous Batching (ORCA) covers the scheduling side of LLM serving that MTP-based systems must integrate with.