Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads

2026-09-30

AI Summary

This study addresses the prohibitive computational cost of multi-token prediction, which typically relies on expensive joint pre-training. To overcome this limitation, this work proposes an efficient approach that enables multi-token prediction through lightweight post-training on a frozen inference model. Specifically, the method fine-tunes using only cross-entropy loss and introduces a chain-aware verification relaxation rule alongside an adaptive controller to dynamically adjust draft length, thereby balancing accuracy and speed. Experimental results demonstrate that the proposed approach achieves or surpasses baseline acceleration ratios while requiring merely one-thousandth of the training data. Furthermore, it improves serving throughput by 12–16% without compromising task accuracy, offering an effective solution for accelerating large language model inference in resource-constrained scenarios.

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 $\approx\!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^3$-$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.

Problem

• Multi-Token Prediction
• Post-Training
• Autoregressive Generation
• Throughput Optimization
• Draft Token Verification

Innovation

• Multi-Token Prediction
• Post-Training
• Draft Token Verification Relaxation
• Adaptive Controller
• Chain-of-Thought