Not All Rollouts are Useful: Down-Sampling Rollouts in LLM Reinforcement Learning

Yixuan Even Xu; Yash Savani; Fei Fang; J. Zico Kolter

Source-linked AI summary

Not All Rollouts are Useful: Down-Sampling Rollouts in LLM Reinforcement Learning

Yixuan Even Xu, Yash Savani, Fei Fang, J. Zico Kolter

arXiv:2504.13818v5Apr 18, 2025cs.LGcs.AIcs.CL

TL;DR

RLVR rollout generation is highly parallel and memory-light, while policy updates are memory- and communication-intensive. PODS generates many rollouts but trains on a max-variance-selected subset with an O(n log n) selection algorithm. Across tested benchmarks and configurations, GRPO-PODS reaches vanilla GRPO’s peak test accuracy at least 1.7× faster, but its evaluation is limited to verifiable-reward tasks and may not suit strict on-policy settings.

Problem
RLVR must reconcile scalable rollout generation with memory- and communication-intensive policy updates that create a computational bottleneck.
Method
PODS generates n rollouts and updates the policy using m < n selected examples, with max-variance down-sampling preserving reward contrast and computable in O(n log n) time.
Results
GRPO-PODS reaches baseline GRPO’s peak test accuracy at least 1.7× faster across math and chemistry tasks, model scales, architectures, and deployment scenarios.
Takeaways & Limitations
Selective rollout training reduces policy-update costs while retaining the reported learning quality and broadening hardware utilization in GRPO-based RLVR.
Takeaways & Limitations
The evaluation focuses on verifiable-reward RLVR tasks, and selective down-sampling changes the rollout distribution so PODS may be unsuitable under strict on-policy guarantees.

Abstract

from arXiv · show

Reinforcement learning with verifiable rewards (RLVR) has emerged as the leading approach for enhancing reasoning capabilities in large language models. However, it faces a fundamental compute and memory asymmetry: rollout generation is embarrassingly parallel and memory-light, whereas policy updates are communication-heavy and memory-intensive. To address this, we introduce PODS (Policy Optimization with Down-Sampling), which decouples rollout generation from policy updates by training only on a strategically selected subset of rollouts, maintaining learning quality while dramatically reducing update costs. We propose a principled subset selection criterion, max-variance down-sampling, that maximizes reward diversity, and provide an efficient $O(n\log n)$ implementation. Empirically, Group Relative Policy Optimization (GRPO) with PODS achieves the peak test accuracy of vanilla GRPO at least $\mathbf{1.7\times}$ faster across the different reasoning benchmarks and hardware configurations we tested.

1 Introduction

RLVR’s inference and policy-update phases have sharply different hardware demands, motivating PODS: generate many rollouts efficiently, then update on a principled informative subset. Max-variance down-sampling provides an efficient selection rule, and GRPO-PODS reaches baseline peak accuracy at least 1.7× faster.

Motivation: Inference is embarrassingly parallel and memory-light, whereas policy updates are memory- and communication-intensive.Batching improves rollout throughput, while policy updates require optimizer states and cross-device gradient and parameter synchronization.
Motivation: 21× lower per-token inference time through batching contrasts with policy updates becoming out of memory beyond 32 rollouts per GPU.The timing evidence uses Qwen2.5-3B-Instruct on GSM8K with 8 A100-80GB GPUs.
PODS: PODS generates n rollouts but updates on only m < n informative samples, avoiding redundant information and update latency.The framework exploits scalable inference while selectively reducing the policy-update workload.
Max-variance down-sampling: Max-variance down-sampling selects the subset with greatest reward variance, preserving contrastive signals and solvable in O(n log n) time.For binary rewards, the rule chooses the m/2 highest-reward and m/2 lowest-reward rollouts.
Results: 1.7× faster peak baseline accuracy is achieved by GRPO with PODS across GSM8K, MATH, and Chemistry evaluations.The comparison spans multiple model and hardware configurations.

2 Related Work

Prior work spans reasoning-oriented RL, prompt selection, and broader data-selection methods. This paper instead down-samples rollouts within each prompt, addressing RLVR’s update bottleneck and complementing prompt-level methods.

Reinforcement learning for LLM reasoning: GRPO is prominent for reasoning tasks because of its simple implementation, competitive PPO performance, and lack of a separate critic network.The related work also situates GRPO among large-scale reasoning-focused RL efforts.
Prompt and rollout selection: Recent RL data-selection work emphasizes selecting or filtering prompts, while this paper selects rollouts within each prompt.The distinction defines the paper’s position relative to prompt-selection methods such as DAPO.
Down-sampling and data selection: Selective sampling has established precedents in clustering, regression, experience replay, foundation-model training, and computational advertising.These examples motivate data management beyond LLM reinforcement learning.
Down-sampling and data selection: The paper presents principled down-sampling of LLM-RL rollouts as a way to mitigate a computational bottleneck while strengthening the learning signal.Its approach can be combined with existing prompt-selection and filtering methods.

3 Down-Sampling Rollouts in GRPO

PODS decouples GRPO rollout generation from policy updates by selecting a smaller, informative subset of rollouts for training. Its max-variance rule preserves contrastive reward signals while enabling efficient down-sampling.

3.1 Preliminaries: GRPO computes rollout rewards and normalized advantages before optimizing a token- and rollout-averaged objective, while PODS recomputes advantages on the selected subset.The frozen policy supplies importance-weighting probabilities, and the asymmetric loss limits increases more than reductions.
3.2 PODS Framework: PODS generates n rollouts, retains m < n using a down-sampling rule, and updates the policy only on the retained subset.The rule outputs m indices for the policy update phase after rewards are computed.
3.2 PODS Framework: Random down-sampling uniformly selects m rollouts and matches standard GRPO with m rollouts in expectation.Percentile sampling instead spreads selected rollouts across the reward distribution, while max-reward sampling focuses on highest-reward examples but loses negative feedback.
3.3 Max-Variance Down-Sampling: Max-variance down-sampling selects the size-m subset that maximizes empirical reward variance, spanning successful and unsuccessful reasoning paths.The resulting contrastive reward distribution is intended to preserve stronger learning signals.
3.3 Max-Variance Down-Sampling: O(n log n) time suffices to find an optimal max-variance subset, replacing an infeasible naive search over subsets.The practical algorithm sorts rewards and evaluates the structured candidate subsets.
3.3 Max-Variance Down-Sampling: For sorted rewards, an optimal subset consists of k highest and m − k lowest rewards for some k.This structure follows from repeatedly replacing selected points with more distant rewards without decreasing variance.

4 Experiments

Experiments evaluate GRPO-PODS across diverse tasks, models, hardware, and training settings, showing faster convergence and practical trade-offs for rollout and update sizes. Max-variance down-sampling remains effective across aggressive reduction ratios, while larger rollout pools eventually yield diminishing returns.

Experimental setup: PODS is evaluated across GSM8K, MATH, and L3 Chemistry using Qwen2.5 and Llama3.2 models from 3B to 7B parameters.The experiments include both single-GPU resource-constrained setups and multi-GPU distributed training.
Experimental setup: The study compares PODS with baseline GRPO under memory-constrained single-GPU settings and with gradient-accumulated GRPO in distributed settings.Single-GPU comparisons use the same training batch size m, while distributed comparisons use GRPO-GA.
GRPO-PODS performance: 1.7× faster convergence: PODS reaches baseline GRPO peak accuracy at least 1.7× faster across the tested configurations and often achieves higher final performance.The evaluation spans Math and Chemistry, 3B–7B models, Qwen2.5 and Llama3.2 architectures, and multiple deployment scenarios.
Rollout and update sizes: n = 64 is an approximate optimal rollout-pool size, while performance degrades beyond n = 128 as inference costs rise and rollout-diversity gains plateau.Increasing n initially improves selection quality, but larger pools eventually add computational overhead without comparable diversity benefits.
Rollout and update sizes: PODS remains robust until very small update batches, supporting down-sampling ratios up to 16 at n = 64 and m = 4.Ratios of 2 to 4 are suggested as an effective general balance between performance and efficiency.
Down-sampling rules: Max-variance down-sampling consistently outperforms max-reward, random, and percentile rules throughout training.This comparison is conducted on GSM8K with one L40S.

5 Conclusion and Discussion

PODS addresses the asymmetry between parallel, memory-light rollout generation and memory-intensive policy updates by training on an informative down-sampled subset. Within GRPO, it matches baseline accuracy faster, while its applicability remains bounded by task, policy, and sampling-setting assumptions.

PODS generates large rollout batches in parallel and updates the policy using only an informative subset selected by max variance.The optimal subset can be found in O(n log n) time.
1.7× speedup is achieved while reaching higher final accuracy across diverse model architectures, scales, and deployment scenarios.The comparison is made against standard GRPO under equal wall-clock budgets.
Performance remains robust across a wide range of down-sampling ratios when m is not too small.
The evaluation focuses on RLVR tasks with verifiable response correctness, while open-ended dialogue may exhibit different dynamics.
PODS behaves off-policy because selective down-sampling changes the training rollout distribution, making it potentially unsuitable when strict on-policy guarantees are required.
The empirical and theoretical claims are specific to PODS within GRPO, and comparable benefits for other RL methods require separate analysis and experimentation.

A.1 Reward Functions

The experiments use rule-based reward components for correctness, format compliance, and thinking-tag placement. These components provide binary or partial-credit signals for evaluating mathematical, chemistry, and formatting behavior.

Correctness rewards assign 1 for correct and 0 for incorrect answers, using symbolic verification or direct answer matching depending on the task.
Format rewards assign 1 for compliant and 0 for non-compliant responses that follow the specified XML structure.
Tag-count rewards provide partial credit from 0 to 1, assigning 0.25 points for each correctly placed thinking or answer tag.

A.2 Hyperparameters

The appendix lists hyperparameters across experimental settings and describes how gradient accumulation comparisons equalize the effective number of generated rollouts. PODS then reduces each rollout batch before updating.

Table 2 lists the key hyperparameters used for different experimental settings.
For settings (e) and (f), GRPO-PODS and GRPO-GA match the total effective rollout count per prompt by equating rollout batch size multiplied by gradient-accumulation steps.
With rollout batch size 128 and 4 gradient-accumulation steps, both methods use an effective n of 512.
GRPO-PODS down-samples each rollout batch by a factor of 4 while preserving equal total rollout processing between variants.

A.3 On Advantage Normalization

GRPO-PODS computes advantage-normalization statistics after down-sampling, using the same batch that updates the policy. Figure 6 compares this design with normalization computed before down-sampling.

GRPO-PODS first down-samples the rollout batch, then computes advantage mean and standard deviation from the down-sampled batch.
This ordering makes normalization consistent with the policy-update data and gives each update batch a total advantage of 0.
Figure 6 compares “After” normalization on the down-sampled batch with “Before” normalization on the full rollout batch.Training time is on the x-axis, test accuracy is on the y-axis, and shading denotes 1.96 times the standard error of the mean.

A.4 PODS’ Speed Up Ratio Over GRPO

GRPO-PODS reaches comparable peak performance faster than standard GRPO across the evaluated settings. Its measured speedup ranges from 1.7× to 3.0×.

1.7×–3.0× speedup is achieved by GRPO-PODS over GRPO across the evaluated settings.The speedup ratio measures the time for each method to reach 0.99× GRPO’s peak performance.
GRPO-PODS consistently outperforms GRPO as training proceeds.
Table 3 reports the speedup ratios for GRPO-PODS over GRPO in Figure 3.

A.5 Additional Evaluation for Experimental Settings (a) and (b)

GRPO-PODS maintains its performance advantage over GRPO on additional test sets in experimental settings (a) and (b). The accompanying figures also report test accuracy and completion length over training time.

Additional evaluation: GRPO-PODS consistently outperforms GRPO on GSM8K Platinum and MATH across experimental settings (a) and (b).GSM8K Platinum is described as a contamination-resistant subset of GSM8K.
Additional evaluation: The additional-test-set results indicate that the improvement is not due to overfitting to a specific test set.
Additional evaluation: Figure 7 plots test-set accuracy against training time for standard GRPO and GRPO-PODS with max-variance down-sampling.
Completion length: Average completion length stays relatively stable over training time in most experimental settings.Figures 8–10 present these completion-length results for the experiments corresponding to Figures 3–5.