Adaptive budget scheduling improves reinforcement learning prompt use
CERO: Where and When to Allocate Rollouts for RL Post-Training
Machine LearningArtificial Intelligence
Summary
When training reinforcement learning models that respond to groups of prompts, it's important to decide which prompts to focus on and when. The authors propose CERO, a method that carefully spreads out a limited budget of tests across prompts during training, adapting which prompts get more attention and how often. This approach helps achieve better overall results on tests involving math reasoning compared to just evenly dividing the budget. The method also comes with theoretical guarantees ensuring its decisions are close to optimal over time.
What this means in practice
- •For machine learning engineers: Improve training efficiency by dynamically allocating testing resources among input prompts to boost reinforcement learning performance.
- •For automated reasoning software teams: Enhance mathematical reasoning system accuracy by adapting prompt selection and response frequency based on performance feedback.
Authors
Yiming Zong, Yige Wang, Xing Hu, Jiashuo Jiang, Zuo-Jun Max Shen
Abstract
Adaptive rollout methods for group-relative reinforcement learning typically allocate a fixed per-update budget across prompts. We instead study how to coordinate a finite rollout budget over the entire training horizon. We formulate this problem using a concave surrogate utility of cumulative prompt exposure and introduce CERO, an online primal dual scheduler for prompt admission and budget pacing. In our experiments, each admitted prompt receives a fixed-size response group. CERO instead adapts which prompts are selected, how often they are revisited across rounds, and how many groups are generated in each round. A compact Fenchel representation linearizes the dependence on cumulative exposure, while projected online gradient descent updates prompt-specific supporting slopes and a shared budget price using reward-variation feedback and budget deviations. We establish pathwise guarantees for the surrogate allocation objective against fixed-rate and same-path time-varying benchmarks, with explicit terms for proxy discrepancy and rate variation. Under matched training-response budgets, CERO attains the highest avg@16 macro-average on each of three backbones across five mathematical reasoning benchmarks. Mechanistic analyses link CERO's prompt choices to within-group reward contrast, while multi-seed ablations show gains from adaptive pacing over both uniform and preset spending schedules.