Source-linked AI summary
Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads
Prachi Badarayani, Aidan Jay, Chenghui Zhou, Dayquan Julienne, Yuan Gao, Tianwei Chen, George Zerveas, Ishmam Zabir, Xiren Zhou, Chris Quirk, Xia Song
TL;DR
The paper asks whether post-training can match the MTP speedup of jointly pretrained heads on a frozen reasoning model, and how inference can be optimized. It trains chained heads on target-generated continuations, relaxes verification using chain-level probabilities, and adapts draft depth to throughput feedback. The reported results include matching or exceeding the joint-pretraining baseline with much less training compute, alongside throughput gains from relaxed verification and adaptive depth.
Problem
Jointly pretrained MTP heads incur pretraining-scale training cost, and it is unclear whether that cost is intrinsic to MTP quality or specific to the recipe.
Source: §1 · p. 1Method
Results
Post-trained heads match or exceed MiMo-7B on mean accepted length at 10^3–10^4× less MTP-training compute, while the verification relaxation improves MAL without compromising task accuracy.
Source: §1 · pp. 1–2, §5 · p. 3Takeaways & Limitations
Post-training permits exploration of MTP architectures and serving strategies, including Mamba-style chaining and adaptive head selection.
Source: §1 · pp. 1–2
Abstract
from arXiv · showhide
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.
1 Introduction
The introduction asks whether MTP quality requires joint pretraining and presents post-trained heads, relaxed verification, and adaptive draft depth as alternatives. It reports matching or better speedup with far less training compute and higher expected throughput without compromising task accuracy.
1 Introduction
- We show that, using backbone-generated chain-of-thought as training data, post-trained MTP heads can match or outperform jointly pretrained MTP heads at 10^3–10^4× less training compute.This enables exploration of post-trained MTP architectures, including Mamba chaining of MTP heads, which improves throughput.Source: §1 · p. 1
- The paper asks whether the cost of jointly pretraining MTP heads is intrinsic to MTP quality or specific to the training recipe.Existing releases train their MTP heads jointly with the backbone at pretraining scale.Source: §1 · p. 1
- A chain-aware relaxation of verification improves expected throughput speedup while preserving task accuracy.It relaxes the exact distributional-match requirement by thresholding the running joint probability of draft tokens.Source: §1 · pp. 1–2
- An adaptive controller selects how many MTP heads to engage using online throughput feedback, improving mean accepted length over a fixed draft length.The motivation is that fixed draft sizes can reduce throughput at high batch sizes and limit use of longer drafts at small batch sizes.Source: §1 · p. 2
2 Post-training on target-generated chain-of-thought
The method trains chained MTP heads on continuations generated by the target model while keeping the backbone frozen. The heads predict successive draft tokens and are trained with per-head cross-entropy on an on-policy corpus.
2 Post-training on target-generated chain-of-thought
- Each MTP head consumes backbone hidden states and the preceding draft token to predict the next token in a three-head chain.The heads produce drafts d_t+1, d_t+2, and d_t+3 respectively.Source: §2 · p. 2
- Training minimizes per-head cross-entropy, with each head sharing the frozen backbone’s language-model unembedding head and adding about 200M trainable parameters.The per-head parameter budget matches MiMo-7B’s.Source: §2 · p. 2
- Only MTP-head parameters are updated; the backbone remains frozen, and each head conditions on the preceding head’s argmax token during training.This matches the inference-time chain; the main empirical results use K=3 heads, with a K=6 variant for the adaptive controller.Source: §2 · p. 2
- The on-policy corpus uses 1.89M prompts and Qwen3-8B-generated chain-of-thought responses, packed into sequences yielding about 2.5B training tokens per epoch.Prompts span general chat, code, math, and instruction following; responses are sampled in thinking mode, up to 8K tokens each.Source: §2 · p. 2
3 Chain-aware approximate verification
The section replaces strict distribution matching with a chain-aware acceptance rule. Rather than assessing each draft position independently, it evaluates the running geometric mean of target probabilities across a draft prefix.
3 Chain-aware approximate verification
- The proposed rule accepts the maximal draft prefix whose running geometric-mean target probability stays above threshold τ.This relaxes the strict distributional-match guarantee using the joint signal across the chain of draft positions.Source: §3 · p. 2
- Classical speculative decoding uses rejection sampling to preserve the target output distribution exactly, a guarantee the authors call stronger than necessary for peaked reasoning targets.The classical rule accepts a draft token based on a uniform draw and the ratio of its draft probability to its target probability.Source: §3 · p. 2
- Per-position relaxation criteria discard joint information in the draft chain, treating one low-probability token in a confident chain differently from consecutive mediocre positions.The proposed threshold instead evaluates the running probability across a prefix.Source: §3 · p. 2
4 Adaptive draft depth
The section motivates adaptive draft depth by the changing throughput tradeoff across batch sizes and text difficulty. It proposes an online bandit controller that chooses among draft depths using observed throughput.
4 Adaptive draft depth
- Runtime throughput can peak at three heads despite monotonically increasing expected speedup, while low-batch predictable text can benefit from more heads and high-batch or difficult text can favor fewer.Deeper verification may become a bottleneck when its acceptance gains do not cover the additional verification costs.Source: §4 · p. 3
- An online lower-confidence-bound bandit controller selects among draft depths K=1 to K=6 using observed batched throughput.It explores all six arms, then exploits the running mean with periodic re-exploration and commits when one arm leads by a relative margin for C windows.Source: §4 · p. 3
- If the controller commits to an arm whose mean is within δ of the best, the resulting loss is bounded by δ.If no arm commits within Wmax windows, it falls back to the arm with the highest lower-confidence bound.Source: §4 · p. 3
5 Experiments
On frozen Qwen3-8B, post-trained MTP heads generally match or exceed MiMo-7B in mean accepted length (MAL). The experiments also evaluate chain-aware verification and adaptive draft depth as ways to improve serving throughput.
5 Experiments
- The main comparison uses frozen Qwen3-8B with K=3 chained heads against MiMo-7B-RL, with both systems evaluated at batch 4 on one A100-80GB GPU in vLLM.Source: §5 · p. 3
- Post-trained heads generally match or exceed MiMo-7B on MAL across benchmarks.Table 1 reports MAL and per-head acceptance rates at matched per-head capacity; exact verification preserves task accuracy.Source: Table 1 · p. 3, §5 · p. 3
- Averaged over four benchmarks, chain-aware verification reaches MAL 3.31, a +14% lift over exact verification with accuracy preserved on every benchmark.The comparison selects the best MAL for each rule subject to task accuracy remaining within ±3 percentage points of exact verification.Source: Table 2 · p. 4, §5 · p. 4
- The adaptive controller captures deep-head gains where K=6 helps and falls back to smaller K when deeper verification costs outweigh acceptance gains.Fixed K=6 forfeits 11–14% throughput at batches 64/128; controller overhead loses up to 6% of speedup ratio on average at low batch size, while recovering most lost throughput at high batch sizes.Source: §5 · p. 4, Table 3 · p. 4
A Ablations
The ablations distinguish the gain from target-generated self-distillation from extensions trained on raw AM-Thinking text. Mamba chaining is the strongest raw-corpus extension, while several other interventions do not improve MAL.
A Ablations
- Self-distillation instead of raw AM-Thinking text gains +0.14 to +0.41 MAL when corpus size and the rest of the recipe are held constant.The lower-block extensions use the raw-corpus baseline, so their results should be compared with Row 2 rather than the self-distilled base.Source: A Ablations · p. 7, Table 4 · p. 7
- Against the raw-corpus baseline, a Mamba gate gains +0.28 MAL on average and up to +0.43 on GSM8K, but trails self-distillation on three of four math benchmarks.Mamba gate chaining produces the most improvement among the ablation extensions over Row 2.Source: A Ablations · p. 7, Table 4 · p. 7
- GRU gating is a small loss, backbone LoRA is neutral to negative, and full-vocabulary KL training does not improve MAL.Source: A Ablations · p. 7
B Coverage outside math datasets
Post-trained heads exceed MiMo-7B on MAL across the listed coding, knowledge, chat, instruction-following, and summarisation benchmarks.
B Coverage outside math datasets
- Post-trained heads exceed MiMo-7B on MAL for MMLU-Redux, MBPP, MT-Bench, Alpaca, and CNN/DailyMail.The H3 gap is +15pp on MT-Bench and +12pp on Alpaca; task-accuracy scoring on MT-Bench and Alpaca uses no LLM-as-judge result here.Source: B Coverage outside math datasets · p. 7
C Full τ sweep for chain-aware verification
The chain-aware τ sweep identifies an operating range that preserves math accuracy while improving expected speedup, but the threshold’s effect differs between exact-answer math grading and ROUGE.
C Full τ sweep for chain-aware verification
- At τ=10^-6, H1 acceptance is 95% but the geometric-mean penalty rejects 49–55% of H2 positions; below τ=10^-12, the clamp makes MAL reach the K+1=4 ceiling.MAL is non-monotone: τ=10^-3 can outperform τ=10^-6 because the chain penalty compounds less harshly.Source: C Full τ sweep for chain-aware… · p. 7
- +0.30 to +0.45 MAL over exact verification is reached at τ=10^-3, while math accuracy matches the exact-verification baseline for τ ∈ [10^-3, 0.90].Below the corridor, AIME24/25 accuracy collapses to 0% and MATH500 drops to 32% at the accept-everything ceiling.Source: C Full τ sweep for chain-aware… · p. 8, Table 5 · p. 8
- On CNN/DailyMail, ROUGE-1/L stays at 13.75–13.87 across the safe corridor and falls to 9.57/6.88 at the accept-everything ceiling.The sweep reports ROUGE-1/2/L against reference highlights.Source: C Full τ sweep for chain-aware… · p. 8, Table 6 · p. 8
D Adaptive draft depth: multi-benchmark coverage
The adaptive controller varies draft depth across workloads, preserving much of fixed K=3 throughput while avoiding losses from fixed K=6 at high batch sizes.
D Adaptive draft depth: multi-benchmark coverage
- The bandit reaches 93–98% of fixed K=3 across workloads and at least 97% at high batch by falling back when deeper heads lose their acceptance-versus-verification trade-off.Fixed K=6 matches K=3 within noise at low batch (+1%) but loses 11–14% at high batch; it loses up to 29% in the strict Pareto comparison.Source: D Adaptive draft depth:… · p. 8
E Baseline choice: MiMo-7B-RL and Qwen3-8B
The authors select a jointly pretrained baseline and a target backbone with similar parameter counts to make the drafter comparison interpretable. The supplied workload results also show why draft-depth comparisons depend on batch size.
E Baseline choice: MiMo-7B-RL and Qwen3-8B
- MiMo-7B-RL is the only publicly available jointly pretrained MTP baseline matching the reporting protocol; its three chained heads and ≈7B-parameter backbone are closest to the comparison target.DeepSeek-V3 has one head, EAGLE-3 uses a recurrent draft module, and larger open MTP checkpoints use backbones an order of magnitude larger.Source: E Baseline choice: MiMo-7B-RL and… · pp. 8–9
- Table 7 compares bandit and fixed-K throughput averaged across ten workloads using paired per-sample-median tok/s; fixed K=6 crashes at high batch.The bandit reaches at least 97% of fixed K=3 there by falling back.Source: Table 7 · p. 9
- At high batch, fixed K=6 loses up to 20–29% on math/MCQ, while the bandit stays close to whichever fixed depth is best.The per-cell view reports K=6 lift over K=3 and bandit percentage of the best fixed depth.Source: Table 8 · p. 9
- Qwen3-8B is chosen as the target because it is within ≈15% of MiMo-7B in parameter count and reports competitive quality on the four math benchmarks.The authors state that using a 70B-class target would change downstream numbers and confound the training-recipe comparison.Source: E Baseline choice: MiMo-7B-RL and… · p. 9
F Reproducibility details
The reproducibility details specify the head architecture, training data and schedule, implementation variants, controller, and evaluation setup. The submission says code and weights are not released, while providing a recipe for independent re-implementation.
F Reproducibility details
- Each MTP head uses one transformer decoder layer at the 4096-wide backbone width, followed by the frozen backbone LM head, and has approximately 200M parameters.Source: F Reproducibility details · p. 10
- Training uses five epochs on ≈2.5B target-generated tokens per epoch, with the backbone frozen and only MTP head parameters updated.The K=6 variant uses the same schedule with two additional heads.Source: F Reproducibility details · p. 10
- The on-policy corpus contains 1.89M AM-Thinking prompts completed by frozen Qwen3-8B in thinking mode, with prompt-response pairs concatenated without filtering.Decoding uses temperature 0.6, top-p 0.95, and top-k 20.Source: F Reproducibility details · p. 10
- The GRU cross-head gate consumes a head’s hidden state and previous argmax-token embedding, then emits a gated state for the next head in place of the base hidden state.The gate has ≈50M parameters per handoff and is trained end-to-end while the backbone remains frozen.Source: F Reproducibility details · p. 10
- The Mamba-style gate replaces the GRU cell with a selective-state-space block over a length-K virtual sequence; its gate parameters are ≈80M per handoff.For K=3 or K=6, the scan is implemented as a small unrolled recursion without the kernel-fused variant.Source: F Reproducibility details · pp. 10–11
- The LoRA variants adapt Q/K/V/O attention projections in either the last 4 or last 18 of 36 backbone layers, using rank 8 and α=16.The adapter totals are ≈1M parameters for shallow LoRA and ≈4.7M for deep LoRA.Source: F Reproducibility details · p. 11
- The full-vocabulary KL ablation replaces per-token hard cross-entropy with DKL(pθ∥p(k)) over the vocabulary at every position.The ablation architectures use the raw AM-Thinking corpus and the base recipe’s schedule and optimizer.Source: F Reproducibility details · pp. 10–11
- The Adaptive-K controller measures batched throughput in eight-step near-full-occupancy windows, explores six draft depths, and commits to a lower-confidence-bound arm before freezing depth for the session.It uses at least two round-robin windows per arm, periodic re-exploration, and a 30-window hard cap; the verification chain threshold is clamped at p ≥ 10^-12.Source: F Reproducibility details · p. 11