AMTP:自适应多 token 预测让 Llama-3.1 解码提速至 4.1 倍
Adaptive Multi-Token Prediction: How far should we go?
arXiv 上这篇论文给 MTP 解码加了个轻量控制器,不用重训练就能自动选草稿长度,Llama-3.1 上最高跑出 4.1 倍解码速度,做推理优化的可以看看。
arXiv 论文提出 Adaptive Multi-Token Prediction(AMTP),把多 token 预测的草稿长度选择转化为在线控制问题,用验证反馈动态决定每轮预测多远。该控制器无需重训练,也不改动验证器本身。论文基于 Llama-3.1-8B-Instruct 搭配 Medusa heads、Meta 的 MTP 模型和 Qwen3-4B-Inst-2507-MTP 构建了五组草稿-目标配对,在 AIME 2025 和 GSM8K 上评测。结果显示 AMTP 相比标准解码最高提速 4.1 倍,10 个设置中的 8 个达到或超过固定窗口方案,吞吐最高提升 10%。
Adaptive Multi-Token Prediction: How far should we go?
Multi-token prediction (MTP) can accelerate language-model decoding by drafting several tokens and verifying them in a single target-model pass. However, the best speculation horizon changes with context: longer drafts offer greater amortization when reliable, but can waste compute on suffixes discarded after a rejection. Unlike autoregressive drafting, where a stopping rule can decide token by token, MTP must fix the horizon before drafting, as all heads fire in the same forward pass. We therefore formulate horizon selection as an online control problem, using verification feedback to learn how far the draft can be trusted in each context, and introduce Adaptive Multi-Token Prediction (AMTP): a lightweight controller that requires no retraining and leaves the exact verifier unchanged. To guide horizon selection, we characterize the expected number of tokens committed at each horizon via acceptance survival, the probability that a draft prefix is fully accepted. Before each round, AMTP predicts how far the draft is likely to survive and picks the horizon with the highest expected tokens per second, learning freely from every verification pass. Theoretically, we show that any online learner's prediction error translates into throughput regret against the best horizon for each context, and in turn into a bound on wall-clock completion time. With online gradient descent, this regret is sublinear. Across five draft-target pairings built on Llama-3.1-8B-Instruct with Medusa heads, Meta's MTP models, and Qwen3-4B-Inst-2507-MTP, evaluated on AIME 2025 and GSM8K, AMTP decodes up to 4.1x faster than standard decoding and matches or outperforms the fixed window in 8 of 10 settings, with throughput gains of up to 10%.