SoftmaxGRPO: Learning to Reason using Softmax Advantage Group Estimation
2026-08-10 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors study a group-based reinforcement learning method called GRPO and find it can struggle to fairly distribute learning effort when prompts have different difficulty levels. They propose a new method, SoftmaxGRPO, that uses a softmax function to better balance learning across prompts, keeping weights from growing too large. The authors show mathematically how their approach works for both binary and scalar rewards, and demonstrate it improves training on tasks like DeepMath and Poetry compared to the original method. Their work offers a more stable and effective way to allocate learning signals in group-based reinforcement learning.
reinforcement learningGRPOgroup normalizationsoftmax functionbinary rewardsscalar rewardslog-moment-generating functiongradient allocationinstruction tuningDeepMath benchmark
Authors
Jefferson Hernandez, Jaywon Koo, Zilin Xiao, Chen Wei, Vicente Ordonez
Abstract
Group-based reinforcement learning objectives such as GRPO can allocate learning signal poorly across prompt difficulty: under binary rewards, group normalization induces a divergent weighting on easy prompts. We introduce Softmax Advantage Group Estimation (SoftmaxGRPO), a drop-in alternative that replaces z-score-normalized group advantages with temperature-scaled softmax advantages, keeping weights bounded regardless of prompt difficulty. For binary rewards, we derive the exact finite-group population objective and identify MaxRL as its low-temperature limit. For bounded scalar rewards, we show that the large-group update exactly optimizes a log-moment-generating-function objective, while a universal finite-group scalar objective cannot exist without additional assumptions on the reward distribution. Empirically, SoftmaxGRPO reallocates measured gradient budget away from near-solved prompts and consistently improves over GRPO under identical rewards. It reaches 51.8% on DeepMath with verifiable rewards and improves a 1.5B instruction-tuned model from 35.0% to 68.0% on Poetry using only lightweight text-similarity rewards.