Group policy optimization improves by calibrating reward and KL updates
When KL Regularization Misfires in Group Policy Optimization
Machine LearningComputation and Language
Summary
When training AI systems that learn by trial and error, sometimes a common technique called KL regularization can actually slow down learning. The authors find several reasons why this happens, such as leftover effects from previous steps and imbalanced influences. They suggest a new method called Zero-Sum Calibrated Policy Optimization that adjusts how rewards and KL regularization work together. Their experiments show this new method helps improve learning in these scenarios.
What this means in practice
- •For reinforcement learning engineers: Improve group-based policy training by balancing reward and KL influences for more stable policy updates.
- •For natural language processing developers: Enhance fine-tuning of language models that use grouped token rewards by calibrating reward coefficients informed by KL divergence.
Authors
Fei Ding
Abstract
Why does removing reference-policy KL regularization sometimes improve group policy optimization? This motivates studying how reference-policy information should enter group-relative updates. We analyze seven potential failure modes in the interactions between KL and rewards: residual KL updates after reward clipping, after gradient cancellation, and in groups with identical rewards; KL growth with response length and an imbalance in its relative contribution; KL concentration on a small number of tokens; and sampling noise when k1 is incorporated into rewards. We propose Zero-Sum Calibrated Policy Optimization (ZCPO), which uses relative drift measured by conditional KL to calibrate within-group reward coefficients and integrates them into the base surrogate. Mathematical reasoning experiments and ablations support this design's effectiveness in our settings.