Adaptive rollout allocation boosts large language model training efficiency
Bayes-Adaptive MDPs for Rollout Allocation to Enhance GRPO Training
Machine LearningArtificial Intelligence
Summary
Training large language models with uniform effort on all questions can waste resources on easier questions and not spend enough on harder ones. The authors address this by creating a new approach that adjusts how much training effort each question gets based on difficulty. They use a mathematical framework called a Bayes-adaptive Markov decision process to decide how to allocate resources before training begins. Tests show this method improves training results without needing extra work during training.
What this means in practice
- •For llm training engineers: Increase training efficiency by allocating training rollouts based on question difficulty without adding computation during training.
- •For machine learning infrastructure teams: Optimize resource use when training large language models by applying offline policies that decide rollout distribution in advance.
Authors
Faeze Moradi Kalarde, Armin Toroghi, Nam Ho Koh, Doo Hwan Hwang, Seungjin Kim, Kee-Eung Kim, Scott Sanner
Abstract
Group Relative Policy Optimization (GRPO) has demonstrated strong performance in LLM post-training. However, it allocates a uniform rollout budget to every training question, ignoring variations in question difficulty. As a result, highly difficult questions may receive too few rollouts to generate informative positive reward signals, limiting their contribution to training, while moderately difficult questions often receive more rollouts than necessary. In this work, we formulate the problem of difficulty-aware rollout budget allocation to maximize training utility under the same total rollout budget as standard GRPO. We propose AdaRoll-GRPO, which reformulates rollout allocation as a sequential decision-making problem and leverages a Bayes-adaptive Markov decision process (MDP) to learn an adaptive allocation policy. The learned policy is computed offline and applied during GRPO training without introducing additional online optimization overhead. Extensive experiments on multiple reasoning datasets and LLM architectures demonstrate that AdaRoll-GRPO consistently improves GRPO training performance under a fixed rollout budget, highlighting the effectiveness of adaptive rollout allocation for GRPO in LLM post-training.