Soft-MTP: Fixing Gradient Noise in Multi-Token Prediction via Cross-Entropy Decomposition

Background: Multi-Token Prediction and Speculative Decoding

Standard next-token prediction (NTP) trains a model to predict the immediate next token using cross-entropy loss L_NTP = −log pθ(yₜ | x₍<t₎). Multi-Token Prediction (MTP) adds auxiliary heads that predict tokens at positions t+1, t+2, t+3, … from a shared Transformer backbone. These extra heads provide longer-range supervision during training and can serve as low-cost draft generators for speculative decoding at inference time. For speculative decoding to work well, the MTP head's output distribution p should closely match the main model's NTP distribution q at the corresponding future position.

Gradient Decomposition in Traditional MTP

The paper analyzes the cross-entropy gradient with respect to the MTP head's hidden representation h (where p = softmax(W h) and W is the shared output projection). For a one-hot target y, the gradient is: g_hard = Wᵀ(p − y) By adding and subtracting the main model's NTP distribution q, this gradient splits exactly into two components:

g_hard = Wᵀ(p − q) + Wᵀ(q − y)
       = g_KL      + g_res

g_KL = Wᵀ(p − q) pushes the MTP distribution p toward the main model's distribution q (equivalent to the gradient of a KL divergence D_KL(q || p)).

g_res = Wᵀ(q − y) is a residual term that depends on the difference between the main model's prediction and the ground-truth token.

The authors argue that g_res injects noise into the MTP head and the shared backbone, and that this noise becomes more pronounced as model size increases (supported by signal-to-noise ratio measurements across 167M–1.2B parameter models).

Why Explicit KL Divergence Is Costly

Computing D_KL(q || p) explicitly requires materializing the full T × V logits matrix (sequence length × vocabulary size). For T = 4,096 and V = 128,000, this matrix contains 524 million elements — over 1 GB in bfloat16 — and would break fused cross-entropy kernels that avoid such large intermediate tensors.

Soft-MTP: Gradient Routing Without Full KL

Soft-MTP achieves the effect of g_KL by manipulating the backward pass only. The forward pass is unchanged. The core idea is a straight-through gradient routing operation on the hidden states: h̃₀ = h₀ + sg(hₖ) − hₖ Here sg denotes stop-gradient. During forward propagation sg(hₖ) − hₖ = 0, so h̃₀ = h₀ numerically. During backpropagation, sg(hₖ) contributes no gradient, while −hₖ sends a negative gradient to the MTP representation hₖ. With proper routing and normalization, this cancels g_res from the gradients flowing to the MTP heads and the shared backbone, leaving only g_KL. The shared output projection W continues to receive the original cross-entropy gradient g_hard.

Implementation Details

The paper provides pseudocode (simplified):

h0[:, k:] = h0[:, k:] + s_k * (hk.detach() − hk)
ntp_loss = ce(lm_head, h0, y)
mtp_mean = mean(mtp_losses)
total_loss = ntp_loss + λ_MTP * mtp_mean

The actual implementation handles token-position alignment across heads and includes normalization factors for the number of supervised positions, number of heads, and loss weights. The official GitHub repository contains gradient tests that match the analytical KL gradient.

Experimental Setup

Four LLaMA-style models: 167M, 332M, 689M, 1.2B parameters.

Three DeepSeek-style MTP heads.

315B FineWeb-Edu tokens, Muon optimizer.

Baselines: standard NTP, standard MTP, explicit KL-MTP, Soft-MTP.

Throughput measured on a single NVIDIA H200 with the 1.2B model.

Results

Speculative Decoding

Evaluated on five domains (FineWeb-Edu, Wikipedia, Books, Science, Math) using three MTP heads. Soft-MTP consistently outperforms standard MTP across all model sizes and domains. For the 1.2B model on FineWeb-Edu:

Standard MTP speedup: 1.74×

Soft-MTP speedup: 1.76× (≈1.1% relative improvement)

Average relative improvement across domains: 1.4%–1.5%.

Language Modeling

On the 1.2B model, average zero-shot accuracy on reported benchmarks rises from 56.80% (standard MTP) to 57.31% (Soft-MTP), a gain of 0.51 percentage points . Next-token perplexity also improves.

Downstream Code and Math Tasks (MTP Heads Removed)

After pre-training, the 1.2B models were fine-tuned on code/math data with standard NTP (MTP heads removed). Soft-MTP pre-training yields:

HumanEval pass@64: 40.85% → 43.90%

MATH-500 accuracy: 23.20% → 25.00%

This suggests the gradient adjustment changes the shared backbone's representations in a way that transfers to downstream tasks even without MTP heads.

Training Overhead

On a single H200 with the 1.2B model:

Standard MTP: 49,783 tokens/sec

Explicit KL-MTP: 46,990 tokens/sec

Soft-MTP: 49,748 tokens/sec

Soft-MTP matches standard MTP throughput (difference ≈0.07%) and peak memory, while explicit KL-MTP is ~5.6% slower.

Limitations

Largest model tested: 1.2B parameters — far smaller than frontier models.

Results from a single preprint (arXiv:2610.38331) without independent large-scale replication.

Gains are modest; the method does not double inference speed.

Conclusion

Soft-MTP shows that decomposing the cross-entropy gradient and canceling the residual component via a simple backward-pass routing operation can improve both language modeling quality and speculative decoding efficiency, with negligible training overhead and no added parameters. The work illustrates how inspecting the gradient structure of existing objectives can reveal removable noise sources.

Paper: https://arxiv.org/abs/2610.38331