# Self-Distillation for Multi-Token Prediction

> MTP-D, a gradient-detached self-distillation method with TopN logits, boosts multi-token prediction acceptance rates by 7.5% and enables cost-efficient scaling to 16 heads for up to 220% faster inference.

- **Source:** [arXiv](https://arxiv.org/abs/2603.23911)
- **Published:** 2026-08-15
- **Permalink:** https://picx.dev/p/swvAfv
- **Whiteboard:** https://picx.dev/p/swvAfv/image

## Summary

## 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):**

| Model | Method | K | Head 1 CAR/AR | Head 4 CAR/AR |
|-------|--------|---|---------------|---------------|
| 2B Dense | MTP | 4 | 81.88/81.88 | 45.47/83.77 |
| 2B Dense | MTP-D | 4 | 85.86/85.86 | 52.96/86.46 |
| A1B MoE | MTP | 4 | 82.38/82.38 | 48.82/85.27 |
| A1B MoE | MTP-D | 4 | 85.67/85.67 | 52.47/85.25 |

**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:**

| Strategy | Main-head Accuracy | MTP-head AR | Mean |
|----------|-------------------|-------------|------|
| MTP-D (default) | 11.68 | 90.06 | - |
| No detach | 10.19 | 94.34 | - |
| TopN=1 | 11.47 | 87.39 | - |
| TopN=1000 | 11.47 | 90.47 | - |
| β=0.1 | 10.99 | 87.66 | - |
| β=1.5 | 11.64 | 90.80 | - |
| Reverse KL | 11.62 | 90.22 | - |
| Hybrid KL | 10.99 | 88.92 | - |

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.

---

_Markdown view of https://picx.dev/p/swvAfv, served by PicX — AI-generated visual whiteboard summaries of research papers._