Correlation normalized method improves multi reward learning for AI models
CorrGRPO: Correlation-Normalized GRPO for Multi-Reward Learning
Machine LearningComputation and Language
Summary
Training AI language models often involves balancing multiple rewards, but some rewards can overpower others when they happen to be on larger scales. The authors found that the usual way of normalizing these combined rewards can suppress important signals from smaller rewards. They propose a new technique called CorrGRPO that normalizes using correlations instead of raw covariances, allowing all rewards to contribute fairly. This method improves training on tasks like code writing, tool usage, and security for AI agents.
What this means in practice
- •For software developers: Enhance AI-assisted coding tools by improving training with balanced multi-objective rewards.
- •For security engineers: Train AI agents more effectively to handle security tasks combining multiple reward signals.
Authors
Wenbin Hu, Huihao Jing, Haochen Shi, Yuxuan Liu, Haoran Li, Yangqiu Song
Abstract
Group Relative Policy Optimization (GRPO) is widely used to train reasoning language models, where it computes advantages by centering and normalizing rewards across rollouts of the same prompt. For multiple rewards, GRPO sums the reward components and normalizes the total reward by its within-group standard deviation. The corresponding variance equals the sum of all pairwise reward covariances. For a fixed centered reward, larger aggregate covariance produces smaller advantages, and vice versa, allowing update magnitudes to adapt to reward dependence. However, correlated rewards with large scales can dominate this normalization and suppress signals from smaller-scale rewards. We propose Correlation-Normalized GRPO (CorrGRPO), which normalizes pairwise covariances into Pearson correlation coefficients. CorrGRPO keeps the centered total reward unchanged while balancing the influence of differently scaled rewards on the correlation-based normalization. This allows advantage magnitudes to adapt to reward correlations without the normalization being dominated by large-scale reward components. We compare CorrGRPO with GRPO and other variants on code generation, tool calling, and agent security, using models ranging from 0.5B to 8B parameters. These tasks all involve multiple rewards that can improve together or present tradeoffs. Results show improvements across three domains, including code generation, tool calling, and agent security. Our code is available at https://github.com/HKUST-KnowComp/CorrGRPO.