OR Else: A Differentiable Trust Region for Policy Optimization

Technical Summary: OR ELSE: A DIFFERENTIABLE TRUST REGION FOR POLICY OPTIMIZATION

Problem Statement
Reinforcement Learning from Human Feedback (RLHF) for Large Language Models (LLMs) commonly relies on Proximal Policy Optimization (PPO) and its variant, Group Relative Policy Optimization (GRPO). Both methods utilize a clipped surrogate objective to constrain policy updates. While the clipping mechanism creates a one-sided zero-gradient region to prevent excessive policy drift in the "favorable" direction, the transition into this region is geometrically abrupt: the derivative of the objective function with respect to the probability ratio changes discontinuously at the clipping boundary. The authors investigate whether Output Reset (OR), a smooth, one-sided saturation rule originally designed for classifier training, offers a superior alternative for LLM post-training. Specifically, the paper asks if replacing the clipped surrogate with a differentiable OR-based loss can improve optimization stability and final reward scores without introducing new pathologies.

Methodology
The paper introduces PPO-OR and GRPO-OR, two variants that replace the standard clipped policy term with an OR-derived squared-margin loss in the space of token log-ratios.

• Core Mechanism: Instead of using the probability ratio $r_t(\theta)$, the methods operate on the per-token log-probability ratio $\rho_t = \log \pi_\theta(a_t|s_t) - \log \pi_{\theta_{old}}(a_t|s_t)$.
• The OR Loss: For a token with a non-zero estimated advantage $\hat{A}_t$, a target margin $\sigma_t \alpha_{OR}$ is set (where $\sigma_t = \text{sign}(\hat{A}_t)$). The loss is a squared one-sided margin: $\ell_{OR,t} = [\max(0, \alpha_{OR} - \sigma_t \rho_t)]^2$.
• Output Reset Behavior: Once $\rho_t$ crosses the margin in the direction favored by the advantage sign, the target is "reset" to the current $\rho_t$. This results in a zero residual and zero direct gradient contribution for that token. The derivative approaches zero continuously at the boundary, making the loss $C^1$ continuous (unlike the $C^0$ continuous but non-differentiable PPO clip).
• Advantage Estimation:
  • PPO-OR: Retains Generalized Advantage Estimation (GAE) with a learned value head.
  • GRPO-OR: Retains group-relative advantages (computed from $G=2$ responses per prompt) without a critic.
• Experimental Setup: The study uses Llama-3.2-1B-Instruct on the Anthropic hh-rlhf dataset. All methods share a single frozen reward model and are evaluated over three seeds. The experiments are designed as two matched within-family comparisons: PPO-clip vs. PPO-OR (under GAE) and GRPO vs. GRPO-OR (under group-relative advantages). Cross-family comparisons are noted as uncontrolled due to differences in advantage estimation, critic usage, and response lengths.

Key Contributions

1. Algorithmic Proposal: The introduction of PPO-OR and GRPO-OR, which substitute the clipped surrogate with a continuously differentiable, one-sided squared-margin loss in log-ratio space.
2. Theoretical Characterization: Formal proof that the OR objective is $C^1$ continuous at the saturation boundary and that the sample-level update direction is strictly determined by the sign of the advantage, vanishing continuously upon crossing the margin.
3. Empirical Evaluation: A rigorous comparison of clipping versus OR within two distinct advantage-estimation families, highlighting that the benefits of OR are not uniform but depend on the surrounding optimization pipeline.
4. Diagnostic Insight: The identification that replacing clipping with OR does not inherently solve the issue of large cumulative policy drift observed in group-relative settings, distinguishing local gradient saturation from global policy control.

Results
The paper reports training-time reward-model scores and optimization diagnostics across three seeds:

• GAE Family (PPO-clip vs. PPO-OR):

  • Reward Score: PPO-OR achieved a higher mean final score (+0.500) compared to PPO-clip (+0.195), a difference of +0.305.
  • Stability: PPO-OR exhibited a larger across-seed standard deviation (0.158) compared to PPO-clip (0.110).
  • Drift: Both methods maintained similar, low rollout-to-current log-ratio displacement (approx. 0.3–0.5).

• Group-Relative Family (GRPO vs. GRPO-OR):

  • Reward Score: GRPO-OR did not improve the mean score; it was slightly lower (+0.457) than GRPO (+0.488).
  • Stability: GRPO-OR showed a significantly smaller across-seed standard deviation (0.027) compared to GRPO (0.080).
  • Diagnostics: GRPO-OR achieved a near-zero terminal OR residual and a declining overshoot fraction, whereas the matched GRPO trace remained large and highly variable.
  • Drift: Crucially, GRPO-OR did not reduce the large rollout-to-current displacement observed in GRPO. Both group-relative methods exhibited substantially higher displacement (approx. 5–13) compared to the GAE methods.

Significance and Claims
The paper concludes modestly, emphasizing that the observed effects of OR are context-dependent:

• Not a Universal Trust Region: The authors explicitly state that the smooth saturation of the OR loss does not constitute a formal trust region. It suppresses the direct gradient contribution of individual tokens after they cross a margin but does not constrain cumulative policy drift caused by shared parameters, auxiliary losses, or changing advantage signs across iterations.
• Pipeline Dependency: The efficacy of OR is tied to the advantage estimation method. In the GAE setting, OR correlated with higher reward scores but increased variance. In the group-relative setting, OR reduced variance and stabilized the loss trace but did not improve the mean reward score or reduce policy drift.
• Limitations: The study acknowledges that the results are specific to a lightweight 1B model, a single frozen reward model, and a constrained group size ($G=2$). The findings characterize this specific experimental setting rather than establishing general superiority over PPO or GRPO. The paper posits that the large policy drift in group-relative methods remains an unresolved problem that OR alone cannot fix, suggesting that future work must address cumulative displacement through stronger reference regularization or explicit displacement penalties.

In summary, the paper demonstrates that Output Reset offers a differentiable alternative to clipping that alters optimization behavior and can improve reward scores in specific contexts, but it does not inherently solve the challenges of policy stability or drift in group-relative RLHF.