Background
Multi-Token Prediction (MTP) is now a common ingredient in language-model pre-training, as it improves downstream performance and enables faster inference through speculative decoding.
Problem Identification
We identify a flaw in standard MTP training, in which a residual component of the MTP gradient interferes with both next-token and multi-token learning, increasingly so with scale. Specifically, the MTP cross-entropy gradient can be decomposed exactly into:
- A KL self-distillation term that pulls each MTP head toward the model's own next-token prediction (NTP) distribution and improves speculative decoding,
- A residual that has a low signal-to-noise ratio and buries the useful KL signal.
Proposed Solution: Soft-MTP
We introduce 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 the KL directly, this trick requires no vocabulary-sized logit matrices and is compatible with standard fused cross-entropy kernels. This enables Soft Multi-Token Prediction (Soft-MTP), a one-line code change to MTP that removes the noisy residual and backpropagates only the KL self-distillation signal by reusing the existing NTP and MTP cross-entropy computations.
Soft-MTP supports both DeepSeek- and Meta-style MTP, adds no parameters, and incurs no measurable overhead.
Experimental Results
From 167M to 1.2B parameters, Soft-MTP improves over standard MTP with:
- Lower NTP perplexity,
- Faster speculative decoding and higher zero-shot accuracy at every scale,
- Lower MTP perplexity from 689M onward,
- Stronger math and code reasoning after fine-tuning at 1.2B.
Soft-MTP also matches, if not exceeds, direct KL divergence minimization on most metrics.
Soft-MTP is, in effect, a free lunch for MTP optimization.
Our code can be found at https://github.com/JaydenTeoh/Soft-MTP.