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

Technical Summary: Reformulating LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

1. Problem Statement

Reinforcement Learning (RL) has become a critical post-training paradigm for enhancing the reasoning capabilities of Large Language Models (LLMs). However, RL post-training frequently suffers from unpredictable sub-optimal performance or complete training collapses. Recent investigations attribute these failures to a pervasive train-inference discrepancy (or mismatch).

This discrepancy arises because industrial demands often force practitioners to use disparate underlying engines and precision formats (e.g., different kernels, parallelization strategies, or numerical precisions like BF16 vs. FP8) for the exact same model weights during training versus inference. This creates a hidden deviation between the action probability distributions of the training policy ($\pi_{Train}$) and the inference policy ($\pi_{Inf}$).

Existing mitigation strategies have significant limitations:

• Uniform Precision Alignment: Simply downgrading training precision to match inference (e.g., to FP16) shrinks the representable value range and exacerbates gradient explosion risks.
• Heuristic Interventions: Methods like masking gradients from high-discrepancy tokens merely suppress the impact of noisy data. They fail to address the root cause and overlook the ultimate goal: optimizing the policy for the inference deployment, not just the training process.
• Infrastructure Alignment: Attempting to align low-level operators and hardware stacks is labor-intensive and often infeasible due to the myriad of corner cases.

The core challenge is to enable the training policy to dynamically perceive and actively nullify its discrepancy with the inference policy without sacrificing the primary reward maximization objective.

2. Methodology: Discrepancy-Constrained MDP (DCMDP)

The authors propose a paradigm shift from standard Markov Decision Processes (MDP) to a Discrepancy-Constrained MDP (DCMDP). This framework treats the train-inference discrepancy as a black-box optimization problem, allowing the system to handle heterogeneous training and inference setups seamlessly.

Key Components:

• Discrepancy Metric:
  The authors empirically evaluate three token-level discrepancy signals: Probability Difference ($\delta_{diff}$), Probability Ratio ($\delta_{ratio}$), and Log-Probability Difference ($\delta_{log}$). They identify that the absolute probability difference ($\delta_{diff} = |p_{Train} - p_{Inf}|$) is the most effective penalty. It offers a "plastic" regularization that corrects deployment-relevant drift without overly constraining reward-driven exploration, unlike ratio-based penalties which can be too conservative for low-probability tokens.

• The Discrepancy Tolerance Region:
  A critical empirical finding is the existence of a discrepancy tolerance region. Aggressively penalizing minor, negligible discrepancies early in training suppresses policy exploration and reduces learning efficiency. Conversely, reducing excessive discrepancy outside this region improves optimization consistency.

  • Mechanism: The DCMDP formulation only applies penalties when the discrepancy exceeds a predefined threshold $c$. Within this tolerance region, the policy is allowed to explore freely to maximize rewards.

• Dual-Objective Optimization via Lagrangian Relaxation:
  The problem is formulated as maximizing the expected reward subject to a constraint on the average token-level discrepancy:
  $$\max_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum r(o_t) \right]$$
  $$\text{s.t. } \frac{1}{|\tau|} \sum |\pi_{Train}(o_t) - \pi_{Inf}(o_t)| \leq c$$

  To solve this, the authors employ Lagrangian Relaxation. A learnable multiplier $\lambda$ dynamically adjusts the weight of the discrepancy penalty:

  • If the discrepancy violates the tolerance boundary ($>c$), $\lambda$ increases, strengthening the penalty to pull the policy back to the safe zone.
  • If the discrepancy is within the boundary, $\lambda$ anneals toward zero, allowing the policy to focus on reward maximization.

  The final objective (Dual-GRPO) injects this adaptive penalty directly into the group-normalized advantage:
  $$\hat{A}^{DCMDP}_{i,t} = \hat{A}_{i,t} - \lambda_k \cdot \delta_{i,t}$$

3. Key Contributions

1. Autonomous Discrepancy Compensation: The paper demonstrates that incorporating a robust, black-box penalty (token-level absolute probability difference) endows the training policy with the capacity to autonomously eradicate the risk of discrepancy-induced training collapse.
2. Paradigm Shift to DCMDP: The authors reformulate LLM RL as a Discrepancy-Constrained MDP. This approach optimally balances the dual objectives of maximizing task reward and minimizing train-inference discrepancy, moving beyond naive scalarization of objectives.
3. Heterogeneous Training Paradigm: The method enables a training setup where models are optimized in high-fidelity, resource-intensive environments (e.g., BF16 precision) while being explicitly aligned for low-cost, resource-constrained inference deployments (e.g., FP8). This turns infrastructure heterogeneity from a liability into a controllable optimization knob.

4. Experimental Results

The method was evaluated on Qwen3-8B (dense) and Qwen3-30B-A3B (Mixture-of-Experts) models across six mathematical reasoning benchmarks (AIME24/25, AMC23, MATH-500, Minerva, OlympiadBench).

• BF16 Training/Inference: DC-GRPO consistently outperformed the vanilla GRPO baseline across all benchmarks for both models, achieving higher average scores (e.g., Qwen3-8B improved from 67.3 to 71.0). This confirms that enforcing discrepancy constraints enhances, rather than hinders, primary reward optimization.
• Heterogeneous Training (BF16) + Deployment (FP8): In this setting, vanilla GRPO suffered significant performance degradation or collapse due to the amplified train-inference gap. DC-GRPO successfully mitigated this, preserving and often exceeding the performance of the BF16-only baseline.
• Ablation Studies: The method proved robust to variations in the initial Lagrangian multiplier ($\lambda_0$) and dual learning rate ($\eta_\lambda$). The tolerance budget ($c$) was found to be critical; setting it too small caused over-regularization, while setting it too large caused the method to degenerate into vanilla GRPO.

5. Significance and Claims

The paper claims that DCMDP provides a unified, algorithmic framework for mitigating train-inference instability without requiring exhaustive, bespoke optimizations of the underlying hardware or software stacks. By treating the discrepancy as a black-box feedback signal, the method allows for:

• Stable Training: Preventing the abrupt collapses common in RL post-training.
• Efficient Exploration: Allowing the policy to explore freely within a safe discrepancy region.
• Resource Efficiency: Enabling the use of high-fidelity training setups to simulate and adapt to constrained deployment environments, thereby facilitating more efficient real-world model deployments.

The authors acknowledge limitations, noting that their empirical scope is currently limited to BF16 and FP8 precisions and specific model families, with future work needed to extend the method to other backend combinations and model architectures.