Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yangyi Fang, Jiaye Lin, Xiaoliang Fu, Cong Qin, Haolin Shi, Chaowen Hu, Lu Pan, Ke Zeng, Xunliang Cai

TL;DR

DynaMO, a theoretically-grounded dual-pronged optimization framework that compensates for gradient attenuation of high-confidence correct actions while utilizing entropy changes as computable indicators to stabilize excessive update magnitudes, is proposed.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has proven effective for Large Language Model (LLM) reasoning, yet current methods face key challenges in resource allocation and policy optimization dynamics: (i) uniform rollout allocation ignores gradient variance heterogeneity across problems, and (ii) the softmax policy structure causes gradient attenuation for high-confidence correct actions, while excessive gradient updates may destabilize training. Therefore, we propose DynaMO, a theoretically-grounded dual-pronged optimization framework. At the sequence level, we prove that uniform allocation is suboptimal and derive variance-minimizing allocation from the first principle, establishing Bernoulli variance as a computable proxy for gradient informativeness. At the token level, we develop gradient-aware advantage modulation grounded in theoretical analysis of gradient magnitude bounds. Our framework compensates for gradient attenuation of high-confidence correct actions while utilizing entropy changes as computable indicators to stabilize excessive update magnitudes. Extensive experiments conducted on a diverse range of mathematical reasoning benchmarks demonstrate consistent improvements over strong RLVR baselines. Our implementation is available at: \href{https://anonymous.4open.science/r/dynamo-680E/README.md}{https://anonymous.4open.science/r/dynamo}.

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

TL;DR

DynaMO, a theoretically-grounded dual-pronged optimization framework that compensates for gradient attenuation of high-confidence correct actions while utilizing entropy changes as computable indicators to stabilize excessive update magnitudes, is proposed.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has proven effective for Large Language Model (LLM) reasoning, yet current methods face key challenges in resource allocation and policy optimization dynamics: (i) uniform rollout allocation ignores gradient variance heterogeneity across problems, and (ii) the softmax policy structure causes gradient attenuation for high-confidence correct actions, while excessive gradient updates may destabilize training. Therefore, we propose DynaMO, a theoretically-grounded dual-pronged optimization framework. At the sequence level, we prove that uniform allocation is suboptimal and derive variance-minimizing allocation from the first principle, establishing Bernoulli variance as a computable proxy for gradient informativeness. At the token level, we develop gradient-aware advantage modulation grounded in theoretical analysis of gradient magnitude bounds. Our framework compensates for gradient attenuation of high-confidence correct actions while utilizing entropy changes as computable indicators to stabilize excessive update magnitudes. Extensive experiments conducted on a diverse range of mathematical reasoning benchmarks demonstrate consistent improvements over strong RLVR baselines. Our implementation is available at: \href{https://anonymous.4open.science/r/dynamo-680E/README.md}{https://anonymous.4open.science/r/dynamo}.
Paper Structure (50 sections, 2 theorems, 44 equations, 6 figures, 4 tables, 1 algorithm)

This paper contains 50 sections, 2 theorems, 44 equations, 6 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Reinforcement Learning for LLMs
  4. Entropy Dynamics in Policy Optimization
  5. Sample Efficiency
  6. Preliminaries
  7. Policy Optimization in RLVR.
  8. Policy Entropy.
  9. Methodology
  10. Overview
  11. Dynamic Rollout Allocation
  12. Optimization Theory
  13. Bernoulli Variance as Practical Proxy
  14. Water-Level Implementation.
  15. Gradient-Aware Advantage Modulation
  16. ...and 35 more sections

Key Result

Lemma 1

Under the GRPO update rule, the per-sample stochastic logit update satisfies: Specifically, for the sampled action $a$ and other actions $a' \neq a$:

Figures (6)

  • Figure 1: Overview of DynaMO, which operates at both the sequence and token levels, enabling fine-grained control over optimization. (i) Left: Dynamic allocation concentrates the rollout budget on high-variance problems. (ii) Right: Gradient-aware advantage modulation compensates for attenuated gradients and stabilizes excessive updates.
  • Figure 2: Impact of DRA across different computational budgets on AIME24 with average Pass@1 (%). Solid lines denote the full DynaMO, and dashed lines are the variant that substitutes DRA with uniform allocation.
  • Figure 3: Hyperparameter sensitivity analysis on Qwen2.5-Math-7B across AIME24 and AMC23.
  • Figure 4: Performance comparison across different LLM scales (1.5B/7B/14B) with DynaMO and GRPO.
  • Figure 5: Training dynamics comparison. DynaMO maintains stable gradient norms and smooth entropy evolution, while GRPO exhibits severe spikes and fluctuations. Shaded regions show raw data range.
  • ...and 1 more figures

Theorems & Definitions (8)

  • Definition 1: Policy Concentration Measure
  • Definition 2: Entropy-Weighted Log-Probability
  • Definition 3: Centered Logit Change
  • Lemma 1: Centered Logit Update for GRPO
  • Remark 1
  • Theorem 1: Factorized Entropy Change
  • proof
  • Remark 2