Core ML

Multi-Token Prediction

During training a model grows several extra heads that guess tokens further ahead at once, not just the next one, then those extra heads get thrown away.

One position producing several prediction heads at once, for t+1, t+2, and t+3; only the t+1 head survives at inference, the rest are discarded

Why Does This Exist?

GPT-style training rewards a model for exactly one thing at every position: guess the very next token correctly. That's an extremely short-horizon signal. A model gets no direct feedback about whether the token it just predicted sets up a coherent sentence five words later — it only ever finds out, indirectly, whether the very next guess was right. Over enough data that turns out to teach a remarkable amount, but it's still a narrow signal by construction, and narrow signals are slower to learn from than they need to be.

Multi-token prediction (MTP) widens that signal without changing what the model does at inference. During training, add a small number of extra prediction heads — instead of just guessing token $t{+}1$ from position $t$ , also guess $t{+}2$ , $t{+}3$ , and so on, each from its own lightweight head reading the same shared representation. The main head stays as it was; the extras exist purely to shape the underlying representation, pushing the model to learn features useful further ahead, not just one step. Once training finishes, every head but the ordinary next-token one gets discarded — a training-time signal boost with zero extra inference cost.

Think of It Like This

Practicing three shots ahead, but only ever taking one

Picture a pool player practicing a specific drill: for every shot, they're asked not just where the ball they're about to hit will land, but also where they expect the next two balls after that to end up, based on the same read of the table. They never actually take those follow-up shots during the drill — only the first one gets played for real. But being forced to reason two shots ahead, every single practice round, sharpens their read of the table in a way that only ever practicing the immediate shot wouldn't.

That's multi-token prediction. The extra heads never influence what actually gets generated — only the real next-token head does that, both in the drill and for real. But training the model to also predict further-ahead tokens, even heads that get thrown away afterward, shapes a sharper internal representation than training on the immediate next token alone.

How It Actually Works

One shared trunk, several heads

The model's main body — the stack of transformer blocks — runs as it would for ordinary next-token training, producing one hidden representation per position. On top of it, MTP attaches $n$ small prediction heads instead of one: head 1 predicts $t{+}1$ as usual, head 2 predicts $t{+}2$ , head 3 predicts $t{+}3$ , and so on. The diagram above shows this fan-out from a single position. Each head is a small, separate set of parameters — a modest addition, since the expensive shared trunk isn't duplicated.

Training on the combined loss

During training, the loss sums the prediction error across all $n$ heads at every position, not just the $t{+}1$ head. Every position contributes $n$ training signals instead of one — a denser gradient per token processed, part of why MTP can improve sample efficiency without changing how much raw text the model sees.

Discarding the extra heads at inference

Once training finishes, generation uses only the ordinary head 1 — predict $t{+}1$ , append it, feed the sequence back in, repeat, exactly like plain autoregressive decoding. Heads 2 through $n$ are simply dropped; they contribute nothing to the model's actual outputs once training stops. This is the detail worth holding onto: MTP changes what happens during training, not what runs during inference. The deployed model is architecturally an ordinary next-token predictor.

Why the discarded heads still help

It's fair to ask why heads that get thrown away would improve the head you keep. The shared trunk — the transformer stack — is trained jointly on all $n$ objectives at once. A representation useful for predicting $t{+}1$ , $t{+}2$ , and $t{+}3$ simultaneously ends up encoding more near-future structure than one only ever asked to support $t{+}1$ . The heads are disposable; the improved representation underneath them survives into the final model.

A second life as a draft model

Because the extra heads exist and are cheap, some systems keep them instead of discarding them, using them to draft several candidate future tokens at once that a separate verification step can accept or reject faster than generating each one in sequence. That technique, speculative decoding, is a distinct topic — MTP's contribution is simply that a model trained this way already has the extra heads available to reuse, which is part of why the technique stuck around in production.

Show Me the Code

Attaching three toy prediction heads to one shared hidden vector, mirroring the fan-out the diagram shows.

1import numpy as np2
3
4def prediction_heads(hidden: np.ndarray, heads: list[np.ndarray]) -> np.ndarray:5    return np.stack([hidden @ h for h in heads])   # one row of logits per head6
7
8rng = np.random.default_rng(0)9hidden = rng.normal(size=(4,))                       # one position's shared representation10heads = [rng.normal(scale=0.3, size=(4, 4)) for _ in range(3)]  # heads for t+1, t+2, t+311preds = prediction_heads(hidden, heads)12print(preds.shape)  # -> (3, 4) — three heads, each predicting over a 4-token vocabulary13print(np.round(preds[0], 4))  # -> [-0.4561  0.0118 -0.1525 -0.0738]  (head 1, t+1)14print(np.round(preds[2], 4))  # -> [ 0.1327 -0.0838  0.2676  0.364 ]  (head 3, t+3, discarded later)

All three heads read the exact same hidden vector — only the weight matrix per head differs, which is what keeps the extra heads cheap relative to the shared trunk.

Watch Out For

Assuming multi-token prediction changes inference-time generation

The heads for $t{+}2$ and beyond are training-only scaffolding in the standard setup — they get discarded, and the deployed model generates one token at a time exactly like an ordinary autoregressive model. If you're benchmarking inference speed or latency, a model trained with MTP shouldn't behave any differently at inference than one trained without it, unless the extra heads are deliberately kept for speculative decoding, which is a separate design choice layered on top.

Expecting every extra head to help equally

Heads predicting further into the future ( $t{+}3$ , $t{+}4$ ) are working with a harder, noisier target than $t{+}1$ — there's more uncertainty about what token sits three positions ahead than one. In practice the benefit from additional heads diminishes past a small number, and too many heads mostly add training cost without a proportional gain. Treat the head count as a real hyperparameter to tune, not something to maximize by default.

The Quick Version

Multi-token prediction adds extra heads during training that predict $t{+}2$ , $t{+}3$ , and beyond, alongside the usual $t{+}1$ head.
All heads read the same shared hidden representation, so the extra cost is small relative to the shared transformer trunk.
The combined loss across heads gives a denser training signal per token than plain next-token prediction alone.
Every head except the ordinary next-token one is discarded at inference — the deployed model behaves like a normal autoregressive model.
The discarded heads can double as a draft model for speculative decoding, which is part of why the technique stuck around in production.