PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Rollout Allocation for Online Reinforcement Learning with Verifiable Rewards

Hieu Trung Nguyen, Bao Nguyen, Wenao Ma, Yuzhi Zhao, Ruifeng She, Viet Anh Nguyen

arXiv:2602.01601v3cs.LGcs.AIcs.CL

TL;DR

Group-based RLVR methods use fixed rollout counts despite prompts contributing different amounts of gradient variance, making sampling efficiency a central challenge. VIP predicts prompt success probabilities with a Gaussian process and optimizes rollout allocation under a fixed budget. Across mathematical and tool-augmented reasoning benchmarks, VIP improves performance over uniform and heuristic strategies with minimal runtime overhead.

  • Problem

    Fixed rollout allocation in group-based RLVR treats prompts uniformly even though their gradient variance differs, while prompt success probabilities change during training and are difficult to estimate.

  • Method

    VIP uses Gaussian-process predictions of prompt success probabilities to estimate gradient variance and solves a convex budget-constrained optimization problem for rollout allocation.

  • Results

    VIP produces consistent gains over uniform or heuristic allocation strategies across mathematical reasoning and tool-augmented benchmarks, while adding only 1.12% and 0.83% training-time overhead for 1.5B- and 7B-parameter models.

  • Takeaways & Limitations

    VIP offers a more sample-efficient and adaptive rollout-allocation approach for group-based reinforcement learning with verifiable rewards.

  • Takeaways & Limitations

    The analysis relies on covariance assumptions whose justification is partly empirical, and the experiments assume equal variances of gradient-norm variables across prompts for allocation.

Abstract

from arXiv · show

Sampling efficiency is a key bottleneck in reinforcement learning with verifiable rewards. Existing group-based policy optimization methods, such as GRPO, allocate a fixed number of rollouts for all training prompts. This uniform allocation implicitly treats all prompts as equally informative, and could lead to inefficient computational budget usage and impede training progress. We introduce VIP, a Variance-Informed Predictive allocation strategy that allocates a given rollout budget to the prompts in the incumbent batch to minimize the expected gradient variance of the policy update. At each iteration, VIP uses a lightweight Gaussian process model to predict per-prompt success probabilities based on recent rollouts. These probability predictions are translated into variance estimates, which are then fed into a convex optimization problem to determine the optimal rollout allocations under a hard compute budget constraint. Empirical results show that VIP consistently improves sampling efficiency and achieves higher performance than uniform or heuristic allocation strategies in multiple benchmarks.

1 INTRODUCTION

Group-based RL methods reduce memory use but require multiple rollouts per prompt, creating substantial computational costs. VIP addresses this by predicting prompt-level gradient variance and allocating rollouts under a fixed budget.

  • Motivation: Group-based RL methods estimate advantages from relative rollout outcomes, avoiding the separate value model required by PPO but increasing generation cost.Stable training may require around 16 generations per example.
  • Motivation: Adaptive training remains limited because existing approaches do not dynamically respond to the model’s evolving ability to solve training problems.Filtering prompts with accuracy near 0 or 1 suggests that prompt informativeness can guide control.
  • VIP: VIP predicts each prompt’s expected gradient variance using a Gaussian process over prompt success probabilities.The model uses past rollout outcomes and prompt similarity for recursive Bayesian updates.
  • VIP: VIP solves a convex optimization problem to allocate rollouts across prompts while minimizing expected gradient variance under a computational budget.An integer-rounding heuristic converts the continuous solution into a feasible allocation.
  • VIP: The paper analyzes how gradient variance relates to success probability for group-based methods including Dr. GRPO and RLOO.This analysis provides the theoretical basis for adaptive budget allocation.

2 RELATED WORK

Related work improves RL training efficiency by filtering or prioritizing prompts according to estimated informativeness. These approaches motivate adaptive control but do not provide the same variance-based allocation framework described by VIP.

  • Prompt selection: Prior work filters prompts whose model accuracy is near 0 or 1 because their rollout batches have zero variance and contribute no gradient signal.This identifies intermediate-difficulty prompts as more useful for training.
  • Prompt selection: Other methods estimate prompt informativeness before rollout generation using recent non-informative outcomes, signal-to-noise ratios, or trial rollouts.These strategies skip or filter prompts rather than directly optimizing a rollout allocation budget.

3 BACKGROUND ON RLVR

RLVR uses verifiable rewards within group-based policy optimization methods such as GRPO, Dr. GRPO, and RLOO. These methods estimate advantages from groups of sampled outputs, with the paper analyzing a setting that removes KL regularization.

  • Group-based RL: Group-based policy optimization includes GRPO and variants such as Dr. GRPO and RLOO, which use group-relative advantage estimation.The objective is defined over a dataset of prompts and their rollout groups.
  • Advantage estimators: Dr. GRPO computes each rollout’s advantage using the mean reward across all rollouts in the group.The estimator assigns the same token-independent advantage to every token in a rollout.
  • Advantage estimators: RLOO computes a rollout’s advantage using the mean reward of the other rollouts, excluding that rollout from the mean.Like Dr. GRPO, its advantage estimator is token-independent.
  • Objective assumptions: The analysis assumes zero KL regularization, so the training objective omits the KL term.The authors argue this is appropriate in RLVR because rule-based verifiers provide verifiable rewards and recent work removes KL while retaining strong performance.

4 ANALYSIS FOR GRADIENT VARIANCE

The paper analyzes the variance of per-prompt gradient estimates generated from rollout rewards and log-likelihood gradients. Under explicit independence and covariance assumptions, the resulting variance depends on rollout count and prompt success probability, motivating adaptive allocation.

  • Setup: The analyzed gradient combines off-policy and on-policy rollouts, but rollout allocation affects only the on-policy component for the current prompt batch.This isolates the portion controlled by VIP.
  • Setup: The per-prompt gradient contribution is represented through rollout rewards and the L2 norm of each rollout’s average log-likelihood gradient.The analysis studies this scalar norm as a tractable proxy for gradient-update scale.
  • Assumptions: The variance analysis assumes rollout rewards and gradient-norm variables are independently and identically distributed across samples, with specified second-order uncorrelatedness.The authors justify these assumptions using fixed-prompt sampling and empirical testing.
  • Variance results: Proposition 4.2 derives the per-prompt projected gradient variance for Dr. GRPO as a function of rollout count and binary-reward success probability.The proposition applies under Assumption 4.1 and a projected-gradient variance condition.
  • Variance results: Proposition 4.3 provides the corresponding per-prompt projected gradient variance result for RLOO under the same binary-reward and assumption framework.Together, the two propositions connect prompt success probability with allocation-relevant gradient variance.

5 PREDICTIVE ROLLOUT ALLOCATION STRATEGY

VIP predicts each prompt’s current success probability and uses those estimates to allocate a shared rollout budget unevenly, targeting lower minibatch gradient variance. Its workflow combines recursive Gaussian-process prediction with constrained optimization and rounding to obtain feasible allocations.

  • 5.1 GP Prior and Recursive Posterior Update: The model’s success probability is nonstatic because evolving weights change the output distribution, while prompt embeddings may provide limited signal for parametric classifiers.These properties motivate an adaptive, nonparametric estimator rather than a fixed predictor.
  • 5.1 GP Prior and Recursive Posterior Update: VIP predicts prompt success probabilities with a Gaussian process over prompt embeddings and uses rollout outcomes for recursive updates.At iteration t, predicted probabilities are obtained from the prior mean and then updated using observed rewards from allocated rollouts.
  • 5.2 Adaptive Budget Allocation for Gradient Variance Minimization: VIP allocates different rollout counts across minibatch prompts to minimize total gradient variance under a hard budget C.Each prompt receives an integer allocation between L and U, with the allocations summing to C.
  • 5.2 Adaptive Budget Allocation for Gradient Variance Minimization: The allocation problem is nonlinear and integer-valued, so VIP solves a continuous relaxation before applying a heuristic rounding procedure.The relaxation uses prompt-specific variance coefficients derived from predicted success probabilities, while rounding produces a feasible integer solution.
  • 5.2 Adaptive Budget Allocation for Gradient Variance Minimization: For Dr. GRPO and RLOO, the relaxed allocation has a unique solution characterized by a budget multiplier λ⋆, after which remaining budget is distributed iteratively.The iterative procedure assigns additional rollouts according to the largest decrease in the relevant variance function while respecting L and U.

6 NUMERICAL EXPERIMENTS

VIP is evaluated on mathematical and tool-augmented reasoning tasks under fixed rollout budgets, with ablations, runtime measurements, and predictor-quality tests. Results show gains over baseline allocation strategies with small computational overhead and improved adaptation to changing success probabilities.

  • Experimental setup: VIP is evaluated on mathematical reasoning and tool-augmented reasoning using fixed total rollout budgets and group-relative baselines.The experiments compare Dr. GRPO and RLOO with and without VIP across mathematical benchmarks, and evaluate tool use on Bamboogle and MuSiQue.
  • Mathematical reasoning: VIP’s relative performance gains are larger for the 1.5B and 3B models than for the 7B model.The authors suggest that budget-aware variance reduction may particularly help weaker backbones that otherwise underutilize the rollout budget.
  • Tool-augmented reasoning: On Bamboogle, Dr. GRPO+VIP raises EM from 20 to 23.2 and RLOO+VIP raises EM from 10.4 to 17.6 under a fixed rollout budget.RLOO+VIP also increases F1@5 from 0.190 to 0.264 and Precision@5 from 0.225 to 0.294; Dr. GRPO+VIP increases them by +0.051 and +0.060, respectively.
  • Ablation studies: RLOO+VIP achieves the best ablation performance, while replacing adaptive allocation causes larger drops than replacing the Gaussian process predictor.The ablation results identify variance-aware budget allocation as the primary contributor, with calibrated uncertainty providing additional improvements.
  • Runtime and prediction quality: VIP adds 1.12% and 0.83% to total RL training time for 1.5B- and 7B-parameter models, respectively, including cached computations.Excluding cached computations reduces the overhead to 0.79% and 0.58%.
  • Runtime and prediction quality: The Gaussian process predictor maintains consistently lower MAE than Moving Average and Ridge Regression for both tested model sizes.The predictor is evaluated over 55 gradient steps using recent samples while per-prompt success probabilities drift during training.

7 CONCLUSION

VIP combines Gaussian-process success-probability prediction with rollout allocation to minimize gradient variance in group-based reinforcement learning. Experiments on mathematical and tool-augmented reasoning benchmarks report consistent gains over heuristic or uniform strategies, while future work targets non-verifiable or noisy rewards.

  • Conclusion: VIP is a framework for minimizing gradient variance in group-based reinforcement learning through predictive rollout allocation.The conclusion describes the method as using Gaussian-process predictions of prompt success probabilities together with optimization-based allocation.
  • Conclusion: VIP uses limited sampling budgets more efficiently by combining success-probability prediction with rollout-allocation optimization.The conclusion characterizes this combination as the basis for more adaptive and resource-efficient training pipelines.
  • Conclusion: Experiments on mathematical reasoning and tool-augmented benchmarks show consistent gains over heuristic or uniform strategies.The conclusion states this result without specifying a single benchmark or metric.
  • Conclusion: Future work will investigate extending VIP to non-verifiable or noisy rewards, including possible applications in RLHF and other alignment paradigms.This is presented as a future direction rather than an evaluated capability.

A.2 PROOFS OF SECTION 5

The proofs establish structural properties of the rollout-allocation optimization problem and characterize its solution through KKT conditions. For feasible budgets, the objective has a unique minimizer, and the associated multiplier can be found by bisection.

  • KKT characterization: KKT conditions characterize each allocation coordinate according to whether it lies at its lower bound, upper bound, or interior.The proof separately analyzes the three cases and uses λ as the multiplier for the sum constraint.
  • Solution procedure: The aggregate allocation function S(λ) is non-increasing in the Lagrange multiplier.This monotonicity enables solving for a multiplier that satisfies the total-budget constraint.
  • Solution procedure: For any feasible budget C, a unique λ⋆ satisfying S(λ⋆) = C exists and can be found by bisection.If C exceeds BU or is below BL, the optimization problem is infeasible.
  • Optimization properties: The allocation objective is convex and decreasing in each rollout-allocation variable on the feasible set.These properties support optimization over bounded allocations.
  • Optimization properties: A unique minimizer exists whenever the feasible set is nonempty, namely when BL ≤ C ≤ BU.The same feasibility condition is stated in the theorem proof.

B.1 FIRST-ORDER CORRELATION TEST VIA FISHER’S METHOD

This section tests whether per-question correlations are collectively consistent with zero using Pearson correlations, per-question p-values, and Fisher’s method. Across 600 questions and training checkpoints, the results support the zero-correlation assumption.

  • Fisher’s method combines the per-question p-values into a single test statistic with a chi-squared reference distribution having 2Q degrees of freedom.The resulting global p-value is used for the decision rule.
  • The test rejects the global null only when pglobal < α; otherwise, it fails to reject the hypothesis that all correlations are zero.The procedure uses the specified significance level α.
  • For each of Q = 600 questions, the analysis computes a Pearson correlation from n = 16 independent rollouts and its corresponding p-value.The per-question p-values are aggregated into a global test.
  • Across Qwen2.5-Math-1.5B checkpoints, all global p-values exceed α = 0.05, supporting zero correlations across questions.The reported checkpoints correspond to 0.0, 0.5, and 1.0 epochs.

B.2 FIRST-ORDER CORRELATION TEST VIA EDGINGTON’S METHOD

This section aggregates per-question correlation and variance tests across questions using Edgington’s method and variance-homogeneity procedures. The reported tests support the assumptions of zero correlations and equal variances across questions.

  • Edgington’s method aggregates the per-question p-values by summing them and uses a normal approximation for large Q.Small values of the sum indicate joint evidence against the global null.
  • Across checkpoints, all Edgington global p-values exceed α = 0.05, supporting the assumption that correlations ρq are zero across questions.The experiment uses Q = 600 questions with n = 16 independent rollouts.
  • Equal variance tests: Levene’s test evaluates whether the variances of ˜Zq are equal across questions using transformed absolute deviations from each question’s median.The resulting statistic is compared with an F-distribution under the equal-variance null hypothesis.
  • Equal variance tests: All reported Levene global p-values exceed α = 0.05, supporting the assumption that the variances of ˜Zq are equal across questions.The test uses Q = 600 questions and n = 16 independent rollouts at Qwen2.5-Math-1.5B checkpoints.
  • Equal variance tests: O’Brien’s test provides a second variance-homogeneity check, transforming observations using each question’s sample mean and unbiased sample variance.Its statistic also follows an F-distribution under the equal-variance null hypothesis.
  • Equal variance tests: All O’Brien global p-values exceed α = 0.05, so the equal-variance null is not rejected across the evaluated checkpoints.The experiment evaluates Qwen2.5-Math-1.5B at 0.0, 0.5, and 1.0 epochs.

C.1 ADDITIONAL INFORMATION ON ABLATION STUDIES

This section describes ablation and implementation details for rollout allocation, predictive modeling, rounding, and continuous-reward extensions. VIP-related procedures use prompt-specific estimates while enforcing budget constraints.

  • Allocation heuristics: Inverse-accuracy allocation assigns more rollout budget to prompts with lower empirical accuracy, using normalized weights under global-budget and per-prompt bounds.The target weights are proportional to 1 − acci + ϵ.
  • Allocation heuristics: Inverse-variance allocation assigns more rollout budget to prompts whose answers exhibit lower variance.Both heuristic strategies use constrained optimization and integer rounding.
  • Allocation heuristics: The allocation comparison examines how VIP differs from inverse-accuracy and inverse-variance baselines as prompt difficulty changes.Figure 2 plots the optimal number of rollouts n⋆ produced by the strategies.
  • Predictive modeling: VIP’s GP predictor models prompt success probabilities from rollout outcomes and updates predictions recursively using prompt similarity.The resulting estimates are used to estimate minibatch gradient variance at each training step.
  • Continuous rewards: The continuous-reward extensions retain variance-based allocation formulations for Dr. GRPO and RLOO, with gradient variance depending on reward variance and rollout count.The rounding procedure applies without modification.
  • Continuous rewards: For continuous rewards, VIP models reward variance with a GP and softplus link, replacing the success-probability model while retaining the posterior-update procedure.The modeled quantity is vq,t = softplus(gt(xq)).

F TRAINING EVOLUTION COMPARISON

The training-evolution comparison evaluates GRPO and GRPO+VIP under matched data and computation. Across most checkpoints, VIP shows higher rewards and advantages and exceeds GRPO on performance metrics at most training steps.

  • The comparison reports mean advantages, mean rewards, and best@32, maj@32, and mean@32 for GRPO and GRPO+VIP.The evaluation uses rollout budgets of 8 and 16 per prompt.
  • All methods use the same 17k training prompts, two epochs, batch size 512, mini-batch size 64, and matched rollout budgets.Each gradient step therefore corresponds to the same amount of data and computation across methods.
  • Across most evaluation checkpoints, VIP produces consistently far-from-zero advantages and higher rewards than GRPO.The authors interpret this as a richer learning signal during training.
  • VIP performance exceeds GRPO at most training steps across the reported AIME 2024 and 2025 accuracy metrics.Figure 7 compares rollout budgets 8 and 16 across multiple accuracy measures.
Loading 2602.01601v3…