Multi-Token Prediction Done Right

What it claims

Standard multi-token prediction (MTP) training sends a noisy gradient into the shared trunk. Jayden Teoh (MIT) shows the MTP cross-entropy gradient at the k-th head equals g_KL + g_res:

• g_KL = Wᵀ(p⁽ᵏ⁾ − q⁽ᵏ⁾) pulls the MTP head toward the model's own position-aligned next-token prediction (NTP) distribution q⁽ᵏ⁾.
• g_res = Wᵀ(q⁽ᵏ⁾ − x) is the NTP cross-entropy gradient against the one-hot target, repeated at the MTP head.

The paper argues that g_res has low signal-to-noise and interferes with both NTP and MTP learning, more so as models grow. The decomposition holds for both parallel (Meta-style) and sequential (DeepSeek-style) MTP heads, provided the LM head W is shared.

How it works

The trick adds a zero-valued straight-through term to the NTP hidden state at each offset position. The NTP logits and NTP loss stay numerically identical. Backprop through the NTP cross-entropy then delivers −g_res into each MTP hidden state, which cancels the residual and leaves only g_KL. The page's pseudocode marks the added line as the only change:

h0[:, k:] += hk.detach() - hk

Explicit KL self-distillation needs T×V logit tensors for both teacher and student and breaks fused cross-entropy kernels. This trick reuses the NTP and MTP cross-entropy computations already present, so it works with fused CE kernels and standard pre-training parallelism.

Gradient routing differs by method:

Soft-MTP keeps data-driven cross-entropy on the LM head and sends self-distillation everywhere upstream. The author says Section 6.3 explains why this beats direct KL minimization, but that section is not in the fetched text.

The paper justifies the trade with a bias-variance argument. Removing g_res removes the irreducible one-hot sampling variance. The cost is a bias toward the model's NTP distribution rather than the data distribution. That bias equals the NTP model's own error, which shrinks with scale and longer training, so the trade should improve as models grow.

Evidence

The fetched text states results only qualitatively, across 167M to 1.2B parameters. Against standard MTP, Soft-MTP gives:

• lower NTP perplexity at every scale;
• faster self-speculative decoding in every domain tested;
• higher zero-shot accuracy at every scale;
• lower MTP perplexity only from 689M upward;
• stronger math and code reasoning after fine-tuning, reported at 1.2B only.

It "matches, if not exceeds" explicit KL-MTP on most metrics. The page reports no measurable compute or memory overhead and no added parameters.

How to use it

Code is at https://github.com/JaydenTeoh/Soft-MTP. Appendix A of the paper has the full PyTorch-style implementation for DeepSeek-style MTP. It includes the λ_MTP/K loss weighting and the per-head rescaling for the different number of targets at each offset; Section B.5 derives that rescaling. Implementations that average each loss over its own target count need that rescaling to recover g_KL exactly.

Limits

The fetched page ends before the experiments section, so no perplexity, acceptance-rate, speedup or accuracy numbers, datasets, token counts or baselines are available here. Experiments stop at 1.2B parameters. The scale argument for larger models is extrapolation. Below 689M, MTP perplexity does not improve over standard MTP. The method requires a linear LM head shared between the NTP and MTP heads. By design it biases MTP heads toward the model's NTP distribution rather than the data.