Soft-MTP Removes Noisy Residuals in Multi-Token Prediction Training

·2026.11.04 15:46

Key point

A new method called Soft-MTP improves standard MTP by filtering out noisy gradients, resulting in lower perplexity and faster speculative decoding from 167M to 1.2B parameters.

Details

Standard Multi-Token Prediction (MTP) training suffers from a residual gradient component that interferes with learning and buries the useful KL self-distillation signal.
The proposed Soft-MTP method uses a KL gradient-recovery trick to backpropagate only the KL signal, requiring no vocabulary-sized logit matrices and no additional parameters.
Experiments from 167M to 1.2B parameters show Soft-MTP achieves lower NTP perplexity, faster speculative decoding, and higher zero-shot accuracy compared to standard MTP.
The approach supports both DeepSeek- and Meta-style MTP architectures and matches or exceeds direct KL divergence minimization on most metrics.
Code is available on GitHub under the name Soft-MTP.

speculative-decoding multi-token-prediction llm-training arxiv-paper gradient optimization

Read originalShare on Threads

This summary was generated automatically by AI. Check the original for the author's claims and context. Copyright belongs to the original author.

Our guide explains how the AI works. Report summary errors, attribution issues, or removal requests via Contact.