Technical Summary: Cross-Epoch Adaptive Rollout Optimization for RL Post-Training (CERO)

1. Problem Formulation

Large Language Model (LLM) post-training via Reinforcement Learning (RL), particularly using algorithms like GRPO, typically relies on sampling a fixed number of rollouts per prompt. This uniform allocation strategy ignores the heterogeneity in training signals provided by different prompts. Some prompts are trivial (yielding consistent success or failure with low variance), while others are too difficult or already mastered, offering diminishing marginal returns. Consequently, a significant portion of the global rollout budget is wasted on prompts with limited learning value.

The paper formulates this as a cross-epoch online resource allocation problem under a fixed global budget $B$. Unlike prior methods that perform within-batch reallocation or post-generation filtering, this work addresses the allocation of rollouts across both prompts and training epochs. The objective is to maximize the cumulative utility of collected samples, where utility is subject to prompt-level diminishing returns. The problem is mathematically defined as maximizing the sum of prompt-specific utility functions $U_i(\sum_k N_{i,k})$ subject to the global budget constraint $\sum_{i,k} N_{i,k} \le B$, where $N_{i,k}$ is the number of rollouts assigned to prompt $i$ at epoch $k$.

2. Methodology: CERO

The authors propose CERO (Cross-Epoch Adaptive Rollout Optimization), a Bayesian primal-dual online algorithm designed to adaptively allocate rollouts based on estimated prompt informativeness.

2.1 Informativeness and Utility Design

• Beta Posterior Modeling: For each prompt $x_i$, CERO maintains a Beta posterior distribution over its latent success probability $p_i$, updated based on the history of rollouts $H_i^k$.
• Informativeness Score ($q_{k,i}$): Instead of using a point estimate of success probability, CERO estimates the expected Bernoulli variance: $q_{k,i} = E[p_i(1-p_i) | H_i^k]$. This score is maximized when $p_i \approx 0.5$ (intermediate difficulty) and vanishes for easy ($p_i \approx 1$) or hard ($p_i \approx 0$) prompts, effectively identifying prompts that offer the most uncertain and thus informative learning signals.
• Utility Function: A concave, saturating utility function is defined as $U_i(n) = 1 - \exp(-\eta q_i n)$, where $n$ is the cumulative number of rollouts. This captures diminishing returns: early rollouts on an informative prompt are highly valuable, but the marginal value decreases as more samples are collected.

2.2 Optimization via Fenchel-Dual Reformulation

The temporal coupling of the utility function (which depends on cumulative allocations) makes direct online optimization difficult. CERO addresses this via a Fenchel-dual reformulation:

• The problem is transformed into a dual objective involving prompt-level dual variables $\theta_i$ and a global budget multiplier $\mu$.
• This reformulation decomposes the temporally coupled problem into per-epoch allocation subproblems coordinated by the dual variables.

2.3 Online Primal-Dual Algorithm

CERO operates as a drop-in layer on top of existing RL pipelines (e.g., GRPO):

1. Allocation Decision: At each epoch, for each prompt, the algorithm solves a linear subproblem. A prompt receives the maximum allowed rollouts ($N_{max}$) if its estimated marginal value (dual variable $\theta_i$) exceeds the current global budget price ($\mu$); otherwise, it receives zero.
2. Feedback Update: After generating responses and observing rewards, the Beta posterior for the prompt is updated, and the informativeness score $q_{k,i}$ is refreshed.
3. Dual Variable Updates:
   • Prompt-level ($\theta$): Updated via projected online gradient descent to balance the allocated rollouts against the utility implied by the current informativeness estimate.
   • Budget-level ($\mu$): Updated based on the discrepancy between the actual rollout consumption and the target budget rate. If consumption exceeds the rate, $\mu$ increases, making future allocation more selective; if below, $\mu$ decreases.

3. Key Contributions

1. Problem Formulation: The paper is the first to formulate LLM RL post-training rollout allocation as a cross-epoch fixed-budget online optimization problem with prompt-specific diminishing returns, explicitly coupling decisions across the training horizon via a shared global budget.
2. Algorithm Design: It introduces CERO, a Bayesian primal-dual algorithm that uses posterior estimates of prompt-level rollout value to guide budgeting. It utilizes the posterior expected Bernoulli variance as a proxy for learnability and employs a Fenchel-dual reformulation to handle temporal non-separability.
3. Theoretical and Empirical Validation:
   • Theory: Under fixed prompt utilities, the authors prove an $O(\sqrt{K})$ regret bound against an offline allocation benchmark, where $K$ is the number of epochs.
   • Experiments: Extensive experiments on mathematical reasoning tasks (AIME, AMC) across four open-weight LLMs (ranging from 1.5B to 7B parameters) demonstrate that CERO consistently outperforms vanilla GRPO.

4. Experimental Results

The authors evaluated CERO on the DAPO-Math-17K dataset using DeepSeek-R1-Distill-1.5B, Qwen3-4B (Base and Instruct), and Qwen2.5-Math-7B.

• Performance: CERO achieved consistent improvements over vanilla GRPO across all models and benchmarks. Average gains ranged from +1.53 to +6.04 points depending on the model.
• Efficiency Analysis: The "Effective Prompt Ratio" (the fraction of prompts in a batch containing both successful and failed responses, which generate non-zero gradients) was significantly higher for CERO than for GRPO.
  • In baseline GRPO, the effective prompt ratio dropped sharply as training progressed, indicating the model had mastered many prompts, leading to degenerate rollout groups.
  • CERO maintained a higher effective prompt ratio by adaptively shifting budget away from mastered prompts toward those with higher uncertainty.
• Overhead: The allocation update mechanism introduces negligible computational overhead (less than one second per update).

5. Significance and Claims

The paper claims that adaptive rollout budgeting is a principled and effective direction for improving the sample efficiency of RL post-training. By moving away from uniform allocation, CERO ensures that the fixed generation budget is spent on prompts that provide the most informative training signals.

The authors position CERO as a "drop-in" adaptive data-collection layer that modifies which prompts receive rollouts and how many, without altering the underlying policy-gradient objective. The work highlights that the inefficiency of current RL post-training stems not just from the policy update rule, but from the suboptimal data collection strategy.

Limitations noted by the authors:

• The theoretical regret guarantee assumes fixed prompt utilities, whereas the implementation uses dynamic, epoch-wise plug-in posterior estimates. Thus, the analysis does not fully capture the non-stationarity induced by policy updates and posterior refreshes.
• Empirical validation is currently limited to mathematical reasoning tasks; broader validation on other domains and reward models is suggested as future work.