Summary (Overview)

• MTP-D Framework: The paper introduces MTP-D, a self-distillation method for Multi-Token Prediction (MTP) in Large Language Models (LLMs) that improves MTP head acceptance rates by +7.5% while preserving main-head performance.
• Looped Extension Strategy: A novel approach for cost-efficient MTP head expansion via continued pre-training, enabling scaling from 4 to 16 MTP heads with further 35.1% speedup.
• Key Techniques: Gradient-detached, TopN-logits-selected distillation from the main head to MTP heads, using KL divergence loss with stop-gradient operations.
• Empirical Results: Achieves up to 220.4% inference speedup over 1-head MTP baseline, validated on 2B Dense and A1B MoE models across seven benchmarks.
• Scalability Insights: Six key insights on MTP scalability, including inherent cascade architecture scalability, superior performance of distillation-based approaches, and data efficiency in looped extension.

Introduction and Theoretical Foundation

Background

Large Language Models (LLMs) demonstrate strong performance across diverse tasks but face inference efficiency bottlenecks. Traditional Next-Token Prediction (NTP) paradigms perform autoregressive, token-by-token generation, incurring high latency and computational costs, particularly for long sequences.

Motivation

Multi-Token Prediction (MTP) extends NTP by training LLMs with multiple heads for parallel future token prediction. However, existing approaches face two critical challenges:

1. Limited acceptance rates of MTP heads, causing exponential decline in cumulative acceptance rate
2. Difficulty in jointly training multiple main and MTP heads due to the "seesaw effect"

Theoretical Foundation

The cascaded MTP architecture (DeepSeek-V3) models inter-token dependencies while preserving the complete causal chain. The MTP loss is formulated as:

$$\mathcal{L}_{\mathrm{mtp}}^{\mathrm{CE}} = \sum_{k=1}^{K} \alpha_{k} \mathcal{L}_{\mathrm{mtp}_{\mathrm{k}}}^{\mathrm{CE}} = \sum_{k=1}^{K} \alpha_{k} \operatorname{CE}\left(\hat{\mathbf{P}}_{k+1:T+1}^{k}, \mathbf{t}_{k+1:T+1}\right) \tag{1}$$

where $\mathrm{CE}(\cdot)$ is cross-entropy loss, $\alpha_{k}$ is the weight coefficient, $\mathbf{t}$ is the ground-truth token sequence, and $\hat{\mathbf{P}}$ is the predicted probability distribution.

Methodology

Self-Distillation for MTP (MTP-D)

The core method introduces a unidirectional KL divergence loss for self-distillation from the TopN logits of the main head to MTP heads:

$$\mathcal{L}_{\mathrm{mtp}}^{\mathrm{KL}} = \sum_{k=1}^{K} \beta_{k} \mathcal{L}_{\mathrm{mtp}_{\mathrm{k}}}^{\mathrm{KL}} = \sum_{k=1}^{K} \beta_{k} \mathrm{KL}\left(\tilde{\mathbf{P}}_{k+1:T+1}^{k}, \mathrm{sg}(\tilde{\mathbf{Q}})_{k+1:T+1}\right) \tag{2}$$

where:

• $\mathcal{I}_{t}^{N} = \operatorname{TopK}\left(\hat{\mathbf{Q}}_{t}, N\right)$ (Equation 3) selects TopN indices
• $\tilde{\mathbf{Q}}_{k+1:T+1} = \sigma\left(\hat{\mathbf{Q}}_{k+1:T+1}\left[\dots, \mathcal{I}_{t}^{N}\right]\right)$ (Equation 4) applies softmax
• $\tilde{\mathbf{P}}_{k+1:T+1}^{k} = \log\left(\sigma\left(\hat{\mathbf{P}}_{k+1:T+1}^{k}\left[\dots, \mathcal{I}_{t}^{N}\right]\right)\right)$ (Equation 5) applies log-softmax

Key Design Choices:

1. Gradient-detached self-distillation: Stop-gradient operation on main head logits $\hat{Q}$ prevents gradient propagation back through the main head, preserving its performance
2. TopN-selected logits: Selecting TopN = 10,000 tokens (from vocabulary of 122,880) ensures efficient, stable distillation while avoiding long-tail distribution issues
3. Final loss: $\mathcal{L}_{\mathrm{mtp}} = \mathcal{L}_{\mathrm{mtp}}^{\mathrm{CE}} + \mathcal{L}_{\mathrm{mtp}}^{\mathrm{KL}}$, combining ground-truth alignment with knowledge distillation

Looped MTP Head Extension

The looped extension strategy:

1. Takes trained MTP heads as a group
2. Copies weights to initialize the next group of heads
3. Updates via continued pre-training with frozen main model and previously trained heads

Training-Free Observations:

• MTP heads at loop connection points show noticeable but acceptable acceptance rate drops
• MTP-D maintains 26.70% cumulative acceptance rate at head 3 vs 0.6% for standard MTP in 1-to-8 loop settings

Empirical Validation / Results

Experimental Setup

• Models: 2B Dense and N10BA1B MoE LLMs
• Data: FineWeb-Edu-350BT dataset (350B tokens pre-training, 70B tokens looped extension)
• Hardware: 256 NVIDIA H20 GPUs
• Benchmarks: AGIEval-en, GSM8K, MATH, NaturalQuestions, SimpleQA, SuperGPQA, TriviaQA

Main Results

Key Results from Table 1 (selected):

Performance Improvements:

• K=1: MTP-D improves acceptance rate by 3.09% (2B Dense) and 4.11% (A1B MoE)
• K=4: 3.91% (2B Dense) and 4.73% (A1B MoE) improvement for 4th head
• Main-head performance preserved: 11.68 avg accuracy (MTP-D) vs 11.28 (MTP) for K=1
• Speedup: 22.9% for 4-head, up to 107.4% vs single-head, 220.4% with looped extension to 16 heads

Ablation Study Results

Key findings from Table 2:

Critical observations:

• Removing detach increases main-head loss (+0.079) and reduces performance (-1.49%)
• TopN=10,000 provides optimal trade-off between efficiency and performance
• β=1.0 balances main-head preservation and MTP improvement
• Forward KL achieves best loss-performance trade-off

Looped Extension Results

Six key insights emerged:

1. Cascaded MTP is inherently scalable due to structural consistency
2. MTP-D shows superior scalability vs standard MTP
3. Grouped MTP heads exhibit stronger loop scalability than single heads
4. Limited data suffices: 70B vs 350B tokens yield marginal differences
5. Ensemble teachers enhance consistency and loop scalability
6. MTP-D supports up to 16 heads with 5-10% CAR maintained

Theoretical and Practical Implications

Theoretical Contributions

• Demonstrates that self-distillation with gradient detachment effectively decouples MTP head improvement from main-head degradation
• Establishes that TopN-selected logits provide sufficient supervision signal despite long-tail distribution
• Reveals inherent scalability of cascaded MTP architectures through structural consistency

Practical Implications

• Inference Acceleration: Significant speedups (up to 220.4%) enable practical deployment of MTP in production LLMs
• Cost Efficiency: Looped extension with minimal continued pre-training (70B tokens) achieves substantial gains
• Architecture Flexibility: Method works across dense and MoE architectures, various model sizes
• Industrial Adoption: Provides guidance for implementing MTP in large-scale LLM pre-training pipelines

Conclusion

The paper presents MTP-D, a simple yet effective self-distillation framework for multi-token prediction in LLMs, along with a looped extension strategy. Key contributions include:

1. MTP-D framework: Gradient-detached, TopN-selected self-distillation that improves MTP head acceptance rates by 7.5% while maintaining main-head performance
2. Looped extension: Cost-efficient scaling from 4 to 16 MTP heads with substantial inference speedups
3. Comprehensive insights: Six validated insights on MTP scalability and optimization strategies

Future Directions:

• Adaptation to post-training scenarios
• Investigation of optimal $\alpha_k$ and $\beta_k$ relationships with K (number of MTP heads)
• Validation on diverse datasets and ultra-large models
• Exploration of theoretical foundations for optimal hyperparameter selection

The work provides valuable guidance for improving LLM pre-training and inference efficiency, with practical implications for industrial-scale deployment of MTP.