Multi-Token Prediction Done Right (Soft-MTP) – Visual Summary

Summary (Overview)

Key contribution: Identifies a flaw in standard Multi-Token Prediction (MTP) training, where a residual component of the MTP gradient interferes with both next-token and multi-token learning, increasingly so with model scale.
The KL gradient–recovery trick: A novel technique that recovers the exact KL self-distillation gradient from cross-entropy gradients when two distributions share a linear output head, using straight-through gradient routing—without materializing vocabulary-sized logit matrices.
Soft-MTP: A one-line code change to standard MTP that removes the noisy residual and backpropagates only the useful KL self-distillation signal, requiring no additional parameters and incurring no measurable overhead.
Empirical results: From 167M to 1.2B parameters, Soft-MTP improves over standard MTP with lower NTP perplexity, faster speculative decoding, higher zero-shot accuracy at every scale, and stronger math and code reasoning after fine-tuning.
Efficiency: Soft-MTP matches or exceeds explicit KL minimization on most metrics while remaining compatible with fused cross-entropy kernels and standard parallelism schemes used in frontier pre-training.

Introduction and Theoretical Foundation

Background and Motivation

Next-token prediction (NTP) has become the standard recipe for training large language models (LLMs). However, Multi-Token Prediction (MTP)—training models to predict multiple future tokens—has been widely adopted in modern LLM training, including in DeepSeek-V3/4, Nemotron 3, GLM-5, and Kimi K3. MTP offers two advantages:

Better model performance: Alleviates the myopic nature of next-token supervision, leading to stronger performance on long-horizon domains like mathematics and code.
Faster inference: MTP heads can be used for speculative decoding.

Despite its widespread adoption, little was understood about how MTP interacts with NTP and which objective is best for training MTP heads. For speculative decoding, a token drafted by MTP is accepted only if it agrees with the model's own next-token prediction—suggesting that MTP heads should be trained to match the NTP distribution via self-distillation. However, naive KL minimization requires materializing vocabulary-sized logit tensors, which is memory-intensive and incompatible with standard fused cross-entropy kernels.

Theoretical Foundation

Consider a model trunk with parameters $\theta$ mapping prefixes $x_{1:t}$ to hidden states $h^{(0)}_t \in \mathbb{R}^d$ . A shared linear output head $W_\theta \in \mathbb{R}^{V \times d}$ projects hidden states to logits. The next-token distribution is:

p^{(0)}_t = \text{softmax}\left(W_\theta h^{(0)}_t\right)

The $k$ -th MTP head produces hidden states $h^{(k)}_t$ to predict token $x_{t+k+1}$ :

p^{(k)}_t = \text{softmax}\left(W_\theta h^{(k)}_t\right)

The full MTP objective is:

\mathcal{L}_{MTP} = \mathcal{L}_{NTP} + \lambda_{MTP} \cdot \frac{1}{K}\sum_{k=1}^{K}\mathcal{L}^{(k)}_{MTP}

Methodology

Gradient Decomposition

The paper decomposes the hard MTP gradient (gradient of standard MTP cross-entropy loss with respect to the $k$ -th MTP hidden state):

g_{\text{hard}} := \nabla_{h^{(k)}_t} \text{CE}\left(x_{t+k+1}, p^{(k)}_t\right) = W^\top_\theta\left(p^{(k)}_t - x_{t+k+1}\right) \tag{1}

This decomposes exactly into:

g_{\text{hard}} = \underbrace{W^\top_\theta\left(p^{(k)}_t - q^{(k)}_t\right)}_{g_{KL}} + \underbrace{W^\top_\theta\left(q^{(k)}_t - x_{t+k+1}\right)}_{g_{\text{res}}} \tag{2}

where $q^{(k)}_t := p^{(0)}_{t+k}$ is the position-aligned NTP distribution (Figure 2 illustrates this alignment).

$g_{KL}$ corresponds to the gradient of the forward KL divergence from a stop-gradient teacher:

\nabla_{h^{(k)}_t} D_{KL}\left(\text{sg}(q^{(k)}_t) \| p^{(k)}_t\right) = W^\top_\theta\left(p^{(k)}_t - q^{(k)}_t\right) = g_{KL} \tag{3}

$g_{\text{res}}$ is the NTP cross-entropy gradient against the data—a residual that has low signal-to-noise ratio.

KL Gradient–Recovery Trick

The core innovation augments NTP hidden states with straight-through gradient routing:

\tilde{h}_t = h^{(0)}_t + \sum_{\substack{k=1 \\ t-k \geq 1}}^{K}\left[\text{sg}(h^{(k)}_{t-k}) - h^{(k)}_{t-k}\right] \tag{4}

Since $\tilde{h}_{t+k}$ is numerically identical to $h^{(0)}_{t+k}$ , the augmented NTP loss produces the same NTP gradients. However, it injects an extra gradient with respect to the MTP representation:

\nabla_{h^{(k)}_t}\tilde{\mathcal{L}}_{NTP} = -W^\top_\theta\left(q^{(k)}_t - x_{t+k+1}\right) = -g_{\text{res}} \tag{6}

Summing the contributions yields the pure self-distillation gradient:

\nabla_{h^{(k)}_t}\left(\tilde{\mathcal{L}}_{NTP} + \mathcal{L}^{(k)}_{MTP}\right) = g_{\text{hard}} - g_{\text{res}} = g_{KL} \tag{7}

Soft-MTP Implementation

The paper provides a minimal code change (see Figure 1):

z = shared_trunk(x)
h0 = ntp_head(z)
for k in range(1, K+1):
    hk = mtp_head[k-1](z, x)
    loss_mtp += ce(lm_head, hk, y[k])

# straight-through gradient routing
h0[:, k:] += hk.detach() - hk  # h0 is numerically unchanged

loss_ntp = ce(lm_head, h0, y[0])
(loss_ntp + loss_mtp).backward()

Bias–Variance View

The residual splits further into NTP error and one-hot sampling noise:

g_{\text{res}} = \underbrace{W^\top_\theta\left(q^{(k)}_t - \pi^{(k)}_t\right)}_{\text{NTP error}} + \underbrace{W^\top_\theta\left(\pi^{(k)}_t - x_{t+k+1}\right)}_{\text{one-hot sampling noise}} \tag{8}

The NTP error shrinks with scale, but the one-hot sampling noise is irreducible—making Soft-MTP's tradeoff increasingly favorable at scale. The gradient routing across methods is summarized in Table 1.

Method	Output head $W_\theta$	MTP heads and shared trunk
Standard MTP	CE gradient	CE gradient
KL-MTP	KL gradient	KL gradient
Soft-MTP	CE gradient	KL gradient

Table 1: Gradient routing under different MTP objectives.

Empirical Validation / Results

Experimental Setup

LLaMA-style models from 167M to 1.2B parameters
DeepSeek-style sequential MTP architecture with $K=3$ heads
Trained on FineWeb-Edu for 60,000 steps with 0.5M tokens/batch, 2048 sequence length (~31.5B tokens)
Muon optimizer with update-RMS rescaling

1. The Residual Interferes with Learning

Key findings (Figure 3):

The residual has ~4× lower SNR than the KL component throughout training
The SNR gap widens with scale: from 4.1× at 167M to 4.6× at 1.2B
Despite low SNR, the residual accounts for ~80% of the MTP gradient reaching the shared trunk

The residual is anti-aligned with NTP optimization (Figure 4a): steps along $g_{KL}$ decrease NTP loss, while steps along $g_{\text{res}}$ increase it. Soft-MTP achieves lower NTP loss, with the gap widening at scale, and lower MTP loss from 689M onward.

2. Soft-MTP Outperforms Standard MTP

Speculative decoding (Table 2): Soft-MTP improves acceptance length by 1.8–2.2% and speedup by 1.4–1.5% at every scale, across FineWeb-Edu, Wikipedia, Books, Science, and Math domains.

Zero-shot downstream performance (Table 3): Soft-MTP improves average accuracy over standard MTP at every scale, with gains growing from ~0.1 points at 167M/332M to ~0.5 points at 689M/1.2B. At 1.2B, Soft-MTP surpasses NTP in average accuracy and achieves best WikiText and LAMBADA perplexity.

Math and code fine-tuning (Table 4): Soft-MTP leads standard MTP in seven of eight cells and is the only objective that never falls below the NTP baseline:

Code: +3.1 percentage points on HumanEval at pass@64
Math: +1.1 points on GSM8K and +1.8 points on MATH-500

Hyperparameter robustness (Figure 5): Soft-MTP Pareto-dominates standard MTP across MTP weights $\lambda_{MTP} \in \{0.25, 0.5, 0.75, 1.0\}$ and horizons $K \in \{3, 4, 5, 6, 7, 8\}$ .

3. Soft-MTP vs. KL-MTP Efficiency

Training cost (Table 5): Soft-MTP has identical FLOPs, MFU, and fused-CE peak memory to standard MTP (throughput within 0.07%), while KL-MTP adds 1.5% FLOPs, lowers throughput by 5.6%, and raises peak memory by 21 GB.

Performance comparison (Table 6): Soft-MTP ties or wins on every perplexity metric, has higher average zero-shot accuracy, and beats KL-MTP on all six code fine-tuning metrics. KL-MTP leads in math fine-tuning.

Why keep CE at the output head: KL-MTP's self-distillation gradient fits the head to the model's own predictions, creating a shortcut where the KL can be reduced by making the readout less sensitive to MTP/NTP hidden state differences. The effective rank of the logit map shows: standard MTP collapses to 131 vs. NTP baseline 496, KL-MTP to 321, while Soft-MTP stays closest at 437.

Theoretical and Practical Implications

Theoretical Implications

Gradient decomposition insight: The MTP cross-entropy gradient inherently contains a self-distillation signal, providing a new theoretical understanding of why MTP improves speculative decoding.
Bias-variance tradeoff: Removing the residual trades unbiasedness for variance reduction. This tradeoff becomes increasingly favorable at scale because the bias term (NTP error) shrinks while the variance term (one-hot sampling noise) is irreducible.
KL gradient–recovery is general: The trick applies more broadly whenever a teacher and student share a linear output head—including post-training methods like MTP Boosting, on-policy self-distillation, and score centering for off-policy RL.

Practical Implications

Free lunch for MTP optimization: Soft-MTP is a one-line code change with no measurable overhead, compatible with fused cross-entropy kernels and standard parallelism (data parallelism, ZeRO, expert parallelism, pipeline parallelism).
Compatibility with parallelism: Only sequence sharding (ring attention, sequence parallelism) needs extra communication—a backward-pass exchange of at most $K$ hidden-state gradients per shard boundary, independent of sequence length.
Frontier-scale relevance: The trick avoids vocabulary-sized tensors, making it preferable to explicit self-distillation, especially for models trained without tensor parallelism (e.g., DeepSeek-V3).

Conclusion

The paper demonstrates that the standard MTP cross-entropy gradient can be decomposed exactly into:

$g_{KL}$ : A KL self-distillation term pulling MTP heads toward the model's own next-token distribution
$g_{\text{res}}$ : A residual carrying most of the gradient but with low signal-to-noise ratio, working against next-token learning

The KL gradient–recovery trick removes this residual by recovering the KL term exactly from existing cross-entropy gradients. Soft-MTP implements this with a one-line code change that:

Adds no parameters
Incurs negligible overhead
Remains compatible with fused cross-entropy kernels
Works with both DeepSeek-style and Meta-style MTP architectures

Future directions: The paper questions whether MTP will remain in frontier pre-training (noting DeepSeek-V4.1-Flash dropped it for post-trained drafter DSpark), but emphasizes that the KL gradient–recovery trick applies broadly beyond MTP—wherever a teacher and student share a linear output head, the forward KL gradient can be recovered from cross-entropy gradients, opening new avenues for efficient self-distillation methods in post-training and reinforcement learning contexts.

Summary