#53 · Multi-Token Prediction

running

◷train_gpt.py▸mtp_weights in forward◎learner＄no GPU

Blocks

Quick summary

Instead of predicting only the next token, MTP predicts the next n tokens simultaneously and takes a weighted sum of…

Full explanation below the code →

fba-lab — train_gpt.py · mtp_weights in forwardexecuting

// block: mtp_weights in forward · lines 1157–1170$ study train_gpt.py --block mtp-forwardInstead of predicting only the next token, MTP predicts the next n tokens simultaneously and takes a… ✓

Explanation

Instead of predicting only the next token, MTP predicts the next n tokens simultaneously and takes a weighted sum of their cross-entropy losses. The model head is shared (untied from embedding), and mtp_weights determines how much each future token's loss contributes.

Think about

Why does predicting future tokens give more learning signal per forward pass?

// architecture

Live diagram

Record #53 · PR #178 · Multi-token prediction · 119.76 s

Loss signal

Each forward pass now predicts the next 3 tokens, not just 1

Standard language model training predicts one token at a time. Each forward pass computes one loss, one gradient. Multi-token prediction (MTP) predicts the next 2 or 3 tokens simultaneously from the same hidden states. The model receives gradient signal from all future positions at once — more learning per step.

The extra predictions are weighted and annealed: full weight on token+1, half weight on token+2, quarter on token+3. As training progresses the weights drop to zero — 3-token early, 2-token at 33%, standard 1-token from 66% onward. MTP helps early training where diverse gradients matter most; by late training, the clean 1-token objective finishes the job.

# lines 1157–1170 — sum weighted losses for future tokens
for i, w in enumerate(mtp_weights):
loss += w * cross_entropy(logits[:, i:], targets[:, i:])

# lines 1521–1531 — anneal from 3-token to 1-token
mtp_weights_schedule = [(1.0, 0.5, 0.25), (1.0, 0.5), (1.0,)]

Why does predicting further ahead help?

Each additional future token provides a gradient signal that forces the model's internal representations to carry longer-range information. Instead of encoding just enough to predict the next word, the model must plan two or three words ahead. This produces richer hidden states that generalize better.

Why anneal back to 1-token?

Late in training, the model needs to sharpen next-token predictions. The gradient noise from t+2 and t+3 targets — which are inherently harder — blurs the final convergence. Annealing to pure 1-token loss in the last third of training lets the model finish cleanly without interference.

What changed vs #46

Select a change to view diff hunks.

Each forward pass now predicts the next 3 tokens, not just 1

Why does predicting further ahead help?

Why anneal back to 1-token?