How to Allocate, How to Learn? Dynamic Rollout Allocation and Advantage Modulation for Policy Optimization

1. Problem Statement

The paper addresses two fundamental challenges in Reinforcement Learning with Verifiable Rewards (RLVR) for Large Language Models (LLMs), particularly in mathematical reasoning tasks:

1. Inefficient Resource Allocation (Sequence Level): Current methods (e.g., GRPO) uniformly distribute the rollout budget (number of generated responses) across all training problems. This ignores gradient variance heterogeneity. Problems with high variance (balanced success/failure rates) offer the most informative learning signals but also introduce noise. Uniform allocation wastes resources on trivial problems (low variance) or overly difficult ones (high noise), failing to optimize the trade-off between informativeness and estimation noise.
2. Unstable Policy Optimization Dynamics (Token Level): The standard softmax policy structure creates a tension in gradient magnitudes:
   • Gradient Attenuation: High-confidence correct actions naturally yield small gradient magnitudes, leading to insufficient learning signals for actions the model is already good at.
   • Training Instability: Conversely, excessive gradient updates can destabilize training. Existing stabilization methods (e.g., ratio clipping) are coarse-grained and lack a unified theoretical grounding, often failing to address the specific interplay between gradient magnitude and policy entropy.

2. Methodology: DynaMO Framework

The authors propose DynaMO, a dual-pronged optimization framework operating at both the sequence and token levels, grounded in variance minimization theory and gradient-entropy analysis.

A. Sequence Level: Dynamic Rollout Allocation (DRA)

• Theoretical Basis: The authors prove that uniform allocation is suboptimal. They derive an optimal allocation strategy by minimizing the total gradient estimation variance under a fixed budget constraint.
• Optimal Formula: The optimal number of rollouts $n_i^*$ for problem $i$ is proportional to its gradient standard deviation $\sigma_i$:
  $$n_i^* \propto \sigma_i$$
• Practical Proxy: Since $\sigma_i$ is unknown, they derive Bernoulli Variance as a lightweight, computable proxy. For binary rewards, variance is $p(1-p)$, where $p$ is the success probability. They estimate this using historical statistics (count of correct responses $k_i$ out of total rollouts $G_i$):
  $$P_i = \frac{k_i(G_i - k_i)}{G_i(G_i - 1)}$$
• Implementation: A "Water-Level" algorithm (Algorithm 1) dynamically allocates rollouts. It concentrates the budget on problems where $P_i$ is high (balanced success/failure, indicating the model is in the "capability gap") while respecting minimum ($G_{min}$) and maximum ($G_{max}$) bounds to ensure coverage and prevent over-concentration.

B. Token Level: Gradient-Aware Advantage Modulation

This component addresses the gradient dynamics within the softmax policy using a unified compensation and stabilization mechanism.

1. Gradient Compensation (for Attenuation):

   • Theory: Theoretical analysis shows that the expected update magnitude is upper-bounded by policy entropy. High-confidence actions (low entropy) suffer from gradient attenuation.
   • Mechanism: A compensation factor $\beta_{comp}$ is applied to positive-advantage tokens. It scales inversely with entropy:
     $$\beta_{comp} = 1 + \alpha \cdot \frac{H_{max} - H}{H_{max} - H_{min}}$$
     This amplifies learning signals for high-confidence correct actions without affecting negative-advantage tokens.

2. Update Magnitude Stabilization (for Instability):

   • Theory: Large entropy changes indicate excessive gradient magnitudes that could destabilize training. The authors derive a factorized form showing that entropy change is driven by policy concentration and the composite update coefficient.
   • Mechanism: An instability indicator $\Xi_{i,t}$ measures the magnitude of entropy change. A stabilization factor $\beta_{stab}$ (a sigmoid-based decay function) reduces the advantage for tokens exhibiting excessive entropy changes, effectively "braking" the update.

3. Integrated Objective:
   The final advantage $A_{final}$ combines both mechanisms:
   $$A_{final} = A_{original} \cdot \beta_{comp} \cdot \beta_{stab}$$
   This is integrated into the standard clipped surrogate objective (GRPO/PPO).

3. Key Contributions

1. Theoretical Proof of Suboptimality: Proved that uniform rollout allocation is suboptimal and derived a variance-minimizing allocation strategy using Bernoulli variance as a computable proxy.
2. Gradient-Entropy Relationship: Established a theoretical link between gradient magnitude bounds and policy entropy, enabling a unified mechanism for compensation (fixing attenuation) and stabilization (preventing explosion).
3. Dual-Level Optimization: Proposed a framework that simultaneously optimizes resource distribution across problems and fine-grained token-level update dynamics.
4. Empirical Validation: Demonstrated consistent improvements across multiple benchmarks and model scales, with ablation studies confirming the necessity of each component.

4. Experimental Results

• Benchmarks: Evaluated on six mathematical reasoning benchmarks: AIME24, AIME25, AMC23, MATH500, Minerva, and Olympiad.
• Models: Tested on Qwen2.5-Math (1.5B, 7B) and Qwen3-14B.
• Performance:
  • DynaMO consistently outperformed strong baselines (GRPO, Clip-Higher, Entropy Loss, W-REINFORCE, etc.).
  • Example (Qwen2.5-Math-7B): On AIME24, DynaMO achieved 34.4% Pass@1, compared to 28.8% for standard GRPO and 32.8% for the next best baseline (KL-COV).
  • Scaling: The performance gap widened as model scale increased (1.5B $\to$ 7B $\to$ 14B), suggesting the method becomes more effective with larger representational capacities.
• Training Dynamics: Visualizations showed that DynaMO maintains stable gradient norms and smooth entropy evolution, whereas baselines (like GRPO) exhibited severe spikes and erratic fluctuations.
• Efficiency: The additional computational overhead (variance estimation, modulation) was negligible, maintaining comparable per-step training times to baselines.

5. Significance

• Resource Efficiency: DynaMO offers a principled way to allocate computational budgets, maximizing learning signal quality per unit of compute by focusing on "hard but solvable" problems.
• Stability & Convergence: By theoretically grounding advantage modulation in gradient-entropy relationships, it solves the trade-off between learning from confident actions and preventing training collapse, leading to more robust RLVR training.
• Generalizability: The method operates purely at the algorithmic level (modifying allocation and advantage computation) without relying on specific architecture features, making it broadly applicable to various LLM families and reasoning tasks.

In conclusion, DynaMO represents a significant step forward in making RLVR for LLMs more efficient, stable, and theoretically sound, addressing critical bottlenecks in both data utilization and optimization dynamics.