Source-linked AI summary
Stabilizing Reinforcement Learning with LLMs: Formulation and Practices
Chujie Zheng, Kai Dang, Bowen Yu, Mingze Li, Huiqiang Jiang, Junrong Lin, Yuqiong Liu, Hao Lin, Chencan Wu, Feng Hu, An Yang, Jingren Zhou, Junyang Lin
TL;DR
RL with LLMs must reconcile sequence-level rewards with token-level policy optimization, but the conditions making that surrogate valid are not fully explicit. The paper derives a first-order formulation and tests stabilization techniques, finding that controlling engine discrepancy and policy staleness supports stable training and comparable final performance after prolonged optimization.
Problem
RL with LLMs assigns rewards to complete responses while commonly optimizing token-level objectives, raising questions about the soundness and stability of this mismatch.
Method
The paper models the token-level policy-gradient objective as a first-order approximation to sequence-level reward optimization and analyzes training–inference discrepancy, policy staleness, and MoE routing.
Results
On-policy policy gradient with importance sampling correction is most stable; off-policy training requires clipping and Routing Replay, while stabilized models reach comparable final performance across cold starts.
Takeaways & Limitations
Importance sampling correction, clipping, and Routing Replay provide a principled recipe for stabilizing RL training across on-policy and off-policy settings.
Takeaways & Limitations
The paper reports that its on-policy findings for R3 contradict prior work, potentially because prior validation used smaller experiments and BF16 rather than the paper’s FP8 inference setting.
Abstract
from arXiv · showhide
This paper proposes a novel formulation for reinforcement learning (RL) with large language models, explaining why and under what conditions the true sequence-level reward can be optimized via a surrogate token-level objective in policy gradient methods such as REINFORCE. Specifically, through a first-order approximation, we show that this surrogate becomes increasingly valid only when both the training-inference discrepancy and policy staleness are minimized. This insight provides a principled explanation for the crucial role of several widely adopted techniques in stabilizing RL training, including importance sampling correction, clipping, and particularly Routing Replay for Mixture-of-Experts (MoE) models. Through extensive experiments with a 30B MoE model totaling hundreds of thousands of GPU hours, we show that for on-policy training, the basic policy gradient algorithm with importance sampling correction achieves the highest training stability. When off-policy updates are introduced to accelerate convergence, combining clipping and Routing Replay becomes essential to mitigate the instability caused by policy staleness. Notably, once training is stabilized, prolonged optimization consistently yields comparable final performance regardless of cold-start initialization. We hope that the shared insights and the developed recipes for stable RL training will facilitate future research.
1 Introduction
The paper formulates when token-level policy-gradient objectives can validly optimize sequence-level rewards, linking this validity to training stability. It also derives practical stabilizing roles for importance sampling, clipping, and Routing Replay, then validates recipes empirically.
- Motivation: The paper frames the central problem as a mismatch between sequence-level rewards and token-level optimization objectives in RL with LLMs.This mismatch raises concerns about the soundness and stability of mainstream methods such as REINFORCE and GRPO.
- Stabilization techniques: Importance sampling correction is inherent to the surrogate objective, clipping limits policy staleness, and Routing Replay reduces both discrepancies for MoE models.Routing Replay fixes routed experts during policy optimization.
- Empirical findings: For on-policy training, policy gradient with importance sampling correction is most stable; off-policy training requires clipping with Routing Replay, while stabilized runs reach comparable final performance across cold starts.The off-policy combination mitigates instability attributed to policy staleness.
- Formulation: The proposed formulation treats the token-level objective as a first-order approximation to the expected sequence-level reward, valid when training–inference discrepancy and policy staleness are minimized.The formulation provides conditions under which token-level optimization is justified.
- Empirical validation: Extensive experiments with a 30B MoE model spanning hundreds of thousands of GPU hours evaluate practical recipes for stable RL training.The experiments cover on-policy and off-policy settings and compare cold-start initializations after stabilization.
2 Formulation for Reinforcement Learning with LLMs
The formulation starts from an intractable sequence-level reward gradient and derives a token-level surrogate through importance sampling and a first-order approximation. Its validity depends on keeping training–inference discrepancy and policy staleness sufficiently small, including in MoE routing.
- 2.1 Notation: An autoregressive LLM is represented as a policy πθ, with x denoting an input prompt, D the prompt set, and y a response.The response likelihood factorizes over tokens, and |y| denotes response length.
- 2.1 Notation: The formulation focuses on sequence-level rewards R(x, y) assigned to complete responses rather than value-based token scores.The paper excludes value-based settings because reliable, general, and scalable value models are difficult to devise.
- 2.2 Expected Sequence-level Reward is Hard to Directly Optimize: The target policy πθ and rollout policy µθold are distinguished because responses are sampled in an inference engine and optimized in a training engine.Importance sampling transforms expectations over the rollout policy into an objective involving the target policy.
- 2.2 Expected Sequence-level Reward is Hard to Directly Optimize: On-policy means rollout and target policies are identical apart from training–inference discrepancy, whereas off-policy means the two policies differ.This distinction is used to characterize policy mismatch in subsequent optimization settings.
- 2.2 Expected Sequence-level Reward is Hard to Directly Optimize: The exact sequence-level gradient is difficult to optimize directly because sequence likelihoods have a large numerical range and high variance.This motivates replacing direct sequence-level optimization with a token-level surrogate.
- 2.3 Token-level Objective as a First-order Approximation to Sequence-level Objective: The surrogate token-level objective is basic REINFORCE equipped with token-level importance-sampling weights.It is introduced as an approximation to the sequence-level objective rather than as an unrelated optimization target.
- 2.3 Token-level Objective as a First-order Approximation to Sequence-level Objective: The approximation neglects second-order and higher-order products of small policy-ratio deviations such as δ_iδ_j.This first-order treatment explains why the surrogate tracks the sequence-level objective when the policies are close.
- 2.3 Token-level Objective as a First-order Approximation to Sequence-level Objective: When πθ is close to µθold, updating parameters with the surrogate gradient can improve the sequence-level objective.The paper explicitly identifies the token-level objective as a first-order approximation under this closeness condition.
3 Challenge for Mixture of Experts, and Routing Replay
MoE routing entangles expert selection with training–inference discrepancy and policy staleness, weakening the first-order approximation behind token-level optimization. Routing Replay fixes routed experts to reduce these discrepancies, but may bias the target policy, with R2 and R3 differing in how they do so.
- Challenge for Mixture of Experts: MoE routing can make the first-order approximation break down because training–inference discrepancy and policy staleness alter the routed experts and final outputs.The discrepancy may arise from inconsistent expert choices across engines, while staleness also changes the activated experts as parameters update.
- Routing Replay: Routing Replay fixes routed experts during policy optimization, allowing the MoE model to be optimized more like a dense model.The approach targets the routing component that undermines the approximation.
- Routing Replay: R2 replays rollout-policy experts in the training engine to mitigate routing-driven policy staleness, whereas R3 replays inference-engine experts to reduce training–inference discrepancy and also policy staleness.R2 and R3 therefore target different discrepancy sources and can alter the target policy differently.
- Routing Replay: Routing Replay reduces discrepancy but implicitly biases the target policy by constraining experts away from their naturally routed choices.The original target policy uses naturally routed experts, while replay changes which experts govern token likelihoods.
- Routing Replay: R2 and R3 may perform differently as off-policiness varies because R2 preserves the original target policy in the first mini-batch while R3 does not.The paper identifies the batch-to-mini-batch ratio as a relevant factor in their differing degrees of bias.
- Routing Replay: The paper leaves the net benefit of Routing Replay unresolved because its approximation-restoring effects must be weighed against the bias introduced by altering routed experts.Further experiments are needed to assess its practical utility and the balance between advantages and disadvantages.
4 Empirical Analyses
Experiments test stable RL recipes under on-policy and increasingly off-policy settings, finding that stability depends on preserving the first-order approximation through importance sampling, clipping, and Routing Replay. Once stabilized, different cold-start initializations achieve comparable final performance.
- 4.3 Results of On-policy Training: MiniRL, the basic policy gradient algorithm with importance sampling correction, achieves the best performance and training stability in on-policy training.Removing training–inference importance sampling correction causes rapid collapse, while length normalization remains stable but is suboptimal.
- 4.3 Results of On-policy Training: Removing training–inference importance sampling correction causes rapid training collapse and a sharp drop in entropy.The experiments identify the importance sampling weight as an inherent component of the first-order approximation.
- 4 Empirical Analyses: These results support designing token-level objectives that preserve the first-order approximation to the expected sequence-level reward.The experiments validate the proposed formulation as an explanation for stable RL training.
- 4.4 Results of Off-policy Training: Under off-policy updates, combining Routing Replay and clipping is necessary for stable training, whereas omitting either causes premature collapse and lower peak performance.Routing Replay alleviates expert-routing effects, while clipping prevents aggressive policy updates and restrains policy staleness.
- 4.4 Results of Off-policy Training: R2 performs better when off-policiness is small, whereas R3 surpasses R2 at larger off-policiness and becomes necessary for stable training.R2 is sufficient at gbs = 2 × mbs, while R3 is favored at gbs = 4 × mbs and gbs = 8 × mbs.
- 4.5 Results of Varying Cold-start Initializations: The three cold-start initializations consistently achieve comparable final performance when trained with a stable RL recipe.The study also reports similar peak performance across stabilized on-policy and off-policy training.
5 Conclusion
The paper frames token-level RL objectives as first-order approximations to sequence-level rewards and identifies techniques that preserve this approximation for stable training. It also reports that stabilized RL converges to similar performance across prolonged training despite different cold-starts.
- The formulation treats the token-level optimization objective as a first-order approximation to the true expected sequence-level reward.
- Importance sampling correction, clipping, and Routing Replay for MoE models stabilize RL training by preserving the approximation’s validity.
- Once RL training is stabilized, prolonged optimization consistently produces similar performance from the same base model across cold-start initializations.
A Comparison of MiniRL against GRPO and CISPO
The paper compares MiniRL with GRPO and CISPO and identifies objective-design differences that affect approximation validity and training stability. In particular, training–inference correction, length normalization, and gradient clipping distinguish the methods.
- GRPO and CISPO differ from MiniRL because their original objectives omit training–inference discrepancy correction and use length normalization.
- Length normalization invalidates the first-order approximation to expected sequence-level reward, producing a biased token-level objective and potentially suboptimal performance.
- CISPO’s omission of clipping for certain token gradients can result in unstable training.
B Detailed Benchmark Results
This section presents detailed benchmark-result figures for on-policy training, multiple off-policy batch configurations, and varying cold-start initializations. The supplied passages identify the benchmark scopes but do not report their numerical outcomes.
- Figure 6 reports detailed benchmark results for on-policy training with gbs = mbs = 1, 024.
- Figures 7–9 report detailed benchmark results for off-policy training with gbs = 2 × mbs = 2, 048, gbs = 4 × mbs = 4, 096, and gbs = 8 × mbs = 8, 192.
- Figure 10 reports detailed benchmark results for varying cold-start initializations.