Source-linked AI summary
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
TL;DR
Jointly training MTP with RL post-training is attractive because both improve LLM capabilities, yet existing practice detaches MTP gradients after joint training degrades performance. This paper explains the degradation through first- and second-order effects, then introduces OCC, which adaptively calibrates MTP coefficients and consistently matches or exceeds detached training across mathematical reasoning benchmarks.
Problem
Joint MTP-RL training degrades performance, leaving unclear why it fails and whether it can be made effective.
Method
The paper decomposes MTP’s per-step RL effect into correlation and perturbation terms, then uses OCC to adaptively calibrate the MTP coefficient through a log-probability proxy.
Results
Across multiple competition-level mathematical reasoning benchmarks, OCC consistently matches or exceeds the Detach baseline.
Takeaways & Limitations
The analysis supports safely reintroducing MTP into RL post-training with adaptive coefficient calibration.
Takeaways & Limitations
The theory relies on L-smoothness, the proxy is a first-order approximation, and experiments focus on mathematical reasoning with verifiable rewards.
Abstract
from arXiv · showhide
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.
1 Introduction
Joint MTP-RL training is motivated by the complementary strengths of RLVR and MTP but has historically degraded main-model performance. The paper analyzes this failure through gradient interactions and proposes OCC, which consistently matches or exceeds detached training across benchmarks.
- RLVR improves LLM reasoning, while MTP supports multi-step representations, downstream accuracy, and speculative decoding.
- Previous joint MTP-RL training degraded main-model performance, motivating gradient detachment in current RL practices.
- MTP’s per-step effect decomposes into a first-order gradient-correlation term and a second-order perturbation penalty.Performance improves when the correlation outweighs the penalty.
- Detach isolates MTP gradients, CE loss produces weak correlation with RL gradients, and policy loss initially aligns but later suffers decaying alignment and persistent perturbation.These regimes explain detached stability, CE underperformance, and policy loss’s rise-then-fall behavior.
- OCC adaptively calibrates the MTP coefficient online using a log-probability proxy with negligible computational overhead.The proxy avoids full-model gradient computation while tracking the theoretically preferred coefficient.
- Across multiple competition-level mathematical reasoning benchmarks, OCC consistently matches or exceeds the Detach baseline.The experiments span different models and algorithms and report stable improvements across tasks.
2 Related Work
MTP and RLVR are established components of LLM training, but their joint use during RL post-training remains unresolved. Prior systems commonly detach or otherwise isolate MTP gradients, while no prior work had provided a theoretical explanation and principled joint-training method.
- Multi-Token Prediction: MTP predicts multiple future tokens and has improved sample efficiency and downstream task performance during pretraining.
- Reinforcement Learning Post-Training: RLVR has become a standard paradigm for enhancing LLM reasoning capabilities, with PPO, GRPO, and DAPO refining RL post-training.
- Joint MTP-RL Training: RL frameworks document severe degradation when MTP gradients reach the main model and recommend gradient detachment as the stable default.
- Joint MTP-RL Training: Prior work had not theoretically explained joint MTP-RL failure or proposed a principled method for enabling joint training.
3 Joint MTP-RL Training
The analysis decomposes MTP’s per-step effect on the RL objective into a first-order gradient-correlation term and a second-order perturbation penalty, explaining why training regimes behave differently. It further shows that policy-loss benefits can reverse during training, motivating adaptive coefficient calibration.
- 3.1 Theoretical Framework: Effect of MTP on RL Objective: MTP changes the RL improvement bound through a first-order correlation between RL and MTP gradients and a second-order perturbation penalty.The correlation can improve the objective, while the penalty is nonpositive; their relative magnitudes determine whether MTP helps or hurts.
- 3.2 Analysis of Three Training Regimes: Detach sets the MTP gradient to zero for the main model, so MTP does not affect its RL update.
- 3.2 Analysis of Three Training Regimes: Cross-entropy loss underperforms because its near-zero expected correlation with RL gradients leaves the second-order penalty dominant.RL weights samples by rewards, whereas cross-entropy treats samples equally, producing conflicting directions and persistent variance.
- 3.3 Degradation Behind Policy Loss: Policy loss initially improves training because its positive gradient correlation dominates, but later degrades it as correlation decays while the MTP gradient norm persists.The resulting phase transition flips the MTP contribution from positive to negative.
- 3.3 Degradation Behind Policy Loss: A fixed MTP coefficient becomes unsuitable as the policy-aligned gain drifts from positive early in training to negative later.The analysis therefore motivates recalibrating the coefficient as training progresses.
- 3.4 Optimal Coefficient Calibration (OCC): OCC tracks the theoretically preferred coefficient online using log-probability changes as a low-cost proxy for gradient alignment and variance.The proxy avoids computing the full-model gradient in large-scale distributed training while estimating the ratio governing the optimum.
4 Experiments
Experiments evaluate OCC across datasets, RL algorithms, model scales, training regimes, proxy fidelity, coefficient choices, and computational overhead. OCC consistently outperforms or matches Detach while tracking training dynamics at nearly Detach’s cost.
- Setup: Experiments use six mathematical reasoning benchmarks, DAPO and GSPO, and models ranging from MiMo-7B-RL to the 106B-A12B GLM-4.5-Air.Results report average accuracy over 32 runs, comparing Detach, CE Loss, Policy Loss, and OCC.
- Main results: 47.7 vs. 58.9 average accuracy: CE loss falls far below Detach, with a -29.2-point gap on AIME24.The result is attributed to weak first-order correlation and an unmitigated second-order penalty.
- Main results: 61.7 average accuracy: OCC exceeds Detach by +2.8 points and Policy Loss by +4.3 points, with gains across several benchmarks.OCC reaches its largest reported margin on AIME25, at +10.0 points over Detach and Policy Loss.
- Generalization: 60.1 average accuracy with GSPO versus 57.8 for Detach (+2.3), while OCC remains uniformly best across benchmarks.The same ordering appears under DAPO, supporting generalization across RL algorithms.
- Generalization: 67.6 average accuracy versus 65.9 for Detach (+1.7) on GLM-4.5-Air, with gains on every benchmark.The reported evaluation spans a dense 7B model and a 106B MoE model.
- Efficiency: 8.07 s per step: OCC is comparable to Detach at 8.11 s, whereas full-model backpropagation requires 45.06 s, a 5.6× slowdown.The log-probability proxy avoids exact full-model gradient computation.
- Ablation study: All fixed Policy Loss coefficients are surpassed by OCC on average and every benchmark, while small coefficients resemble Detach and large coefficients collapse late.The sweep uses λ ∈ {0.1, 0.2, 0.5, 1.0}; OCC’s advantage is attributed to dynamic adjustment rather than constant tuning.
5 Conclusion
The paper explains joint MTP-RL failure through a first-order correlation and second-order penalty, then introduces OCC to adapt the MTP coefficient online. Across multiple benchmarks, OCC matches or exceeds Detach, showing that MTP can be reintroduced into RL post-training with consistent gains.
- Conclusion: A per-step improvement bound decomposes MTP’s effect into first-order correlation and second-order penalty terms.The framework unifies Detach, Cross-Entropy, and Policy Loss and exposes a training phase transition that static coefficients cannot track.
- Conclusion: OCC sets the MTP coefficient to its closed-form optimum at every step using a log-probability proxy at negligible cost.The method is designed to avoid full-model gradient computation in large-scale training.
- Conclusion: OCC consistently matches or exceeds the Detach baseline across multiple benchmarks.The reported conclusion is that MTP can be safely reintroduced into RL post-training while delivering consistent gains.
6 Limitations
The analysis has several scope and modeling limitations: its smoothness-based guidance is qualitative, its proxy may lose fidelity under aggressive updates, and empirical coverage is limited.
- The theoretical bound assumes L-smoothness, which may not hold exactly in nonconvex deep-learning landscapes and provides qualitative rather than quantitative guidance.
- The log-probability proxy is a first-order approximation that may lose fidelity under large learning rates or aggressive policy updates.
- Experiments focus on mathematical reasoning with verifiable rewards, leaving generalization to RLHF and open-ended generation unvalidated.
- The predefined ratio λ+ absorbs an unknown smoothness-dependent prefactor, so its optimal choice may vary across training configurations.
A Detailed Proof of the Policy Improvement Lower Bound
The proof applies the L-smoothness condition to the combined RL and MTP update, expands its inner-product and squared-norm terms, and establishes positivity of the standard RL improvement term under η < 1/L.
- A Detailed Proof of the Policy Improvement Lower Bound: The proof assumes η < 1/L so the quadratic approximation remains valid within the update neighborhood.
- A Detailed Proof of the Policy Improvement Lower Bound: Setting θ′ = θt+1 = θt + η(gRL + λgMTP) inserts the combined RL–MTP update into the smoothness bound.
- A Detailed Proof of the Policy Improvement Lower Bound: The inner product expands into the standard RL term ∥gRL∥2 and the first-order interaction λ⟨gRL, gMTP⟩.
- A Detailed Proof of the Policy Improvement Lower Bound: The squared norm is expanded and grouped by ∥gRL∥2, ⟨gRL, gMTP⟩, and ∥gMTP∥2 to isolate the update contributions.
- A Detailed Proof of the Policy Improvement Lower Bound: Under η < 1/L, coefficients (1 − Lη) and (1 − Lη/2) are positive, making the standard RL improvement term strictly positive when gRL ≠ 0.
B Vanishing First-Order Term under Cross-Entropy MTP
Under approximate independence and zero-mean group-normalized advantages, cross-entropy MTP contributes no expected first-order correlation with the RL update, leaving a quadratic perturbation penalty.
- B Vanishing First-Order Term under Cross-Entropy MTP: In standard RLVR batches, cross-sample gradient inner products are small and vanish in expectation when summed against zero-mean advantages.
- B Vanishing First-Order Term under Cross-Entropy MTP: Assuming advantages are approximately independent of per-sample gradient norms lets the expectation factor into an advantage mean and a gradient-norm term.
- B Vanishing First-Order Term under Cross-Entropy MTP: The group-normalized advantage satisfies E[Ai] = 0 by construction, causing the expected first-order term to vanish.
- B Vanishing First-Order Term under Cross-Entropy MTP: CE loss injects pure perturbation without correlation gain, with damage scaling quadratically with λ.
C Derivation of the Optimal Adaptive Weight
The adaptive-weight derivation treats the MTP effect as a concave quadratic in λ and selects the coefficient at its maximum, using the current correlation and MTP-gradient magnitude.
- C Derivation of the Optimal Adaptive Weight: Treating ΔMTP as a function of λ yields a concave quadratic because the λ² coefficient is negative.
- C Derivation of the Optimal Adaptive Weight: The quadratic depends on c = ⟨gRL, gMTP⟩ and v2 = ∥gMTP∥2, representing gradient correlation and MTP-gradient magnitude.
- C Derivation of the Optimal Adaptive Weight: Setting f′(λ) = 0 identifies the coefficient that maximizes the per-step perturbation improvement.
D Experiment Details
Experiments compare four joint MTP-RL training regimes under shared configurations, using competition-level mathematical reasoning evaluation with avg@32 accuracy.
- Training setup: The four regimes differ only in MTP gradient flow into the main model and the MTP loss coefficient.The regimes are Detach, CE loss, policy loss, and OCC.
- Reported metrics: Table 4 reports final avg@32 accuracy (%) on MiMo-7B-RL + DAPO, with the best result shown in bold.The table compares the clipping variants in the adaptive coefficient experiment.
- Datasets and licensing: The evaluation resources include AIME24, AIME25, MATH, MATH-500, Minerva Math, OlympiadBench, AMC23, and the DeepScaleR framework.The passages specify licenses for these datasets and evaluation resources, with AMC23 lacking an explicit license declaration on its current dataset card.
- Evaluation protocol: Evaluation uses 32 independent samples per prompt, with temperature 1.0 and top-p = 0.7.Reported results are averages over 32 decodes using the DeepScaler framework.
E Clipping the Adaptive Coefficient
The clipping ablation tests whether OCC should preserve signed alignment estimates or clip them to non-negative values. Preserving the sign performs better because negative coefficients can reverse harmful MTP directions.
- Ablation result: OCC-NoClip outperforms OCC-Clip by +1.3 average points, including a +2.6 gap on AIME25.The comparison is conducted on MiMo-7B-RL + DAPO.
- Coefficient design: The adaptive coefficient is proportional to the signed alignment ratio ˆct/ˆv2_t, where ˆct measures RL–MTP gradient alignment.This motivates retaining negative alignment values rather than clipping them.
- Mechanism: When alignment is negative, a negative MTP coefficient flips the auxiliary gradient so it aligns with the RL direction.Clipping the estimate to zero prevents these corrective updates, whereas OCC-NoClip retains them.
- Scope caveat: The paper notes that RLVR-based improvements may increase sensitivity to reward-signal design and quality, motivating robustness checks.This is identified as a potential broader-setting risk rather than a clipping result.