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

• URL: http://arxiv.org/abs/2605.28184v1
• Date: Wed, 27 May 2026 09:07:06 GMT
• Title: Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration
• Abstract summary: Multi-Token Prediction (MTP) has been a widely adopted module in pretraining.<n>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 penalty.<n>We propose an adaptive scheme that tracks the optimal coefficient online via a log-probability proxy at negligible cost.
• Score: 61.46060073417047
• License: http://creativecommons.org/licenses/by/4.0/
• 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 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.

Related papers

• Breaking Entropy Bounds: Accelerating RL Training via MTP with Rejection Sampling [87.16803442525755]
  Reinforcement learning (RL) has become a key component in modern large language models, yet the rollout stage remains the key bottleneck in RL training pipelines.<n>Although Multi-Token Prediction (MTP) offers a natural solution to accelerate rollouts through speculative decoding, many studies have observed that MTP acceptance rates degrade significantly during RL training.<n>We present Bebop, a systematic study of MTP in LLM post-training, and offer practical recipes to integrate MTP into large-scale RL pipelines.
  arXiv  Detail & Related papers  (2026-06-10T17:36:45Z)
• Momentum for Reasoning: Dense Intrinsic Signals in Policy Optimization [31.547614109943158]
  Reinforcement learning with verifiable rewards (RLVR) has emerged as a powerful paradigm for eliciting long-chain reasoning in large language models.<n>We propose ISPO (Intrinsic Signal Policy Optimization), which combines a sequence-level signal measuring how informative the thinking trajectory is for the final answer, with a token-level directional reward whose hallucinated-certainty hinge penalizes confidently-wrong predictions at critical decision tokens.<n>ISPO consistently outperforms competitive baselines, with the largest gains on the hardest benchmarks where zero-advantage collapse is most frequent, and training-dynamics diagnostics confirm that both failure modes are decreased.
  arXiv  Detail & Related papers  (2026-06-07T20:08:36Z)
• DISA: Offline Importance Sampling for Distribution-Matching LLM-RL [56.9445657766829]
  We introduce DISA, which moves this calibration problem outside the RL loop.<n> DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate.<n>On two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline Flow.
  arXiv  Detail & Related papers  (2026-05-17T07:14:44Z)
• Alternating Reinforcement Learning for Rubric-Based Reward Modeling in Non-Verifiable LLM Post-Training [29.56905427210088]
  gradient-ARM is a framework that jointly optimize a rubric generator and a judge using reinforcement learning from preference feedback.<n>We show that gradient-ARM achieves state-of-the-art performance among baselines on benchmarks and significantly improves downstream policy alignment in both offline and online reinforcement learning settings.
  arXiv  Detail & Related papers  (2026-02-02T00:50:53Z)
• CurES: From Gradient Analysis to Efficient Curriculum Learning for Reasoning LLMs [53.749193998004166]
  Curriculum learning plays a crucial role in enhancing the training efficiency of large language models.<n>We propose CurES, an efficient training method that accelerates convergence and employs Bayesian posterior estimation to minimize computational overhead.
  arXiv  Detail & Related papers  (2025-10-01T15:41:27Z)
• WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training [64.0932926819307]
  We present Warmup-Stable and Merge (WSM), a framework that establishes a formal connection between learning rate decay and model merging.<n>WSM provides a unified theoretical foundation for emulating various decay strategies.<n>Our framework consistently outperforms the widely-adopted Warmup-Stable-Decay (WSD) approach across multiple benchmarks.
  arXiv  Detail & Related papers  (2025-07-23T16:02:06Z)
• Unearthing Gems from Stones: Policy Optimization with Negative Sample Augmentation for LLM Reasoning [41.83677588934301]
  We propose Behavior Constrained Policy Gradient with Negative Sample Augmentation (BCPG-NSA)<n>BCPG-NSA is a fine-grained offline framework that encompasses three stages: 1) sample segmentation, 2) consensus-based step correctness assessment combining LLM and PRM judgers, and 3) policy optimization with NSA designed to effectively mine positive steps within negative samples.<n> Experimental results show that BCPG-NSA outperforms baselines on several challenging math/coding reasoning benchmarks using the same training dataset.
  arXiv  Detail & Related papers  (2025-05-20T14:16:49Z)
• Prior Constraints-based Reward Model Training for Aligning Large Language Models [58.33118716810208]
  This paper proposes a Prior Constraints-based Reward Model (namely PCRM) training method to mitigate this problem. PCRM incorporates prior constraints, specifically, length ratio and cosine similarity between outputs of each comparison pair, during reward model training to regulate optimization magnitude and control score margins. Experimental results demonstrate that PCRM significantly improves alignment performance by effectively constraining reward score scaling.
  arXiv  Detail & Related papers  (2024-04-01T07:49:11Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.