Rollout-GDRO: Dynamic Compute Allocation

• Rollout-GDRO is a compute-neutral strategy that dynamically allocates policy rollouts based on prompt difficulty to efficiently reduce gradient variance.
• The method employs a game-theoretic formulation with a no-regret primal–dual controller that adapts rollout counts, yielding 9-11% accuracy gains in hard reasoning tasks.
• Empirical results demonstrate significant robustness improvements with reduced variance proxies, ensuring optimal use of a fixed mean rollout budget in LLM post-training.

Rollout-GDRO is a compute-neutral allocation strategy introduced within the Group Distributionally Robust Optimization (GDRO) framework for reinforcement learning post-training of LLMs. It dynamically reallocates the number of policy rollouts per prompt across difficulty-defined groups, maximizing gradient variance reduction for hard reasoning tasks under a strict mean rollout budget. Rollout-GDRO contrasts with standard uniform rollout strategies by adapting computation to the evolving difficulty distribution of prompts, yielding significant robustness gains on heterogeneous, heavy-tailed reasoning data (Panaganti et al., 27 Jan 2026).

1. Motivation for Dynamic Rollout Allocation

Traditional Group Relative Policy Optimization (GRPO) allocates a fixed number of rollouts $n$ per sampled prompt, maintaining uniform exploration across the dataset. In LLM reasoning, the data distribution is heavy-tailed: most prompts become “solved” (i.e., exhibit low variance) early, whereas a minority of hard prompts continue to present high estimation uncertainty. Persisting with uniform rollout allocation leads to wasted compute on simple, well-understood prompts and inadequate coverage of the hard tail, thereby hindering policy improvement precisely where it is most needed. Rollout-GDRO addresses this inefficiency by making the rollout count a function of the prompt's difficulty group, while strictly preserving the total mean rollout budget $\bar n$.

2. Game-Theoretic Formulation of Rollout-GDRO

At each training step, an online pass@k classifier partitions prompts into $B$ dynamic “difficulty bins.” Let $\hat q_t(b)$ denote the empirical proportion of prompts in bin $b$ for the current batch. The goal is to allocate, for each bin $b$, an integer number of rollouts $n_b \in \{ n_{\min}, \dots, n_{\max} \}$ so as to maximize empirical utility under a compute constraint: $$\max_{\{n_b\}} \sum_{b=1}^B \hat q_t(b)\, \hat J_b(\theta; n_b) \quad \text{s.t.} \quad \sum_{b=1}^B \hat q_t(b)\, n_b = \bar n,$$ where $\hat J_b(\theta; n_b)$ is the empirical bin-utility (negative loss) defined as

$$\hat J_b(\theta; n_b) = -\frac{1}{|\mathcal{B}_{b,t}|} \sum_{x_i \in \mathcal{B}_{b,t}} \left[ \frac{1}{n_b} \sum_{j=1}^{n_b} \ell_{i,j}(\theta) \right],$$

with $\bar n$0 the GRPO loss for rollout $\bar n$1 and prompt $\bar n$2. The corresponding Lagrangian introduces a shadow price $\bar n$3 (compute dual variable): $\bar n$4 Each bin $\bar n$5 faces a penalized bandit loss: $\bar n$6

3. Variance Proxy and the Square-Root Allocation Law

The effect of $\bar n$7 on policy improvement manifests through variance reduction in Monte Carlo gradient estimation. Under mild bounded-difference conditions, the per-prompt gradient variance for bin $\bar n$8 is bounded as: $\bar n$9 where $B$0 is the intrinsic gradient variance for bin $B$1. The overall batch-level variance proxy is: $B$2 Relaxing $B$3 to $B$4 leads to a convex optimization problem whose Karush-Kuhn-Tucker (KKT) solution yields the "square-root law": $B$5 implying that bins with larger variance $B$6 are assigned disproportionately more rollouts, scaling as $B$7 (Panaganti et al., 27 Jan 2026).

4. No-Regret Primal–Dual Controller and Algorithmic Implementation

The allocation is realized as a discrete zero-sum game between a primal player, distributing rollout arms via an EXP3P/entropic mirror-descent update, and a dual player, updating the shadow price $B$8 by projected gradient ascent on the budget violation. At each step:

• Each bin $B$9 samples rollout arm $\hat q_t(b)$0 from a mixture of current distribution $\hat q_t(b)$1 and uniform exploration.
• The exact mean budget constraint is enforced (via dynamic programming if necessary).
• Bin-utility $\hat q_t(b)$2 is estimated.
• Primal updates use bandit loss $\hat q_t(b)$3.
• The dual variable $\hat q_t(b)$4 is updated as $\hat q_t(b)$5, projected to a prescribed interval. The algorithm guarantees (Theorem B.11 in the paper) that the saddle-point gap of the averaged iterates is of order $\hat q_t(b)$6 after $\hat q_t(b)$7 training steps, ensuring no-regret and near-optimality for the original budgeted variance minimization.

5. Pseudocode Summary

A compact version of Rollout-GDRO is as follows:

$b$1

6. Empirical Evaluation and Performance

Rollout-GDRO was evaluated on the DAPO-14.1k math reasoning dataset, using pass@8 as the online difficulty classifier and Qwen3-Base models at 1.7B, 4B, and 8B parameter scales. Under a mean budget of $\hat q_t(b)$8 (with allowed arms $\hat q_t(b)$9), Rollout-GDRO achieved substantial improvements in pass@8 accuracy at no additional sampling compute:

Model Size	GRPO	Rollout-GDRO	Relative Gain
1.7B	50.74%	56.14%	+10.64%
4B	56.31%	62.27%	+10.59%
8B	62.04%	67.75%	+9.20%

7. Emergent Allocation Patterns and Qualitative Analysis

Rollout-GDRO exhibits several qualitatively distinct behaviors:

• Budget Frontier: The allocation decouples compute from data frequency; rare, hard bins can receive 3×–10× the rollouts assigned by uniform baselines.
• Variance Reduction: The weighted standard-error proxy (WSE) is reduced by 37.1%, 22.6%, and 33.4% for 1.7B, 4B, and 8B models, respectively, relative to uniform rollout allocation.
• Staircase Allocation: The shadow price $b$0 induces abrupt transitions between discrete rollout arm choices for bins (“staircase” patterns in allocation snapshots).
• Multiplier Effect: At mid-training, bins containing fewer than 20% of prompts may command over 80% of the total rollout budget.

These findings collectively support the conclusion that dynamically adapting the rollout budget using a GDRO-style shadow-price controller induces an emergent curriculum, shifting computational resources toward evolving hard-task frontiers and yielding large robustness improvements for LLM reasoning, all without increasing total sampling compute (Panaganti et al., 27 Jan 2026).