Source-linked AI summary

Prune as You Generate: Online Rollout Pruning for Faster and Better RLVR

Haobo Xu, Sirui Chen, Ruizhong Qiu, Yuchen Yan, Chen Luo, Monica Cheng, Jingrui He, Hanghang Tong

arXiv:2603.24840v1cs.CL

TL;DR

RLVR rollout-based training is expensive and can produce weak learning signals when binary rewards within groups are nearly uniform. ARROL uses an online quality head to prune partial rollouts toward balanced survivors, and reuses the head for test-time weighting. Across GRPO and DAPO, it improves accuracy while accelerating training and test-time scaling.

  • Problem

    GRPO and DAPO require many rollouts per prompt, while nearly all-correct or all-incorrect groups have low reward variance and weak learning signals.

  • Method

    ARROL trains a lightweight quality head to estimate partial-rollout success, prunes during generation toward correctness-balanced groups, and uses the scores for test-time candidate weighting.

  • Results

    +2.30 to +2.99 average accuracy in GRPO/DAPO training, up to 1.7x training speedup, and +8.33 average-accuracy gains in test-time scaling.

  • Takeaways & Limitations

    Online rollout pruning can reduce RLVR computation while strengthening group-normalized learning signals through more balanced retained rollouts.

  • Takeaways & Limitations

    The study evaluates mathematical RLVR tasks and does not validate other reward-based domains; pruning also requires generation to reach Ldetect = 512 before taking effect.

Abstract

from arXiv · show

Reinforcement Learning with Verifiable Rewards (RLVR) has significantly advanced the reasoning capabilities of Large Language Models (LLMs). However, methods such as GRPO and DAPO suffer from substantial computational cost, since they rely on sampling many rollouts for each prompt. Moreover, in RLVR the relative advantage is often sparse: many samples become nearly all-correct or all-incorrect, yielding low within-group reward variance and thus weak learning signals. In this paper, we introduce arrol (Accelerating RLVR via online Rollout Pruning), an online rollout pruning method that prunes rollouts during generation while explicitly steering the surviving ones more correctness-balanced to enhance learning signals. Specifically, arrol trains a lightweight quality head on-the-fly to predict the success probability of partial rollouts and uses it to make early pruning decisions. The learned quality head can further weigh candidates to improve inference accuracy during test-time scaling. To improve efficiency, we present a system design that prunes rollouts inside the inference engine and re-batches the remaining ones for log-probability computation and policy updates. Across GRPO and DAPO on Qwen-3 and LLaMA-3.2 models (1B-8B), arrol improves average accuracy by +2.30 to +2.99 while achieving up to 1.7x training speedup, and yielding up to +8.33 additional gains in average accuracy in test-time scaling. The code is available at https://github.com/Hsu1023/ARRoL.

1 Introduction

RLVR training methods generate many rollouts, creating high computational cost and often weak learning signals when group rewards are imbalanced. ARROL prunes rollouts online using quality predictions, balances survivors, and reports accuracy and speed gains.

  • Motivation: GRPO-style RLVR is costly because each prompt requires many rollouts, while post-generation pruning does not reduce rollout-generation time.Highly imbalanced binary rewards can also produce weak or vanishing group-normalized learning signals.
  • Method: ARROL scores partial rollouts with a lightweight quality head, prunes during generation, and selects a correctness-balanced subset for training.The scores estimate rollout success and can also support test-time candidate weighting.
  • Results: +2.30 to +2.99 average accuracy for GRPO/DAPO training on Qwen-3 and LLaMA-3.2 models (1B-8B) with ARROL.The method reduces compute while explicitly controlling within-group reward balance.
  • Results: +8.33 gains in average accuracy when the learned quality head supplies voting weights for test-time scaling.The quality head is reused to improve final-answer aggregation.
  • Results: 1.6−1.7× speedup is achieved by pruning inside the generation backend and re-batching surviving rollouts in the frontend.The system reallocates freed backend capacity and filters survivors before log-probability computation and optimization.

2 Related Work

Prior RLVR efficiency methods reduce rollout construction, reuse historical information, prune after generation, or accelerate decoding with speculative or low-precision techniques. ARROL instead targets online pruning guided by rollout-quality prediction.

  • Efficient RLVR: Prior methods prune completions or downsample rollout pools after generation, so they do not directly reduce token-generation cost during rollout.Some update-phase methods can even increase rollout-generation cost.
  • Efficient RLVR: Speculative-decoding approaches rely on historical sequences that may be unsuitable for commonly adopted small-epoch settings.Other generation accelerators use low-precision or quantized rollouts rather than online quality-guided pruning.
  • Test-time Scaling: Test-time scaling methods allocate additional inference compute through sequential or parallel strategies, including extended reasoning or answer revisiting.The related-work discussion distinguishes sequential scaling from parallel approaches.

3 Preliminaries

RLVR uses verifiable binary rewards for generated answers, while GRPO learns from relative performance among rollouts sharing a prompt. Its computational cost and sparse within-group signals motivate rollout pruning.

  • Trace Confidence: Trace confidence averages token-level confidence or aggregates window-level uncertainty along a generated reasoning trace.Window-based variants can use the minimum window value or the average of the bottom 10% windows.
  • RLVR: RLVR assigns reward 1 when a generated answer is equivalent to the ground-truth answer and 0 otherwise.The reward is represented as an indicator over answer equivalence.
  • GRPO: GRPO estimates advantages from relative performance within each prompt’s rollout group without using a value model.The objective uses importance ratios and group-relative rewards.
  • GRPO: GRPO must generate and process many rollouts per prompt, and all-0 or all-1 reward groups can make normalized advantages zero.This creates substantial time cost and potentially vanishing gradients.

4 Method

ARROL balances binary rewards by pruning rollouts online with quality predictions, then uses the learned scores to select survivors and support inference-time weighting. Its design combines posterior-guided pruning, quality prediction, calibration, and an intermediate detection length.

  • 4.1 Pruning Improves Sample Balance: ARROL targets a within-group positive ratio near ρ = 0.5, where binary-reward variance and the associated learning signal are maximized.Theoretical results state that accurate posterior-guided pruning reduces deviation from the target ratio and can improve training effectiveness.
  • 4.2 Quality Prediction Head: A quality head scores partial rollouts and maps the scores to estimated success probabilities for early pruning.The method uses hidden representations because token-likelihood confidence can misalign with final correctness.
  • 4.2 Quality Prediction Head: Ldetect = 512 balances pruning reliability and time cost because correlation plateaus as generation length reaches 512 while time cost continues increasing.The detection length is an intermediate point at which partial-rollout quality is evaluated.
  • 4.2 Quality Prediction Head: Quality-head scores correlate more reliably with final correctness than trace confidence, whose signals can be distorted by reflection-related or formula-heavy tokens.The comparison is reported across Math500 and Dapo17k; quality-head scores achieve consistently higher Spearman correlation.
  • 4.2 Quality Prediction Head: Online score calibration normalizes quality scores into bins and estimates posterior success probabilities from historical positive and negative rollouts in a sliding buffer.The calibrated posterior feeds the pruning design.
  • 4.1 Pruning Improves Sample Balance: Each rollout receives a survival probability based on its estimated success probability, controlling both the expected keep ratio κ and the retained positive ratio ρ.The survival function is monotonic in the difference between the target ratio and the estimated success probability.

Appendix A.4.

The system places rollout pruning inside a frontend-backend inference architecture, re-batching surviving rollouts for downstream computation. The quality head is trained online from final rewards and also supports test-time candidate scoring.

  • System Design: Rollout pruning is implemented inside the generation backend, while the frontend handles log-probability computation and policy updates.The architecture uses verl with vLLM as the high-throughput rollout-generation backend.
  • System Design: Each training step generates rollouts, computes log-probabilities and advantages, then updates the policy.
  • System Design: The quality head is trained online from rollout final rewards using cross-entropy while stopping gradients to the backbone.This design keeps the additional overhead negligible.
  • System Design: At test time, quality-head scores are converted into rank-based weights for aggregating completed reasoning candidates.Scores are sorted and linearly rescaled to the [0, 1] range because they are uncalibrated logit-like values.

5 Experiment

Across GRPO and DAPO experiments, ARROL improves accuracy while reducing training cost, balancing within-group rewards and extending the quality head to test-time voting.

  • Performances on GRPO: The quality head reaches approximately 80% prediction accuracy, including 82.37% on Qwen3-1.7B, supporting its use for early pruning decisions.The head is trained during the RLVR pipeline.
  • Performances on DAPO and GRPO: +2.99 average accuracy with DAPO and 1.70× end-to-end speedup, while GRPO gains +2.30 to +2.87 on Qwen-3 and +2.86 on LLaMA-3.2.ARROL generally outperforms vanilla training across the evaluated benchmarks and model settings.
  • Test-time Scaling: Up to +8.33 additional average-accuracy improvement over DeepConf shows that the learned quality head can improve test-time voting across datasets and models.The quality head is used as a voting weight alongside vanilla majority voting and DeepConf.
  • Comparison with Random Pruning: ARROL moves rollout groups toward a 0.5 positive ratio and increases E[ρ̂(1 − ρ̂)], which is proportional to binary within-group reward variance.The balanced subset is associated with stronger non-degenerate learning signals and better final performance.
  • Efficiency Decomposition: Log-probability computation and model updates become about 2× faster, while rollout generation reaches 1.46× speedup because pruning occurs after an initial detection stage.The efficiency decomposition attributes the larger update-side savings to discarding roughly half of the rollouts.
  • Ablation Study on keep ratio κ: A keep ratio κ = 0.5 provides the reported accuracy–efficiency trade-off, whereas κ = 0.25 removes too many rollouts and slightly reduces performance.Smaller κ values increase speedup, but excessive pruning can hurt accuracy.

6 Conclusion

ARROL prunes rollouts online using a lightweight quality head to retain more correctness-balanced samples, improving RLVR efficiency and learning signals. It also reuses the head for test-time voting, achieving accuracy gains while delivering up to 1.7× training speedup.

  • Conclusion: ARROL prunes rollouts during generation, steers survivors toward balanced 0/1 rewards, and trains a quality head to predict partial-rollout success.The system integrates pruning into generation and re-batches surviving rollouts for subsequent computation.
  • Conclusion: Across models, ARROL improves accuracy with up to 1.7× training speedup, while the quality head provides additional gains when used as test-time voting weights.The paper characterizes this as a “less rollouts, more learning” approach for RLVR training and scaling.

Limitations

The study is limited to mathematical RLVR tasks with verifiable rewards and does not validate the approach in other reward-based domains. Training-time pruning also incurs a detection-length cost before pruning can begin, which can reduce rollout-generation speedups.

  • Validation is limited to mathematical RLVR tasks with verifiable rewards, leaving domains such as UI interaction and tool-use agents untested.
  • Pruning begins only after partial rollouts reach the intermediate detection length Ldetect = 512, so generation speedups can be smaller than later-phase savings.

A.1 Proof of Theorems in Sec. 4.1

The appendix formalizes posterior-guided pruning under binary rollout labels and establishes that the selected prune remains close to the target positive ratio when posterior estimates are accurate and samples concentrate.

  • Concentration: Under conditionally independent binary labels, Hoeffding concentration bounds the realized retained positive ratio around its conditional expectation after the data-dependent prune.
  • Posterior-guided pruning: Posterior-guided pruning chooses the removal index that minimizes the estimated deviation between the retained positive ratio and target ρ.
  • Improving prune: If the current true positive ratio exceeds ρ and a rollout has posterior probability above that ratio, removing it strictly reduces deviation toward ρ.
  • High-probability guarantee: Theorem 4.2 combines posterior-estimation accuracy, near-optimal pruning, and concentration to give high-probability closeness of the retained ratio to target ρ.
Loading 2603.24840v1…