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

Problem Statement

Reinforcement Learning from Verifiable Rewards (RLVR) has become the standard paradigm for enhancing the reasoning capabilities of Large Language Models (LLMs). Concurrently, Multi-Token Prediction (MTP), an auxiliary objective trained during pretraining to predict multiple future tokens, has demonstrated significant benefits in sample efficiency and serving as a draft model for speculative decoding.

Intuitively, jointly optimizing MTP and RL objectives during post-training should allow models to benefit from multi-token supervision while acquiring reasoning abilities. However, empirical evidence shows that joint MTP-RL training often leads to significant performance degradation on the main model. Consequently, current practices (e.g., in frameworks like veRL and slime) typically detach MTP gradients from the main model during RL training. This detachment forces MTP parameters to adapt only through independent post-hoc fine-tuning, sacrificing the potential gains from simultaneous multi-token supervision. The core problem addressed is: Why does joint MTP-RL training fail, and can a principled strategy be designed to enable effective joint training?

Methodology

Theoretical Analysis

The authors approach the problem from an optimization perspective, analyzing the per-step effect of MTP on the RL objective $J(\theta)$. Assuming an $L$-smooth objective function, they derive a lower bound for the per-step improvement. This decomposition reveals that the effect of MTP is governed by two competing terms:

1. First-order Correlation Term: Determined by the directional agreement (inner product) between the RL gradient ($g_{RL}$) and the MTP gradient ($g_{MTP}$).
2. Second-order Perturbation Penalty: A non-positive term reflecting the perturbation introduced by the magnitude of the MTP gradient.

Performance improvement occurs only if the correlation term outweighs the penalty term.

Analysis of Existing Regimes

Using this decomposition, the paper unifies and explains three existing training regimes:

• Detach: $g_{MTP} = 0$ regarding the main model. $\Delta_{MTP} = 0$. This is safe but forfeits potential gains.
• Cross-Entropy (CE) Loss: The correlation term is near zero in expectation because RL up-weights high-reward samples while CE treats all samples equally. Consequently, the negative second-order penalty dominates, causing degradation.
• Policy Loss: Uses the same RL objective formulation for MTP. Initially, gradients are aligned, and the correlation term dominates, yielding gains. However, as training progresses and the model enters flat regions of the loss landscape, the correlation decays while the gradient norm (and thus the penalty) persists. This leads to a "rise-then-fall" performance curve where a fixed coefficient cannot track the phase transition.

Proposed Solution: Optimal Coefficient Calibration (OCC)

To address the phase transition, the authors propose Optimal Coefficient Calibration (OCC), an adaptive scheme that tracks the optimal MTP coefficient ($\lambda$) online.

1. Closed-Form Optimum: The analysis yields a closed-form expression for the optimal $\lambda^*$ at each step, which is proportional to the ratio of the gradient correlation ($c$) to the squared gradient norm ($v^2$).
2. Log-Probability Proxy: Computing full-model gradients for $c$ and $v^2$ is prohibitively expensive in large-scale distributed systems. OCC utilizes a log-probability proxy derived from the change in log-probability under small-step updates ($\delta \approx \log \pi_\theta - \log \pi_{old}$).
   • Alignment Proxy ($\hat{c}_t$): Estimates the inner product.
   • Variance Proxy ($\hat{v}^2_t$): Estimates the squared norm.
3. Adaptive Update: The coefficient is updated as $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$. This allows OCC to dynamically adjust the weight of MTP, maintaining positive gains throughout training without computing full gradients.

Key Contributions

1. Theoretical Unification: The paper provides a theoretical analysis that unifies Detach, Cross-Entropy, and Policy Loss regimes, explaining their success or failure based on the interplay between first-order correlation and second-order perturbation.
2. OCC Algorithm: It proposes Optimal Coefficient Calibration, an adaptive procedure that calibrates the MTP coefficient online using a log-probability proxy. This method requires negligible computational overhead and avoids full-model gradient computation.
3. Empirical Validation: Extensive experiments across multiple models (MiMo-7B, GLM-4.5-Air) and algorithms (DAPO, GSPO) on competition-level mathematical reasoning benchmarks demonstrate that OCC consistently matches or exceeds the Detach baseline.

Experimental Results

The authors evaluated OCC on six benchmarks: AIME24, AIME25, AMC, MATH, Minerva, and OlympiadBench.

• Performance: Across all benchmarks and model scales (7B dense and 106B MoE), OCC consistently outperformed the Detach baseline and the fixed-coefficient Policy Loss.
  • On MiMo-7B with DAPO, OCC achieved an average accuracy of 61.7%, surpassing Detach (58.9%) and Policy Loss (57.4%).
  • On GLM-4.5-Air, OCC achieved 67.6% average accuracy versus 65.9% for Detach.
• Training Dynamics: While Policy Loss showed a rise-then-fall trend (degrading in later stages), OCC maintained stable improvements throughout training. CE Loss consistently underperformed Detach.
• Efficiency: OCC introduced negligible overhead. The wall-clock training time per step for OCC (8.07s) was comparable to Detach (8.11s), whereas full-model backpropagation required 45.06s (a 5.6× slowdown).
• Ablation: Sensitivity analysis showed that the predefined scaling factor $\lambda_+$ is robust across a range of values, with $\lambda_+ = 1.0$ yielding the best average performance.

Significance and Claims

The paper claims that the failure of joint MTP-RL training is not inherent but stems from a mismatch between static coefficients and the dynamic drift of gradient alignment during training. By providing a theoretical decomposition of the optimization landscape and proposing OCC, the authors demonstrate that MTP can be safely re-introduced into RL post-training.

The significance lies in enabling models to simultaneously acquire reasoning capabilities via RL and leverage multi-token representations without sacrificing performance or incurring prohibitive computational costs. The work suggests that adaptive calibration is essential for joint training in non-convex deep learning landscapes, offering a principled alternative to the current industry standard of gradient detachment.

Limitations

The authors note that their theoretical analysis relies on $L$-smoothness assumptions which may not hold exactly in non-convex landscapes, providing qualitative rather than quantitative guidance. The log-probability proxy is a first-order approximation that may lose fidelity under large learning rates. Furthermore, the experiments are focused on mathematical reasoning with verifiable rewards; generalization to settings like RLHF with learned reward models or open-ended generation remains to be validated.