Glonce

Soft-MTP fixes multi-token prediction training with a one-line change

researchThe Glonce editorsHugging Face🔥61Glonce rating61Impact45Novelty65Relevance100Surprise35Players10Credibility70
Soft-MTP fixes multi-token prediction training with a one-line change

Multi-Token Prediction (MTP) is now a common ingredient in language-model pre-training. It improves downstream performance and enables faster inference through speculative decoding. A new paper, listed on Hugging Face papers, argues that the standard way of training it has a flaw.

The authors show that the MTP cross-entropy gradient can be decomposed exactly into two parts. The first is a KL self-distillation term that pulls each MTP head toward the model's own next-token prediction (NTP) distribution; this is the part that improves speculative decoding. The second is a residual with a low signal-to-noise ratio that buries the useful KL signal. According to the paper, this residual interferes with both next-token and multi-token learning, and does so increasingly as models grow.

The obvious fix is to minimize the KL directly, but that needs vocabulary-sized logit matrices. The paper instead introduces what it calls 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 direct KL minimization, the trick needs no vocabulary-sized logit matrices and works with standard fused cross-entropy kernels.

That yields Soft Multi-Token Prediction (Soft-MTP). It is a one-line code change to MTP that removes the noisy residual and backpropagates only the KL self-distillation signal, reusing the NTP and MTP cross-entropy computations that already exist. The authors say it supports both DeepSeek-style and Meta-style MTP, adds no parameters, and has no measurable overhead.

The reported results cover models from 167M to 1.2B parameters. Against standard MTP, Soft-MTP gives lower NTP perplexity, faster speculative decoding and higher zero-shot accuracy at every scale tested. MTP perplexity is lower from 689M parameters onward. At 1.2B, the model shows stronger math and code reasoning after fine-tuning. Soft-MTP also matches, if not exceeds, direct KL divergence minimization on most metrics. The authors call it, in effect, a free lunch for MTP optimization. Code is on GitHub at JaydenTeoh/Soft-MTP.

Key facts

  • The paper says the MTP cross-entropy gradient splits exactly into a useful KL self-distillation term and a low signal-to-noise residual that interferes with learning, more so at larger scale.
  • Soft-MTP is a one-line change that keeps only the KL term, using a KL gradient-recovery trick built on straight-through gradient routing.
  • Unlike direct KL minimization, it needs no vocabulary-sized logit matrices and works with standard fused cross-entropy kernels.
  • From 167M to 1.2B parameters it beats standard MTP on NTP perplexity, speculative decoding speed and zero-shot accuracy at every scale; MTP perplexity improves from 689M onward.
  • It supports DeepSeek- and Meta-style MTP, adds no parameters and has no measurable overhead, according to the authors.

Why it matters

MTP is described as a common ingredient in language-model pre-training, valued for better downstream performance and faster inference through speculative decoding. If the standard loss really carries a noisy residual that grows with scale, a fix that costs one line of code is notable. The paper's framing is that the useful part of the gradient is the KL self-distillation term, and the rest mostly gets in the way.

Who it affects

Teams that pre-train language models with MTP heads, in either the DeepSeek style or the Meta style, are the direct audience. Anyone relying on speculative decoding with MTP heads is also in scope, since the paper reports faster speculative decoding.

How to use it

The authors released code at https://github.com/JaydenTeoh/Soft-MTP. The change is described as one line in an MTP setup. It reuses the existing NTP and MTP cross-entropy computations, adds no parameters, and is compatible with standard fused cross-entropy kernels, so it should not require a different training stack.

How solid is it

The decomposition of the gradient is presented as exact, and the gradient-recovery trick as exact when two distributions share a linear output head. The empirical results are self-reported by the authors and cover only models from 167M to 1.2B parameters. The abstract gives no numeric size for the gains, and names no benchmarks or datasets. The comparison with direct KL minimization holds on most metrics, not all. No independent replication is mentioned.

Risks and caveats

Nothing is reported above 1.2B parameters, so it is unclear whether the gains hold at the scales where frontier models are trained. MTP perplexity improves only from 689M onward, and the math and code reasoning gain is reported at 1.2B only. The claim that the residual worsens with scale rests on the paper's own analysis. Phrases like free lunch are the authors' own framing.

“Soft-MTP is, in effect, a free lunch for MTP optimization.”

— Paper abstract, Hugging Face papers 2610.38331