Soft-MTP: a one-line change recovers the KL gradient buried in standard MTP training — better drafts, same FLOPs, 167M to 1.2B

MIT researcher Jayden Teoh's preprint "Multi-Token Prediction Done Right" argues that standard multi-token prediction (MTP) training already computes the exact gradient from self-distilling each MTP head toward the model's own next-token distribution — it is just buried under a low-signal residual that makes up 77–80% of the gradient reaching the trunk. Soft-MTP, a one-line straight-through routing change, recovers that KL gradient exactly. The paper reports lower next-token loss, faster self-speculative decoding, and higher zero-shot accuracy at every scale from 167M to 1.2B parameters, at no measurable training cost.

What shipped

Multi-token prediction is now a standard pre-training ingredient — DeepSeek-V3/4, Nemotron 3, GLM-5, and Kimi K3 all train extra heads to predict K future tokens, and those heads double as drafters for speculative decoding. The paper identifies a flaw in how those heads are trained. The MTP cross-entropy gradient decomposes exactly into two terms: a KL self-distillation term that pulls each MTP head toward the model's own next-token (NTP) distribution — precisely the alignment that determines speculative-decoding acceptance — and a residual equal to the NTP cross-entropy gradient against the data. The residual has low signal-to-noise ratio and buries the KL signal, increasingly so with scale.

The paper quantifies the imbalance: the residual's SNR is about 4× lower than the KL component's (4.1× at 167M, widening to 4.6× at 1.2B), yet the residual accounts for 77–80% of the MTP gradient reaching the shared trunk.

The fix is the KL gradient-recovery trick: when two distributions share a linear output head, the gradient of their KL divergence can be recovered exactly by combining their cross-entropy gradients through straight-through gradient routing. Unlike minimizing KL directly, it needs no vocabulary-sized logit matrices and stays compatible with the fused cross-entropy kernels used in modern pre-training.

In the released code, standard MTP (hard_mtp) and Soft-MTP (soft_mtp) share the same combined_mtp_loss() function and differ only by an is_soft flag. The core change is one line:

ntp_hidden[:, head:] = ntp_hidden[:, head:] + kl_recovery_scale * (head_hidden.detach() - head_hidden)

The forward pass is numerically unchanged; in the backward pass the negative NTP cross-entropy gradient flows into each MTP head, cancelling the residual and leaving only the KL gradient. The scale factor compensates for the different loss averaging (the 1/T vs 1/(T−k) means and the 1/K head average). The repo includes a pytest that verifies the recovered gradient matches explicit KL(q‖p) minimization to machine precision in float64.

Results, all author-reported from a single-seed run per cell:

• Next-token loss ends lower at all four sizes, with the gap growing with scale: −0.002, −0.004, −0.009, −0.014 at 167M/332M/689M/1.2B.
• MTP loss is higher at 167M and 332M, lower from 689M onward.
• Self-speculative decoding at 1.2B on FineWeb-Edu: speedup over a KV-cached autoregressive baseline rises from 1.74× (standard MTP) to 1.76× (Soft-MTP), with mean acceptance length 1.67 → 1.70. Across the five test domains (fineweb, wiki, books, science, math), Soft-MTP improves average acceptance length by 1.8–2.2% and average speedup by 1.4–1.5% at each scale.
• Zero-shot accuracy (average over LAMBADA, PIQA, HellaSwag, WinoGrande, ARC-easy, ARC-challenge, OpenBookQA, SciQ) is higher at every scale; at 1.2B it is 57.31 vs 56.80.
• After discarding the MTP heads and fine-tuning with NTP loss only, Soft-MTP leads standard MTP in seven of eight code/math cells and never falls below the NTP baseline.

Against explicit KL self-distillation, Soft-MTP matches or exceeds it on most metrics — it wins all six code fine-tuning metrics — with direct KL ahead on both math fine-tuning cells (GSM8K 50.34 vs 49.66; MATH-500 26.00 vs 25.00) and marginally (≤0.01) on FineWeb-Edu self-speculative decoding.

The cost argument is the point for practitioners. At 1.2B with three heads on one H200, the paper reports MTP and Soft-MTP are identical: 10.28 GFLOP/token matmul, 51.7% MFU, 107.0 GB peak memory, and 49,783 vs 49,748 tokens/s throughput (0.07% apart). Direct KL, which materializes full-vocabulary softmaxes, reports 10.43 GFLOP/token, 46,990 tokens/s, and 128.1 GB peak memory.

The code is MIT-licensed, with configs for four sizes (167M→178M, 332M→357M, 689M→745M, 1174M→1273M with K=3 heads), 12 pretrained checkpoints on Hugging Face, and a training setup of 31.5B FineWeb-Edu tokens at 2048 context with the Muon optimizer (PyTorch 2.9+).

Why it matters

The interesting claim is not that self-distillation helps MTP — that was known — but that standard MTP training already computes the self-distillation gradient for free. The residual that buries it is the NTP cross-entropy gradient against the data, and the straight-through routing cancels it exactly using computations the training loop already performs. Within the tested envelope, that is a genuine free lunch: same FLOPs, same memory, same parameters, better drafts.

For engineers training models with MTP heads, the change is trivially adoptable: one flag in the released loss function, no new kernels, no vocab-sized tensors, and a machine-precision test to confirm the gradient identity on your own setup. It also removes the main reason to reach for explicit KL distillation, which costs real memory at scale (128 GB vs 107 GB peak at 1.2B).

The mechanism also explains where the gains show up. Speculative decoding accepts a draft only when it agrees with the verifier, and the paper reads the acceptance gains as the mechanism, not a side effect: training drafts toward the model's own NTP distribution is what raises acceptance length.

The paper's effective-rank analysis of the shared logit map (496 for the NTP reference, 437 for Soft-MTP, 321 for direct KL, 131 for standard MTP) supports the design choice of keeping cross-entropy at the shared head so the NTP readout is preserved — though the paper notes this does not show that is the only cause of the differences.

Caveats

• This is a single-author preprint, not peer-reviewed, and every number above is author-reported with no independent replication.
• The evidence stops at 1.2B parameters, 31.5B training tokens, three MTP heads, and one seed per cell. Whether the trend — the NTP-loss gap and the SNR imbalance both widen with scale — continues at frontier scale is untested.
• Direct KL still wins both math fine-tuning cells, so "matches or exceeds" is not universal.
• "One-line change" is shorthand: the full implementation adds per-head scaling to account for loss weights and differing numbers of supervised targets.
• Removing the residual has a bias cost: it removes one-hot sampling variance but also an NTP-error correction, leaving the upstream gradient biased toward the model's own prediction.
• DeepSeek-style sequential MTP is what the main runs use; Meta-style parallel heads are supported, but only one 332M Meta-style run is reported.
• Specific figures above are drawn from the paper's figures and tables as summarized by alphaXiv's page-cited overview; the paper's own tables were not directly retrievable at the time of writing.

Sources

• arXiv abstract: https://arxiv.org/abs/2610.38331
• arXiv full text: https://arxiv.org/html/2610.38331v1
• GitHub repo: https://github.com/JaydenTeoh/Soft-MTP
• alphaXiv overview: https://www.alphaxiv.org/abs/2610.38331
• Hugging Face collection: https://huggingface.co/collections/JaydenTeoh/soft-multi-token-prediction