Multi-Token Prediction (MTP)

From next-token to next-n-tokens

A causal language model factorizes the probability of a sequence x_1 … x_T autoregressively: p(x) = ∏_t p(x_{t+1} | x_{≤t}). Training maximizes that likelihood, which in practice means minimizing the cross-entropy of the one token that actually follows each position. Every position t contributes exactly one supervised target: the ground-truth token x_{t+1}.

Multi-token prediction keeps the same trunk and the same hidden states, but attaches several output heads to each position. Head k is asked to predict the token k steps into the future, x_{t+k}, for k = 1 … n. Head 1 is the ordinary next-token head; heads 2 through n are the new arrivals. The model still reads only the past — it never peeks at future tokens as inputs — but it is now scored on how well it can anticipate a short horizon of what is coming. Nothing about the causal mask changes; what changes is how many targets each position is held accountable for. That single generalization — from one future target to n of them — is the entire idea, and everything else follows from it.

The training objective, written out

The next-token loss for a single position is the negative log-probability of the true next token:

L_next(t) = − log p(x_{t+1} | x_{≤t})

L_next     = − Σ_t  log p(x_{t+1} | x_{≤t})

MTP generalizes the inner term into a sum over the n future offsets, one per head:

L_MTP(t) = − Σ_{k=1..n}  λ_k · log p_k(x_{t+k} | x_{≤t})

L_MTP    = − Σ_t  Σ_{k=1..n}  λ_k · log p_k(x_{t+k} | x_{≤t})

Here p_k is the distribution produced by head k, and λ_k is an optional per-head weight (often λ_k = 1, or a mild decay so far-future heads count for less). Setting n = 1 recovers the ordinary next-token loss exactly — MTP is a strict superset. Each head has its own cross-entropy against its own shifted target, and the position’s total loss is their (weighted) sum. Because all n targets are just the input sequence shifted by 1, 2, …, n, no extra labels are needed — the denser supervision is free, extracted from text you already have.

The architecture: one trunk, several heads

The compute is deliberately lopsided. A shared trunk — the full stack of transformer layers — reads the context and produces a hidden state h_t at each position, shape [T, d]. That trunk is where essentially all the parameters and FLOPs live. On top of it sit n lightweight heads, each mapping h_t to a distribution over the vocabulary of size V.

The cheapest design (Meta’s formulation, Gloeckle et al. 2024) makes the n heads parallel and independent: every head reads the same trunk output h_t and predicts its offset directly, so head k does not feed head k+1. Each head is often little more than a small transformer/MLP block plus an unembedding projection [d, V]. A richer design (DeepSeek-V3) instead chains sequential MTP modules that each add a transformer block and predict one token deeper while preserving the causal chain, so later predictions are conditioned on earlier ones rather than all sharing a single snapshot. Both share the key economics: the expensive trunk runs once, and the extra heads are a small surcharge bolted onto its output.

Why extra heads barely cost anything

The reason MTP is practical is a shape argument. The trunk’s cost is dominated by L layers of attention and feed-forward over T positions with width d — roughly O(L · T · d^2) for the projections plus the O(T^2 d) attention term. A single prediction head is one matmul from d to V, costing O(T · d · V) per head. Adding n heads multiplies only that last, comparatively thin term by n — it does not touch the deep trunk at all.

With a typical n of 2 to 4, the added forward FLOPs are a few percent, not a few multiples. The real pressure is memory: materializing n logit tensors of shape [T, V] at once, with a large V, can dominate activation memory. The standard fix is to run each head’s forward-and-backward sequentially, accumulating gradients into the shared trunk and freeing each head’s logits before computing the next — so peak memory stays close to the single-head case while all n losses still contribute. You pay a little more compute time; you do not pay n× the logit memory.

Densifying the training signal

The intuitive payoff is that MTP squeezes more learning out of every token of text. Under next-token training, position t receives exactly one bit of feedback: was x_{t+1} predicted well? Under MTP with n = 4, the same position receives four gradients — on x_{t+1}, x_{t+2}, x_{t+3}, and x_{t+4}. The corpus has not grown, but the number of prediction–target pairs the model trains on has, by up to a factor of n.

That density does more than add examples; it changes what the representation must encode. To predict three or four tokens ahead from a single hidden state, h_t can no longer be a myopic summary tuned only to the immediate next word — it has to carry enough about the near future to seed several steps of it. In effect the far-future heads act as an auxiliary objective that regularizes the trunk toward longer-range, more ‘planful’ features. The reported empirical consequence is better sample efficiency: for a fixed amount of training data, MTP models tend to reach a given quality sooner, with the benefit most visible at larger model scales and on generative and reasoning-flavored tasks rather than pure multiple-choice benchmarks.

A worked numeric example

Take the fragment ‘the cat sat on the mat’ and stand at the position right after ‘the cat’. A next-token model is asked for one thing: x_{t+1} = ‘sat’. An MTP model with n = 4 is asked for four: ‘sat’, ‘on’, ‘the’, ‘mat’ — the whole short continuation, from one hidden state.

Suppose the four heads assign the correct tokens probabilities p_1 = 0.50, p_2 = 0.30, p_3 = 0.20, p_4 = 0.10 — sharp nearby, fuzzier further out, exactly as you would expect. With λ_k = 1 the position’s loss is

L = −(ln 0.50 + ln 0.30 + ln 0.20 + ln 0.10)
  = −(−0.693 − 1.204 − 1.609 − 2.303)
  = 5.809 nats

The next-token-only loss at the same position is just the first term, −ln 0.50 = 0.693 nats. The extra three heads contribute the remaining 5.116 nats of gradient — a much larger, richer error signal from the identical context, and one that specifically punishes the model for being unable to see past the very next word.

Shapes, masks, and target construction

Implementation is mostly bookkeeping on shifts. Given an input sequence of length T, head k’s targets are the input shifted left by k: target_k[t] = x_{t+k}. The last k positions of each sequence have no valid target for head k (they run off the end), so they are masked out of that head’s loss. Head 1 masks one trailing position, head n masks n of them — a negligible fraction for realistic sequence lengths.

The trunk keeps its ordinary causal mask; MTP does not let position t attend to x_{t+1}, because that token is a prediction target, not an input. The heads only ever see h_t, which was computed strictly from x_{≤t}. In shapes: trunk output [B, T, d]; each head emits logits [B, T, V]; each is compared to a shifted, masked target [B, T] via cross-entropy; the n scalar losses are combined by the weights λ_k. Because the targets are pure shifts of the input, there is no new data pipeline — a couple of roll-and-mask operations turn any next-token dataset into an MTP one.

Self-speculative decoding: heads that draft

Here is where the training trick pays a second dividend. Ordinary autoregressive decoding is stubbornly serial: one full forward pass yields exactly one token, and decode is memory-bandwidth-bound — you reload the whole model from memory to emit a single token. Speculative decoding attacks that by drafting several candidate tokens cheaply and then verifying them with one pass of the big model, which can check many positions in parallel.

Classic speculative decoding needs a separate small draft model. An MTP model already has the drafter built in: its extra heads predict x_{t+1}, x_{t+2}, …, x_{t+n} in a single pass, so it can propose a short block of future tokens by itself. This is self-speculative decoding — the model speculates about its own continuation, no second network required. The heads guess the next few tokens; the model then runs one verification pass to check whether it would have generated exactly those tokens; the ones that match are accepted in bulk, and generation resumes from the first mismatch. When the guesses are good, you emit multiple tokens per forward pass instead of one.

Drafting, verifying, and the accept rule

The correctness of self-speculation rests on one rule: a drafted block is accepted only up to the first token the base model would not have chosen under its own decoding rule. Concretely, propose ŷ_{1}, …, ŷ_{m} from the heads, then run a single forward pass that computes the base distribution at each of those positions. Walk left to right; keep ŷ_i while it agrees with what the base model would have produced, and stop at the first disagreement, where you instead take the base model’s own token.

The guarantee is that the accepted output is distributionally identical to plain autoregressive decoding — speculation is a speed optimization, not a quality change; you never emit a token the base model would not have emitted. The speedup is governed by the acceptance rate: if, on average, a of each drafted block survives verification, you emit roughly 1 + a tokens per pass. Since far-future heads are less accurate, acceptance falls off with depth, which is exactly why n is small and why reported end-to-end wins land in the roughly 2–3× range rather than scaling linearly with the number of heads.

Tree drafting: hedging the guesses

A single linear draft commits to one guess per future position, so one bad token early truncates the whole block. The refinement is to draft a small tree of candidates instead of a line: at each speculative position, keep the top few token options, forming branching continuations, and verify many of them at once using a specially shaped attention mask (often called tree attention) that lets one forward pass score all branches in parallel.

The payoff is a higher effective acceptance rate for a modest increase in verification width: because several alternative continuations are checked simultaneously, the odds that some path matches the base model deeper into the block go up. The cost is more logits to compute per pass and more bookkeeping to reconstruct which branch was accepted. Tree drafting is not unique to MTP — it is the same machinery that makes Medusa fast — but it composes naturally with MTP’s heads, which already emit per-offset distributions from which the top-few candidates for each branch are read off directly. It is the standard way to squeeze more accepted tokens out of the same set of speculative heads.

Bridge to Medusa: the sibling that skips retraining

Medusa is MTP’s inference-first sibling. It starts from an already trained next-token model and simply bolts on extra decoding heads — the ‘Medusa heads’ — that read the backbone’s last hidden state and predict tokens at offsets +2, +3, …. Those heads are then used for exactly the self-speculative, tree-verified decoding described above. If MTP is ‘train with extra heads and, as a bonus, decode faster,’ Medusa is ‘take a finished model and add the extra heads only to decode faster.’

The distinction is when and how the heads are learned. Medusa-1 freezes the backbone and trains only the new heads, so the base model’s quality is provably untouched and training is cheap. Medusa-2 fine-tunes the heads and backbone together for higher acceptance, trading a little training cost for more speed. Either way the architecture is the MTP architecture — a shared trunk feeding parallel offset heads — repurposed as a retrofit. Medusa shows that the multi-head idea is valuable even if you never pretrained with it: you can graft speculative heads onto any existing checkpoint.

Bridge to EAGLE: speculating in feature space

EAGLE is the other close sibling, and it sharpens the drafter rather than the training loss. Instead of predicting future tokens from the last hidden state, EAGLE runs a small autoregressive drafter at the feature (hidden-state) level: it takes the sequence of the backbone’s own hidden features, plus the embedding of the last sampled token, and predicts the next feature — then reuses the backbone’s existing LM head to turn that predicted feature into a token distribution.

The insight is that a model’s hidden features are smoother and more predictable than its discretized token outputs, so drafting in feature space and unrolling autoregressively yields more accurate multi-step guesses than a bank of independent offset heads that all condition on a single frozen snapshot. Higher draft accuracy means a higher acceptance rate, which means bigger speedups. EAGLE still verifies with the base model and still uses tree-style candidates — the safety and the parallel check are identical — but its drafter is a tiny recurrent-in-features network rather than parallel token heads. It is the same family solving the same problem, moved one representation deeper.

How the three relate: one idea, three angles

Seen together, MTP, Medusa, and EAGLE are three cuts of a single idea: let the model predict more than one step so it can emit more than one token per pass. They differ mainly in where the multi-step capability comes from. MTP bakes it into pretraining, so the extra heads improve the model itself and drafting is a free byproduct. Medusa adds it after training as a lightweight retrofit, purely for decoding speed, with the backbone optionally frozen. EAGLE upgrades the drafter’s representation, autoregressing over hidden features for more accurate speculation.

All three lean on the same verification contract — draft cheap, verify with the full model in one parallel pass, accept the longest matching prefix — so all three are lossless with respect to the base model’s output distribution. And all three inherit the same governing quantity, the acceptance rate, which is why their designs converge on tree drafting and shallow horizons. The practical takeaway is that ‘predict several tokens’ is not one technique but a small design space: choose it at pretraining time (MTP), as a retrofit (Medusa), or as a smarter drafter (EAGLE), depending on whether you control the training run and how much acceptance you need.

Implications for small models and CPU inference

For small language models and CPU-bound serving, the MTP family is unusually well matched to the bottleneck. On a CPU, decode is dominated by memory bandwidth: each generated token drags the model’s weights through the cache hierarchy, and arithmetic units sit idle waiting on memory. Emitting several accepted tokens per weight-load directly amortizes that dominant cost, so self-speculative decoding can deliver a real wall-clock win precisely where a GPU-style throughput trick would not help.

The catch is that the win is entirely a function of acceptance, and acceptance is workload-dependent: highly predictable, templated, or code-like text drafts well and accelerates a lot, while high-entropy creative text accepts fewer drafted tokens and gains less. For a small model where every head still costs a real fraction of a tiny trunk, keep n modest — often 2 or 3 — so the drafting overhead does not eat the savings on rejected blocks. Medusa-style retrofitting is especially attractive here: you can add speculative heads to an existing small checkpoint without a full retrain, measure the acceptance rate on your actual traffic, and keep the heads only if the numbers justify them.

Common pitfalls and honest caveats

Several traps recur. First, logit memory: naively holding n tensors of shape [B, T, V] can blow up activation memory for large vocabularies — compute the heads’ losses sequentially and free each before the next. Second, expecting linear speedups: n heads do not give n× throughput, because deeper heads are less accurate and rejected drafts still cost a verification pass; the realistic range is a small multiple, and only on drafting-friendly text.

Third, conflating training and inference benefits: MTP’s sample-efficiency gains and its decoding speedups are separate effects — a model can help with one and not the other, and reported quality improvements are clearest at scale and on generative tasks, not universal. Fourth, weighting the heads: if far-future heads are given too much weight they can drag on the primary next-token objective, so many recipes down-weight or decay λ_k with k. Finally, remember that self-speculation is lossless only if the accept rule is exact — a sloppy verification that accepts near-misses silently changes the output distribution and forfeits the one guarantee that makes the whole scheme safe to deploy.