Reward weighting method improves learning speed in multi-goal AI training
Make Sparse Rewards Count: Density-Aware Reward Aggregation for Multi-Reward RL
Machine LearningComputation and Language
Summary
Training AI models to meet several goals at once can be tricky because some goals might get learned faster than others. The authors figured out that this happens because some rewards only activate rarely in practice, so their importance is underappreciated during training. They introduced a new method, called Density-Aware Reward Aggregation (DARA), which adjusts how rewards are combined by giving more weight to those that show up less often. This helps the AI learn the different goals more quickly without changing the core training rules. Experiments show DARA speeds up learning for tasks like calling tools and solving math problems.
What this means in practice
- •For ai engineers: Accelerate training of language models to better balance multiple behavior goals using adaptive reward weighting.
- •For machine learning platform teams: Incorporate density-aware reward adjustments to improve efficiency of reinforcement learning workflows with sparse, multiple rewards.
Authors
Tong Zheng, Skylar Zhai, Zhan Cheng, TianMing Sha, Youling Huang, Shuo Zhou, Shaotong Qi, Jingcheng Liang, Xuwei Ding, Pengcheng Xu
Abstract
Multi-reward reinforcement learning trains large language models to satisfy multiple behavioral objectives simultaneously. Reward-wise normalization, as used in GDPO, preserves reward-specific relative information within rollout groups, but different objectives can still exhibit uneven learning progress. We study this behavior through advantage energy, the sum of a reward's squared advantages over a batch. Under idealized GDPO normalization, we show that this energy is proportional to active-group density: the fraction of rollout groups in which the reward provides nonzero relative advantages. This reveals a residual batch-level signal imbalance and provides a basis for calibrating reward contributions. Based on this relation, we propose Density-Aware Reward Aggregation (DARA). We derive an inverse-square-root density correction that gives greater weight to signals from less frequently active rewards. DARA computes its weights from each rollout batch, adapting to changes in reward activity throughout training without modifying the underlying policy optimization objective. Experiments on tool calling and mathematical reasoning show that DARA learns the targeted behaviors faster than GDPO, reaching high format compliance in up to 26% fewer training steps on tool calling and near-saturated length compliance in up to 65% fewer steps on mathematical reasoning, while remaining competitive in final performance. Our code is available at https://github.com/zhaihaotian/DARA.