Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration

Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin


Abstract
Reinforcement Learning from Verifiable Rewards (RLVR) has emerged as the standard paradigm for improving reasoning capability of large language models, while Multi-Token Prediction (MTP) has been a widely adopted module in pretraining. Combining the two is a natural approach, yet current RL practices detach MTP gradients because joint training degrades the performance. We revisit this failure from an optimization perspective. We show that the per-step effect of MTP on the RL objective can be decomposed into two terms: a first-order correlation and a second-order perturbation penalty. This decomposition unifies three MTP training regimes: Detach, Cross-Entropy loss, and Policy loss, and explains why each succeeds or fails. Further analysis of policy loss reveals that, although it aligns with intuition, performance still degrades: the correlation term decays while the quadratic penalty persists. Guided by the analysis, we propose Optimal Coefficient Calibration (OCC), an adaptive scheme that tracks the optimal coefficient online via a log-probability proxy at negligible cost. Across six competition-level mathematical reasoning benchmarks, OCC consistently matches or exceeds the detach baseline, delivering improved joint MTP-RL training performance.
Anthology ID:
2026.findings-emnlp.671
Volume:
Findings of the Association for Computational Linguistics: EMNLP 2026
Month:
October
Year:
2026
Address:
Budapest, Hungary
Editors:
Hung-Yi Lee, Sunipa Dev, Jessy Li, Desmond Elliott
Venue:
Findings
SIG:
SIGDAT
Publisher:
Association for Computational Linguistics
Note:
Pages:
14094–14108
Language:
URL:
https://aclanthology.org/2026.findings-emnlp.671/
DOI:
Bibkey:
Cite (ACL):
Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, and Guojun Yin. 2026. Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration. In Findings of the Association for Computational Linguistics: EMNLP 2026, pages 14094–14108, Budapest, Hungary. Association for Computational Linguistics.
Cite (Informal):
Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration (Wang et al., Findings 2026)
Copy Citation:
PDF:
https://aclanthology.org/2026.findings-emnlp.671.pdf