It's Not You, It's Clipping: A Soft Trust-Region via Probability Smoothing for LLM RL
Madeleine Dwyer, Adam Sobey, Adriane Chapman
TL;DR
PSPO addresses the instability and information loss caused by ratio clipping in RL for LLM fine-tuning by introducing probability smoothing that blends the current policy with the behaviour policy, creating a soft trust region. The method yields a simple, drop-in replacement for clipping with a single hyperparameter $\alpha$, while preserving gradients and implicitly shrinking KL and TV distance to the behaviour policy. Empirically, GR-PSPO outperforms GRPO-clip and sigmoid-based alternatives, with pronounced gains on math-focused and larger models, exemplified by $79.9\%$ GSM8K and $59.6\%$ MATH on Qwen2-Math-1.5B at $T=0$. The results imply that stability-aware refinement of knowledge-rich models benefits most from probability smoothing, supporting its practical adoption in post-SFT RL pipelines for LLMs.
Abstract
Training large language models (LLMs) with reinforcement learning (RL) methods such as PPO and GRPO commonly relies on ratio clipping to stabilise updates. While effective at preventing instability, clipping discards information, introduces gradient discontinuities and can prevent exploration of better policies. Inspired by label smoothing, we propose Probability Smoothing Policy Optimisation (PSPO). PSPO smooths current policy probabilities toward the behaviour policy before computing importance ratios, creating a soft trust region that preserves gradients while preventing destabilising updates. Unlike prior soft clipping approaches that use sigmoid-based transformations which can suffer from vanishing gradients and saturation, our method uses a linear interpolation, providing simpler and more robust gradient preservation. Empirically, GR-PSPO outperforms clipping and sigmoid-based alternatives on mathematical reasoning benchmarks when refining models with prior domain knowledge, achieving an accuracy of 79.9% on GSM8K and 59.6% on MATH for Qwen2-Math-1.5B.