This is a Plain English Papers summary of a research paper called Soft-MTP Isolates the KL Signal to Improve Multi-Token Prediction. If you like these kinds of analyses, you can find more AI and machine-learning research on AIModels.fyi or follow us on Twitter.
Soft-MTP isolates the useful part of MTP training
Soft Multi-Token Prediction (Soft-MTP) changes standard MTP training by routing only a KL self-distillation gradient through the MTP heads and shared trunk. In experiments on LLaMA-style models from 167M to 1.2B parameters, it improves next-token perplexity, speculative-decoding acceptance, and zero-shot accuracy over standard MTP, without adding parameters or measurable overhead.
The central technical claim is that a standard MTP cross-entropy gradient splits exactly into two terms: a KL gradient that pulls an MTP prediction toward the model’s own position-aligned next-token distribution, and a residual cross-entropy gradient against the sampled data token. The authors report that the residual has lower signal-to-noise ratio and can interfere with both next-token and multi-token learning.
Soft-MTP recovers the KL gradient by combining existing next-token and MTP cross-entropy gradients through straight-through gradient routing. This avoids constructing vocabulary-sized logit matrices, remains compatible with fused cross-entropy kernels, and is presented as a one-line change to standard MTP code....
