SoftmaxGRPO: Learning to Reason using Softmax Advantage Group Estimation

2026-08-10

AI Summary

This work addresses the inefficiency of conventional group-based reinforcement learning methods, such as GRPO, which assign overly strong learning signals to easy prompts under binary rewards, thereby hindering effective training. To overcome this limitation, the authors propose SoftmaxGRPO, which introduces a temperature-scaled softmax advantage function into group policy optimization, replacing the standard z-score normalized advantage. This modification enables difficulty-adaptive weighting while ensuring bounded advantages. Theoretically, the method is shown to be equivalent to optimizing a log-moment-generating-function objective, with MaxRL emerging as its low-temperature limit. Empirically, SoftmaxGRPO achieves 51.8% accuracy on the DeepMath benchmark and, using only a lightweight text similarity reward, substantially improves performance on the Poetry task from 35.0% to 68.0% with a 1.5B-parameter model.

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.

Problem

• group-based reinforcement learning
• prompt difficulty
• binary rewards
• learning signal allocation
• advantage weighting

Innovation

• SoftmaxGRPO
• group-based reinforcement learning
• softmax advantage
• temperature scaling
• reward redistribution