Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads
cs.LG, cs.AI
Submitted: 2026-10-01
Updated: 2026-10-01
Code: https://github.com/tatsu-lab/stanford_alpaca
License: http://creativecommons.org/licenses/by-nc-nd/4.0/
The gist: Multi-token prediction (MTP) improves the throughput of autoregressive generation by enabling the language model to draft multiple next tokens per forward pass, while a verification step over draft
Abstract
Multi-token prediction (MTP) improves the throughput of autoregressive generation by enabling the language model to draft multiple next tokens per forward pass, while a verification step over draft tokens ensures that token distribution of the backbone is preserved. Every open MTP-family release (MiMo-7B, DeepSeek-V3, Qwen3) trains its heads jointly with the backbone over the full pretraining run of tens of trillions of tokens, thus setting the drafter quality at pretraining time. We ask whether a lightweight post-training pass on target-generated chain-of-thought is enough to reach the same expected throughput speedup on a frozen reasoning model, and study how a serving-time system built on such a checkpoint can be optimized. We present three findings. 1) On a frozen Qwen3-8B with K = 3 chained MTP heads, we show that a post-training recipe with plain cross-entropy on about!2.5 B tokens reaches or exceeds the expected speedup of jointly trained MiMo-7B on math, coding and knowledge benchmarks. Our post-training recipe utilizes 10 cubed - 10 4 times less MTP-training tokens as compared with joint pre-training of MiMO-7B MTP baseline. 2) We propose a chain-aware relaxation of draft token verification rule that allows a bounded drift from backbone language model token distribution. We show that this relaxation lifts expected speedups by +12 to +16% per benchmark while preserving task accuracy. 3) We propose an adaptive controller that dynamically chooses the number of MTP heads to be engaged at inference time and demonstrate recovery of upto 11 -- 14% loss in speedup using fixed maximum MTP draft length.
Sources
- Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding
- Program Synthesis with Large Language Models
- Accelerating Large Language Model Decoding with Speculative Sampling
- Evaluating Large Language Models Trained on Code
- Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation
- DeepSeek-V3 Technical Report
- Are We Done with MMLU?
- Better & Faster Large Language Models via Multi-token Prediction
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- LoRA: Low-Rank Adaptation of Large Language Models
- Nightjar: Dynamic Adaptive Speculative Decoding for Large Language Models Serving
- EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test
- Let's Verify Step by Step
- MiMo: Unlocking the Reasoning Potential of Language Model -- From Pretraining to Posttraining
- MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark
- Qwen3 Technical Report
- Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks