Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration
On the map: LLM reasoning → Reinforcement learning for reasoning
Abstract (arXiv)
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 them 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.
Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning. Builds on 18 of its 27 references that are in the map. First posted on arXiv 27 May 2026.
Builds on
18 of 27 references.
Closest ancestors · 2
Same direction, newest first
- 2026FIPO: Eliciting Deep Reasoning with Future-KL Influenced Policy Optimization20 Mar
- 2017Proximal Policy Optimization Algorithms20 Jul
Other references · 16
- 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,371
- 2015Trust Region Policy Optimizationcited by 2,580
- 2024HybridFlow: A Flexible and Efficient RLHF Frameworkcited by 2,057
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
- 2022Fast Inference from Transformers via Speculative Decodingcited by 1,653
- 2022Solving Quantitative Reasoning Problems with Language Modelscited by 1,392
- 2023SGLang: Efficient Execution of Structured Language Model Programscited by 1,141
- 2024OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problemscited by 1,136
- 2017Visualizing the Loss Landscape of Neural Netscited by 1,130
- 2026DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligencecited by 1,040
- 2023Accelerating Large Language Model Decoding with Speculative Samplingcited by 1,006
- 2026GLM-5: from Vibe Coding to Agentic Engineeringcited by 587
4 more
- 2024Better & Faster Large Language Models via Multi-token Predictioncited by 360
- 2018Blockwise Parallel Decoding for Deep Autoregressive Modelscited by 323
- 2019Global Convergence of Policy Gradient Methods to (Almost) Locally Optimal Policiescited by 124
- 2026Beyond Precision: Training-Inference Mismatch is an Optimization Problem and Simple LR Scheduling Fixes Itcited by 9
Built on by
Citing papers per year: 2026: 2.
2 papers.
Similar, not cited
Close in wording, no citation either way
- 2025Diversity or Precision? A Deep Dive into Next Token Prediction
- 2026Breaking Entropy Bounds: Accelerating RL Training via MTP with Rejection Sampling
- 2026Which Tokens Matter? Adaptive Token Selection for RLVR with the Relative Surprisal Index
- 2025Stabilizing Knowledge, Promoting Reasoning: Dual-Token Constraints for RLVR
- 2026Decoupling Reasoning and Confidence: Resurrecting Calibration in Reinforcement Learning from Verifiable Rewards
- 2026DelTA: Discriminative Token Credit Assignment for Reinforcement Learning from Verifiable Rewards
- 2025Accelerating RL for LLM Reasoning with Optimal Advantage Regression
- 2026Don't Mix Rewards, Mix Policies: Policy Decomposition and Optimization for Multi-Reward RL