Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

Abstract: Reinforcement Learning (RL) has emerged as a pivotal post-training paradigm, yet it frequently suffers from unpredictable sub-optimum performance or even training collapses. Recent findings attribute these failures to a hidden train-inference discrepancy (or mismatch), stemming from the disparate underlying engines and architecture. We find that the training policy can actively self-correct such a discrepancy when provided with an appropriate learning signal. Then, we further empirically identify a discrepancy tolerance region: within this region, aggressively narrowing the discrepancy can suppress policy exploration and reduce learning efficiency, whereas outside this region, reducing excessive discrepancy improves optimization consistency and raises the achievable local performance ceiling. According to such findings, we formulate this problem as a Discrepancy-Constrained Markov Decision Process (DCMDP), where reward maximization is coupled with a constraint that aligns training-Inference behavior, achieving stable dual-objective optimization. To adaptively balance performance improvement and discrepancy control, we introduce a Lagrangian relaxation mechanism that dynamically adjusts the relative weight of the two objectives according to the current degree of discrepancy violation. This enables stable dual-objective optimization: the policy is allowed to explore freely within the tolerance region, while being guided back when the discrepancy exceeds the safe boundary. Empirically, DCMDP significantly improves the performance of 8B dense model (Qwen-3-8b) and 30B Mixture-of-Expert model (Qwen-3-30bA3b), and enables a heterogeneous training paradigm, where LLMs can be optimized in high-fidelity training setup while being explicitly aligned for low-cost, resource-constrained inference deployment.

• The paper introduces the DCMDP framework to mitigate train-inference discrepancies by employing token-level probability difference penalties.
• It demonstrates that adaptive discrepancy regularization improves training stability and final accuracy across various LLM architectures and quantization regimes.
• Empirical results validate enhanced reward optimization and elimination of training instabilities in both dense and MoE LLMs deployed under heterogeneous conditions.

Discrepancy-Constrained Markov Decision Process for Robust LLM RL Optimization Under Black-Box Train-Inference Mismatch

Introduction

The paper "Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy" (2606.08779) addresses a critical and under-explored instability in post-hoc RL optimization of LLMs: the pervasive mismatch between the training and inference policies, induced by infrastructural and architectural divergences such as disparate numerical precisions and backend kernel implementations. This so-called train-inference discrepancy is shown to be a dominant cause of training collapse and sub-optimal deployment performance, especially as models and inference pipelines become increasingly decoupled for practical, large-scale deployments. The authors dissect the failure modes of RL post-training, formalize the mismatch challenge as a constrained optimization problem, and introduce the Discrepancy-Constrained Markov Decision Process (DCMDP) as a theoretically principled and empirically validated solution.

Discrepancy-Aware RL: Core Formulation and Theoretical Insights

The key empirical motivation is the observation that aligning training and inference engines at the system level is labor-intensive and often infeasible due to nontrivial kernel and parallelization differences (e.g., FSDP vs. Megatron-LM, model quantization, rollout determinism). Importantly, mere precision matching (as with uniform FP16) either damages training stability due to saturating representable ranges or fails to bridge operator-induced mismatches.

The paper leverages three token-level discrepancy measures—probability difference, probability ratio, and log-probability difference—to quantitatively capture divergence between backend-executed policies. Their theoretical analysis, including closed-form stationary points for each regularizer in a bandit setting, highlights the critical geometric difference: the probability difference penalty enables more adaptive, tail-seeking corrections without strictly suppressing reward-improving updates in low-probability actions, in contrast to the conservative regularization exerted by the probability ratio (see Figures 3 and 4 for comparison of induced optimal policies). This has direct implications for the learnability and plasticity of RL-trained transformers.

Figure 1: Effect of different discrepancy penalties under varying mismatch budgets; the probability difference penalty supports aggressive tail reallocation enabling better alignment with high-advantage, low-probability actions.

Figure 2: One-sided variant of the analysis; penalty selection determines how adaptively the optimizer reallocates mass under constrained drift.

Empirically, experiments confirm that probability difference regularization provides the best tradeoff between reward maximization and discrepancy minimization over mathematical reasoning benchmarks, outperforming ratio and log-difference penalties in both stability and final accuracy metrics (see Figure 3, right panel).

Figure 3: Performance and discrepancy evolution under different penalties; probability difference achieves superior avg@32 benchmark performance and effective train-inference gap suppression.

Discrepancy Tolerance and the DCMDP Paradigm

A core contribution is the empirical identification and exploitation of the discrepancy tolerance region: policies exhibit an inherent robustness to minor train-inference deviations, and over-regularization at low discrepancy hinders exploratory search and thus caps attainable reward. Imposing penalties only when the discrepancy surpasses an adaptive threshold enables the optimizer to preserve exploratory reward-seeking within the safe tolerance region (see Figures 2 and 5).

The authors formalize this as a constrained optimization—maximizing the primary reward under a hard ceiling on the average per-token probability discrepancy, instantiated as:

$$\max_\theta \mathbb{E}_{\tau} \left[ \sum_t r(o_t \mid q, o_{<t}) \right] \quad \text{s.t.}\quad \frac{1}{|\tau|}\sum_t |\pi^{\text{Train}}(o_t) - \pi^{\text{Inf}}(o_t)| \leq c$$

To enforce this adaptively and efficiently, they instantiate a dual-objective formulation via Lagrangian relaxation, where the penalty weight $\lambda$ is dynamically tuned. The penalty is only activated when the discrepancy constraint is violated, and the multiplier anneals away in the tolerance interior, thus unifying reward maximization and robust alignment control in a mathematically stable fashion.

Figure 4: (Left) DCMDP directs the training policy towards inference-oriented optimal solutions by injecting real-time discrepancy feedback. (Right) DCMDP enables seamless black-box handling for heterogeneous infrastructure scenarios.

Efficient Algorithm: DC-GRPO

The proposed method, DC-GRPO, implements DCMDP atop the Group Relative Policy Optimization backbone. The implementation requires forward passes under both training and inference backends to extract the black-box token-level discrepancies, then injects the adaptively-weighted penalty signal on the group-level normalized advantage. The dual variable $\lambda$ is updated via dual gradient ascent at each iteration, maintaining average discrepancy control at scale.

The overhead is minor: only an extra pass per trajectory on the training backend, and a single scalar variable learned per training step.

Figure 3: The abscissa shows training steps; curves depict the efficacy of temperature-controlled and discrepancy-aware exploration in stabilizing reward and train-inference gap.

Empirical Results

Evaluation is conducted on two open-source LLM architectures: dense (Qwen3-8B) and MoE (Qwen3-30B-A3b), spanning major system backends (FSDP2, Megatron-LM) with corresponding distinct mismatch patterns. Benchmarks include six rigorous mathematical reasoning tasks. Key findings include:

• Substantial improvement in average accuracy on all benchmarks versus baseline GRPO under both BF16 and heterogeneous BF16/FP8 regimes
• Complete elimination of instability and collapse events even when the train-inference gap is intensified by quantized deployment backends
• The method consistently controls final train-inference discrepancy to within the empirical tolerance region while delivering enhanced primary reward compared to over-regularized or unconstrained approaches

Ablation studies show that the approach is robust to $\lambda$ initialization, dual learning rate, and tolerance budget parameters—the adaptive Lagrangian control adequately compensates for a range of settings, as long as the budget is not trivially tight or loose.

Theoretical and Practical Implications

The DCMDP-based approach constitutes a practical framework for black-box, infrastructure-agnostic RL post-training under unavoidable train-inference mismatch, without relying on costly, brittle system-level alignment. By introducing a tolerance-aware discrepancy constraint directly at the algorithmic layer, the method realizes several practical advancements:

• Unlocks heterogeneous training/deployment paradigms: High-fidelity, resource-heavy training backends can now be used to optimize LLMs for low-cost, quantized inference environments without performance collapse or instability
• Provides a versatile template for dealing with diverse architectural and infrastructural sources of deployment drift, as the discrepancy is treated as a dynamic black-box signal, abstracting away details of model implementation or backend code

Broader Implications and Future Directions

DCMDP reframes system-level RL optimization in LLMs as a constrained, dual-objective black-box problem, offering a robust and generalizable abstraction that is orthogonal to backend specificity or reward type. While the empirical focus is on BF16 and FP8, the architecture supports extension to other quantization regimes and hardware-specific dynamics.

Future research directions include:

• Extension to more diverse backend scenarios, such as aggressive model pruning, mixed-precision activation, or extreme expert routing non-determinism in MoEs
• Application of DCMDP constraints in long-horizon reasoning and multi-agent RL setups where inference drift may be more subtle and compounding

Conclusion

The paper provides a theoretically principled, empirically validated, and practically deployable framework for mitigating train-inference discrepancy in RL-based LLM training. By formulating and solving the optimization as a DCMDP and employing token-level probability difference in a dynamically controlled dual-objective schema, the method stabilizes training, enables robust reward maximization, and prescribes a scalable path for deployment-robust post-training pipelines, especially in highly heterogeneous, large-scale environments (2606.08779).