Soft-MTP Recovers the Useful Signal in Multi-Token Prediction

Introduction

Multi-Token Prediction (MTP) has become a familiar component of language-model pre-training. Instead of supervising the model only on the next token, MTP adds heads that predict tokens farther ahead. This can improve downstream behavior and produce stronger draft predictions for speculative decoding. Yet the benefit depends not only on adding prediction heads, but also on the quality of the gradients sent through the shared model trunk.

The problem with standard MTP

The study identifies a hidden weakness in the usual MTP objective. Its cross-entropy gradient can be exactly decomposed into two parts. One is a KL-divergence self-distillation signal: it encourages each MTP head to match the model’s own next-token prediction distribution. This alignment is useful for speculative decoding because the draft heads become more consistent with the main model. The other part is a residual with a low signal-to-noise ratio.

That residual is not harmless. It can interfere with both next-token learning and multi-token learning, and the interference becomes more pronounced as model scale increases. In other words, adding the standard MTP loss may also send noisy supervision into the shared trunk instead of isolating the part of the objective that is most useful.

How Soft-MTP works

The proposed solution is a KL gradient-recovery trick. NTP and MTP share a linear output head, so their existing cross-entropy gradients contain enough information to reconstruct the gradient of the KL term. With straight-through gradient routing, the method combines those gradients and retains the self-distillation component without explicitly materializing vocabulary-sized logit matrices.

Soft-MTP is consequently described as a one-line change to standard MTP. It reuses the existing NTP and MTP cross-entropy computations, remains compatible with fused cross-entropy kernels, adds no parameters, and introduces no measurable overhead. The approach works with both DeepSeek-style and Meta-style MTP. The reported evaluations cover models from 167M to 1.2B parameters and show improvements over standard MTP.

Why it matters

The broader lesson is that better auxiliary training may come from filtering gradients rather than adding more capacity. If the most valuable role of MTP is to make future-token heads imitate a distribution already produced by the main model, then preserving that KL signal may be more effective than backpropagating every component of the original loss.

The engineering implications are also notable. Soft-MTP does not require a new architecture, additional trainable parameters, or a replacement for optimized cross-entropy kernels. Teams with an existing MTP pre-training pipeline could therefore treat it as a lightweight objective-level modification. The available material does not include the complete experimental tables, so results at larger scales, with more prediction heads, or across additional tasks still require verification from the full paper.

Source: Hugging Face Daily Papers