Balanced Aggregation: Understanding and Fixing Aggregation Bias in GRPO
| Authors | Zhiyuan Zeng et al. |
| Year | 2026 |
| HF Upvotes | 4 |
| arXiv | 2605.04077 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
Reinforcement learning with verifiable rewards (RLVR) has become a central paradigm for improving reasoning and code generation in large language models, and GRPO-style training is widely adopted for its simplicity and effectiveness. However, an important design choice remains underexplored: how token-level policy gradient terms are aggregated within each sampled group. Standard GRPO uses sequence aggregation, while recent work has advocated token aggregation as a better alternative. We show that these two rules induce different optimization biases: token aggregation introduces sign-length coupling, while sequence aggregation implicitly downweights longer responses through sequence-level equal weighting. To address this tension, we propose Balanced Aggregation (BA), a simple drop-in replacement that computes token-level means separately within the positive and negative subsets and then combines them with sequence-count-based weights. Experiments with Qwen2.5-Math-7B and Qwen3-1.7B on DAPO-17k and Polaris, evaluated on six reasoning and coding benchmarks, show that BA consistently improves training stability and final performance over standard token and sequence aggregation. Our analysis further shows that the relative effectiveness of token and sequence aggregation is largely governed by response-length variation and the positive-negative length gap, highlighting aggregation as a critical design dimension in GRPO-style RLVR.
Engineering Breakdown
Plain English
This paper addresses a critical but overlooked design choice in GRPO (Group Relative Policy Optimization) training: how to aggregate policy gradients across tokens within a sampled sequence. The authors demonstrate that the standard sequence-level aggregation approach downweights longer responses unfairly, while the alternative token-level aggregation introduces problematic sign-length coupling. They propose Balanced Aggregation (BA), a simple replacement that fixes both biases by properly normalizing gradient contributions regardless of sequence length, enabling fairer and more stable training for reasoning and code generation tasks in large language models.
Core Technical Contribution
The core innovation is identifying and formally characterizing two distinct aggregation biases in GRPO-style reinforcement learning: sequence-level equal weighting implicitly penalizes longer correct solutions, while token-level aggregation couples gradient signs with response lengths in ways that distort optimization. The authors propose Balanced Aggregation as a principled solution that decouples these effects through proper normalization of per-token gradients, ensuring that optimization pressure is distributed fairly across sequences of different lengths. This is a simple but theoretically motivated fix that requires minimal code changes while addressing a fundamental flaw in how current GRPO implementations handle multi-token trajectories.
How It Works
Standard GRPO computes policy gradients at each token position within a sampled sequence, then aggregates these per-token gradients using one of two methods: (1) sequence aggregation sums or averages gradients across all tokens in a single sequence before the parameter update, or (2) token aggregation treats each token as an independent sample and updates on the aggregate across all tokens from all sequences. The paper shows sequence aggregation inadvertently creates an implicit 1/length weighting because each sequence gets equal weight regardless of its token count, penalizing longer (potentially correct) responses. Balanced Aggregation fixes this by normalizing each token's gradient contribution by the sequence length, then aggregating across all tokens from all sequences—this preserves the signal strength of longer sequences while avoiding the sign-length coupling problem of naive token aggregation. The result is a simple drop-in replacement that reweights gradients more fairly without requiring architectural changes or additional computational overhead.
Production Impact
For teams training reasoning models with GRPO, this fix directly addresses a silent but significant performance drain: current implementations systematically bias the model away from longer correct reasoning traces, even when those traces produce higher rewards. In production, this means adopting Balanced Aggregation could improve solution quality on benchmarks like AIME and CodeForces by removing the artificial penalty for exploration of longer solution paths, without requiring new infrastructure or data. The implementation cost is negligible—just a change to the gradient aggregation logic in your training loop—making this a high-ROI optimization for any existing GRPO pipeline. The trade-off is minimal: no computational overhead, no hyperparameter tuning required, and full backward compatibility with existing checkpoints and learning rate schedules. For large-scale training runs (like those spanning weeks of compute), this fix compounds to measurable improvements in final model performance without increasing training time.
Limitations and When Not to Use This
The paper assumes that longer sequences are not inherently lower quality, which may not always hold in practice—some tasks genuinely produce shorter optimal solutions, and the fair weighting could overvalue unnecessary length expansion. The work focuses specifically on GRPO-style training and may not generalize directly to other RL approaches like PPO or DPO variants, each of which may have their own aggregation-related biases not addressed here. The paper does not deeply explore interaction effects with other hyperparameters (batch size, learning rate scheduling, reward scaling) that might also compensate for or exacerbate the original bias, meaning practitioners may still need tuning. Finally, the evaluation appears limited to code and reasoning tasks; the approach's behavior on open-ended generation, dialogue, or other domains where response length distribution is different remains unclear and would benefit from follow-up work.
Research Context
This work builds directly on the GRPO framework (recent work simplifying group-relative policy optimization for LLMs) and contributes to the broader shift toward RL with verifiable rewards for reasoning—following successes like OpenAI's o1 approach and related work on verifiable reasoning. It fits into a research direction examining hidden biases and implicit inductive biases in RL for language models, complementary to work on reward modeling, outcome supervision, and process supervision. The paper advances practical understanding of how gradient aggregation choices in sequence-level RL affect optimization trajectories, opening the door to examining similar biases in other aggregation-based training schemes. By showing that a single line of code (normalization) can remove a systematic bias, it demonstrates the value of careful algorithmic analysis in mature frameworks and encourages similar scrutiny of other standard practices in LLM RL.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::