OR Else: A Differentiable Trust Region for Policy Optimization

2026-07-20

AI Summary

This work addresses the issue of abrupt gradient shifts in PPO and GRPO caused by clipping mechanisms when policy updates saturate in favorable directions. To mitigate this, the authors propose an Output Reset (OR) mechanism that replaces hard clipping with a smooth, one-sided saturation function and introduces a sign-aware squared-margin loss in the token log-ratio space, yielding a differentiable trust-region optimization method. This is the first application of such smooth saturation techniques to large language model post-training. The approach integrates GAE, group-relative advantages, the Llama-3.2-1B-Instruct backbone, and a shared reward model for policy optimization. Experiments show that PPO-OR achieves a 0.305 higher average final reward under GAE compared to PPO-clip, while GRPO-OR, though not improving mean performance, substantially reduces inter-seed variance, drives terminal residuals toward zero, and lowers overshoot ratios.

Abstract

PPO and the GRPO baseline studied here use clipped surrogate objectives whose favorable-direction saturation introduces an abrupt change in the scalar objective's derivative. We ask whether Output Reset (OR), a smooth one-sided saturation rule, offers a useful alternative for large language model post-training. PPO-OR and GRPO-OR replace the clipped policy term with an OR squared-margin loss in rollout-relative token log-ratio space; the advantage sign determines the update direction, and a token contributes zero direct OR residual after crossing the favorable margin. We compare PPO-clip with PPO-OR under generalized advantage estimation (GAE), and GRPO with GRPO-OR under group-relative advantages, using \texttt{Llama-3.2-1B-Instruct} on Anthropic \texttt{hh-rlhf} with one shared reward model and three seeds per method. Under GAE, PPO-OR has a mean final training-time reward-model score $0.305$ higher than PPO-clip, with a larger observed across-seed spread. Under group-relative advantages, GRPO-OR does not have a higher mean score, but shows a smaller observed spread, a near-zero terminal OR residual, and a declining overshoot fraction, while the matched GRPO clipped-objective trace remains variable. Both group-relative methods exhibit substantially larger rollout-to-current log-ratio displacement than the GAE methods, and OR does not consistently reduce it. Thus, OR changes optimization behavior in both matched comparisons, but the observed reward effect differs between them. At $G=2$, the GRPO-OR diagnostics do not translate into a reward-score gain. Whether larger groups change this outcome remains open. The reported scores are training-time reward-model measurements, not held-out human-preference performance.

Problem

• policy optimization
• clipped surrogate objective
• saturation
• large language model post-training
• trust region

Innovation

• Output Reset
• differentiable trust region
• policy optimization
• smooth saturation
• squared-margin loss