Multi-Token Prediction Done Right

Soft-MTP removes the noisy residual from the standard multi-token prediction gradient by recovering the KL self-distillation gradient exactly from existing cross-entropy computations, improving NTP, speculative decoding, and zero-shot accuracy without added cost.

Standard MTP training's cross-entropy gradient contains a residual component that interferes with next-token and multi-token learning and becomes increasingly noisy with scale. The useful KL self-distillation signal is buried by this low-SNR residual. Direct KL minimization is costly because it requires vocabulary-sized logit matrices and is incompatible with fused cross-entropy kernels.

The paper decomposes the hard MTP gradient exactly into a KL self-distillation term and a residual term. It introduces the KL gradient-recovery trick: by adding stop-gradient routed MTP hidden states to the NTP hidden states, the augmented NTP loss injects minus the residual gradient at the MTP hidden state. Summing this with the MTP cross-entropy loss leaves exactly the KL gradient at the MTP head. Soft-MTP applies this routing so the output head still receives cross-entropy gradients while MTP heads and shared trunk receive only the recovered KL gradient. It supports DeepSeek- and Meta-style MTP and adds no parameters.

• At 1.2B, Soft-MTP lowers teacher-forced KL divergence and TV distance of the MTP distribution from the model's own NTP verifier compared with standard MTP. §6.2
• Across 167M to 1.2B parameters, Soft-MTP wins every metric in the self-speculative decoding table at every scale. §6.2
• Soft-MTP matches standard MTP in training cost: identical matmul FLOPs, MFU, and fused cross-entropy peak memory, with throughput within 0.07%. §6.3
• The residual carries about 80% of the MTP gradient in the shared trunk at every scale, despite having substantially lower signal-to-noise ratio than the KL component. §6.1
• With Meta-style MTP heads at 332M, K=3, lambda_MTP=1, and 7.3B FineWeb-Edu tokens, Soft-MTP attains lower NTP and MTP loss than standard MTP. §C

• Soft-MTP is biased toward the NTP distribution rather than the data distribution, as the paper states in its bias-variance view; this is desirable for speculative decoding but is a stated tradeoff.
• The experiments cover 167M to 1.2B parameters; the paper discusses frontier-scale relevance in §7 but reports no frontier-scale experiments, leaving that regime empirically open.
• Meta-style MTP is tested only at 332M in Appendix C, so architecture coverage beyond DeepSeek-style MTP is limited in the experiments.
• The paper reports lower MTP perplexity only from 689M onward, so the MTP-loss benefit is not established at the smaller 167M and 332M scales.

• The hard MTP gradient decomposes exactly into a KL self-distillation component and a residual component: g_hard = g_KL + g_res. §4
• The g_KL component equals the gradient of the forward KL divergence from a stop-gradient aligned NTP distribution to the MTP distribution. §4.1
• The KL gradient-recovery trick augments NTP hidden states with routed MTP hidden states so the augmented NTP loss injects minus g_res at the MTP hidden state; summing with MTP loss leaves exactly g_KL. §5.1
• Soft-MTP sends cross-entropy gradients into the shared output head but sends the recovered KL gradient into MTP heads and the shared trunk, avoiding vocabulary-sized logit matrices. §5.2
• The residual contains NTP prediction error plus one-hot sampling noise; Soft-MTP removes the sampling variance at the cost of bias toward the NTP distribution. §5.3
• Soft-MTP supports both DeepSeek- and Meta-style MTP, adds no parameters, and incurs no measurable overhead. §Abstract
• Soft-MTP matches standard MTP in FLOPs, MFU, and fused cross-entropy peak memory, and is more efficient than direct KL-MTP on every reported cost axis. §6.3
• Soft-MTP improves NTP perplexity, zero-shot accuracy, and self-speculative decoding at every tested scale from 167M to 1.2B. §6.2