MTP-RL: Acceleration of Reinforcement Learning Rollouts with Policy-Aligned Multi-Token Prediction
Paper · 2026 · Ke Wang, Aohan Zeng, Zhengxiao Du et al.
Summary
Rollout generation consumes most of the compute in reinforcement-learning fine-tuning of LLMs, especially when models produce long reasoning chains. Speculative decoding with multi-token prediction (MTP) could speed this up, but it faces two obstacles. Most open models have no MTP head. Where one exists and is kept frozen, the evolving policy drifts away from it, and the draft acceptance length falls sharply. The authors call this "acceptance collapse."
Pilot experiments show that freezing the MTP, training it with plain cross-entropy, or cutting its gradients into the main model all fail to prevent this collapse. MTP-RL therefore works in two stages:
- Equipping stage: a lightweight, multi-layer, parameter-sharing MTP module is trained from scratch so that any model can draft tokens.
- Joint RL stage: the policy and the MTP are optimized together. The MTP objective combines cross-entropy with a KL term weighted by advantages, which favours high-reward trajectories and down-weights poor ones.
With this setup, acceptance length grows during RL instead of decaying, and rollouts become considerably faster at little cost to accuracy. A further benefit is that the trained model keeps speculative decoding ability for downstream use.
Key points
- Experiments use Qwen models of several sizes, including Qwen2.5-7B-Instruct and Qwen2.5-14B-Instruct, on math and code tasks. The RL algorithms tested are GRPO and GSPO, with rollouts served through SGLang.
- Compared with baselines, average rollout time drops by 23.1%–55.3%.
- On Qwen2.5-14B-Instruct with GSPO, the speedup peaks at 3.15× and averages 2.12×.
- End-to-end cost was measured including the extra MTP-equipping stage. Total training time falls by about 16.4% for 7B and 35.6% for 14B under GRPO. Math accuracy shifts slightly: 76.8 → 76.2 for 7B and 79.8 → 79.2 for 14B.
- Stated limitations:
- The MTP-equipping stage adds engineering complexity and overhead.
- The gains were measured at batch size 256 and remain well below what speculative decoding achieves with small batches.
Institutions
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works
Cite
Cite as Ke Wang, Aohan Zeng, Zhengxiao Du, Yuxuan Hu, Bohan Zhang, Xinyi Wang, Jie Tang and Jing Zhang. 2026. MTP-RL: Acceleration of Reinforcement Learning Rollouts with Policy-Aligned Multi-Token Prediction. Findings of the Association for Computational Linguistics: ACL 2026.
@inproceedings{wang2026mtp,
title = {{MTP-RL: Acceleration of Reinforcement Learning Rollouts with Policy-Aligned Multi-Token Prediction}},
author = {Ke Wang, Aohan Zeng, Zhengxiao Du, Yuxuan Hu, Bohan Zhang, Xinyi Wang, Jie Tang and Jing Zhang},
booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
year = {2026},
doi = {10.18653/v1/2026.findings-acl.1871},
url = {https://doi.org/10.18653/v1/2026.findings-acl.1871}
}References
- Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. 2024. Medusa: Simple llm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401.10774.
- Yuxuan Cai, Xiaozhuan Liang, Xinghua Wang, Jin Ma, Haijin Liang, Jinwen Luo, Xinyu Zuo, Lisheng Duan, Yuyang Yin, and Xi Chen. 2025. Fastmtp: Accelerating llm inference with enhanced multi-token prediction. arXiv preprint arXiv:2509.18362. — [arXiv:2509.18362](https://arxiv.org/abs/2509.18362) find
- Mark Chen. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374.
- Qiaoling Chen, Zijun Liu, Peng Sun, Shenggui Li, Guoteng Wang, Ziming Liu, Yonggang Wen, Siyuan Feng, and Tianwei Zhang. 2025. Respec: Towards optimizing speculative decoding in reinforcement learning systems. arXiv preprint arXiv:2510.26475. — [arXiv:2510.26475](https://arxiv.org/abs/2510.26475) find
- Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. 2023. Enhancing chat language models by scaling high-quality instructional conversations. Preprint, arXiv:2305.14233.
- Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, and 1 others. 2025. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948.
- Jingkai He, Tianjian Li, Erhu Feng, Dong Du, Qian Liu, Tao Liu, Yubin Xia, and Haibo Chen. 2025. History rhymes: Accelerating llm reinforcement learning with rhymerl. arXiv preprint arXiv:2508.18588. find
- Zhenyu He, Zexuan Zhong, Tianle Cai, Jason Lee, and Di He. 2024. Rest: Retrieval-based speculative decoding. In Proceedings of the 2024 conference of the North American chapter of the association for computational linguistics: Human language technologies (volume 1: long papers), pages 1582–1595. find
- Yuxuan Hu, Ke Wang, Xiaokang Zhang, Fanjin Zhang, Cuiping Li, Hong Chen, and Jing Zhang. 2025. Sam decoding: Speculative decoding via suffix automaton. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 12187–12204. find
- Wei Huang, Yi Ge, Shuai Yang, Yicheng Xiao, Huizi Mao, Yujun Lin, Hanrong Ye, Sifei Liu, Ka Chun Cheung, Hongxu Yin, and 1 others. 2025. Qerl: Beyond efficiency–quantization-enhanced reinforcement learning for llms. arXiv preprint arXiv:2510.11696. find
- Ruoyu Qin, Weiran He, Weixiao Huang, Yangkun Zhang, Yikai Zhao, Bo Pang, Xinran Xu, Yingdi Shan, Yongwei Wu, and Mingxing Zhang. 2025. Seer: Online context learning for fast synchronous llm reinforcement learning. arXiv preprint arXiv:2511.14617. — [arXiv:2511.14617](https://arxiv.org/abs/2511.14617) find
- Yaniv Leviathan, Matan Kalman, and Yossi Matias. 2023. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pages 19274–19286. PMLR. find
- Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. 2023. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36:53728–53741.
- Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. 2024a. Eagle-2: Faster inference of language models with dynamic draft trees. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 7421–7432.
- John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. find
- Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. 2024b. Eagle: Speculative sampling requires rethinking feature uncertainty. arXiv preprint arXiv:2401.15077.
- Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, and 1 others. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.
- Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. 2025. Eagle-3: Scaling up inference acceleration of large language models via training-time test. arXiv preprint arXiv:2503.01840. — [arXiv:2503.01840](https://arxiv.org/abs/2503.01840) find
- Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023. Let’s verify step by step. In The Twelfth International Conference on Learning Representations.
- Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, and 1 others. 2024. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437. find
- Liyuan Liu, Feng Yao, Dinghuai Zhang, Chengyu Dong, Jingbo Shang, and Jianfeng Gao. 2025. Flashrl: 8bit rollouts, full power rl. find
- Xingzhou Lou, Junge Zhang, Jian Xie, Lifeng Liu, Dong Yan, and Kaiqi Huang. 2024. Spo: Multi-dimensional preference sequential alignment with implicit reward modeling. arXiv preprint arXiv:2405.12739. — [arXiv:2405.12739](https://arxiv.org/abs/2405.12739) find
- Xianzhen Luo, Yixuan Wang, Qingfu Zhu, Zhiming Zhang, Xuanyu Zhang, Qing Yang, and Dongliang Xu. 2025. Turning trash into treasure: Accelerating inference of large language models with token recycling. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume…
- Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori B Hashimoto. 2025. s1: Simple test-time scaling. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages…
- Wentao Shi, Mengqi Yuan, Junkang Wu, Qifan Wang, and Fuli Feng. 2024. Direct multi-turn preference optimization for language agents. arXiv preprint arXiv:2406.14868. find
- Heming Xia, Tao Ge, Peiyi Wang, Si-Qing Chen, Furu Wei, and Zhifang Sui. 2023. Speculative decoding: Exploiting speculative execution for accelerating seq2seq generation. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 3909–3925.
- LLM Xiaomi, Bingquan Xia, Bowen Shen, Dawei Zhu, Di Zhang, Gang Wang, Hailin Zhang, Huaqiu Liu, Jiebao Xiao, Jinhao Dong, and 1 others. 2025. Mimo: Unlocking the reasoning potential of language model– from pretraining to posttraining. arXiv preprint arXiv:2505.07608. find
- Zhangchen Xu, Yang Liu, Yueqin Yin, Mingyuan Zhou, and Radha Poovendran. 2025. Kodcode: A diverse, challenging, and verifiable synthetic dataset for coding. arXiv preprint arXiv:2503.02951. find
- Shu Yang, Junchao Wu, Xin Chen, Yunze Xiao, Xinyi Yang, Derek F Wong, and Di Wang. 2025. Understanding aha moments: from external observations to internal mechanisms. arXiv preprint arXiv:2504.02956. — [arXiv:2504.02956](https://arxiv.org/abs/2504.02956) find
- Aohan Zeng, Xin Lv, Qinkai Zheng, Zhenyu Hou, Bin Chen, Chengxing Xie, Cunxiang Wang, Da Yin, Hao Zeng, Jiajie Zhang, and 1 others. 2025. Glm-4.5: Agentic, reasoning, and coding (arc) foundation models. arXiv preprint arXiv:2508.06471.
- Hanchen Zhang, Xiao Liu, Bowen Lv, Xueqiao Sun, Bohao Jing, Iat Long Iong, Zhenyu Hou, Zehan Qi, Hanyu Lai, Yifan Xu, and 1 others. 2025. Agentrl: Scaling agentic reinforcement learning with a multi-turn, multi-task framework. arXiv preprint arXiv:2510.04206. — [arXiv:2510.04206](https://arxiv.org/abs/2510.04206) find
- Yao Zhao, Zhitian Xie, Chen Liang, Chenyi Zhuang, and Jinjie Gu. 2024a. Lookahead: An inference acceleration framework for large language model with lossless generation accuracy. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 6344–6355. find
- Yuze Zhao, Jintao Huang, Jinghan Hu, Xingjun Wang, Yunlin Mao, Daoze Zhang, Zeyinzi Jiang, Zhikai Wu, Baole Ai, Ang Wang, Wenmeng Zhou, and Yingda Chen. 2024b. Swift:a scalable lightweight infrastructure for fine-tuning. Preprint, arXiv:2408.05517. — [arXiv:2408.05517](https://arxiv.org/abs/2408.05517) find
- A Train-from-scratch configurations. The MTP module is initialized and trained from scratch. Key hyperparameters, including the optimizer, learning rate, batch size, training epochs and the loss-weighting hyperparameters for each MTP layer are summarized in Table 5. Chujie Zheng, Shixuan Liu, Mingze… — Group Sequence Policy Optimization
- Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Livia Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, and 1 others. 2024. Sglang: Efficient execution of structured language model programs. Advances in neural information processing systems, 37:62557– 62…
- Zilin Zhu, Chengxing Xie, Xin Lv, and slime Contributors. 2025. slime: An llm post-training framework for rl scaling. https://github.com/THUDM/slime. GitHub repository. Corresponding author: Xin Lv. ↗
- Configurations for Training Hyperparameter Value Epoch batch size Temperature Learning rate Optimizer(AdamW) λ 1 , λ 2 , λ 3 5 16 0.8 5e-5 (β 1 , β 2 ) = (0.9, 0.95) 0.6, 0.3, 0.1 *Table 5: Train-from-Scratch hyperparameters for the MTP module* RL configurations. During RL, the MTP module is trained… find
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/wang-2026-mtp-rl-acceleration-reinforcement-learning/