PaperlayerView exact revision on arXiv

Source-linked AI summary

Learn from the Gap: Differential-Aware Advantage Pruning with Adaptive Rollout Sampling for GRPO

Jiahua Yang, Zhiwei Yang, Xianpeng Zhang, Dongyu Chen, Xing Chen, Tianhuang Su, Haonan Lu, Quanlong Guan, Kai Tang, Chuangchuang Wang

arXiv:2609.36932v3cs.AI

TL;DR

GRPO-style training incurs multi-rollout and repeated token-evaluation costs, while homogeneous trajectories can weaken learning signals. FastRL combines differential-aware pruning with adaptive rollout sampling based on historical pruning statistics. On Geometry3K and GeoQA8K-R1V, it reports speedups of up to 2.14× and 2.22×, respectively, alongside average-accuracy improvements of 1.31% and 1.09%.

  • Problem

    GRPO-style methods require multiple sampled completions per prompt and repeated per-token evaluation, while low-information or homogeneous trajectories can weaken learning signals and hinder optimization.

    Source: §1 · pp. 1–2

  • Method

    FastRL combines differential-aware pruning of trajectories using intra-group advantage distributions with adaptive rollout-scale adjustment based on historical pruning statistics.

    Source: §1 · p. 2

  • Results

    FastRL speeds up training by up to 2.14× on Geometry3K and 2.22× on GeoQA8K-R1V, improving average accuracy by 1.31% and 1.09%, respectively.

    Source: §1 · p. 2, §4.2 · p. 8

  • Takeaways & Limitations

    FastRL is presented as a plug-and-play framework for GRPO variants that targets training speed and policy-learning performance together.

    Source: §1 · p. 2

    MDAP is designed primarily for discrete reward signals; extending it to continuous rewards requires modifying reward representation, with discretization requiring task-specific choices and tuning.

    Source: D Limitations and Future Work · p. 19

Abstract

from arXiv · show

Recently, Group Relative Policy Optimization (GRPO) and its variants have been developed for policy optimization and demonstrated notable performance gains. However, these methods usually incur substantial computational overhead due to per-question multi-rollout sampling and repeated per-token probability evaluation across rollouts. Furthermore, low-information or highly homogeneous trajectories can degrade downstream learning signal efficiency, hindering model optimization and limiting final performance. To address these issues, we propose FastRL, a novel plug-and-play reinforcement learning framework that simultaneously improves training efficiency and the effectiveness of policy learning. Specifically, 1) We introduce an advantage-aware pruning strategy to selectively preserve high-advantage trajectories while maximizing inter-trajectory gradient diversity. 2) Then, we design an adaptive rollout sampling mechanism to dynamically adjust the sampling scale across different training stages based on historical pruning distributions, balancing exploration adequacy and computational efficiency. Experiments demonstrate that FastRL can be seamlessly integrated into GRPO, DAPO, and GSPO variants, achieving an average 2.07$\times$ training speedup on Geometry3K and GeoQA8K-R1V, along with an approximately 1.64\% improvement in average accuracy on visual reasoning benchmarks. Source codes will be available at https://github.com/Nicozwy/FastRL.

1 Introduction

GRPO reduces PPO’s memory burden, but per-question multi-rollout sampling and repeated token evaluation remain costly, while homogeneous trajectories weaken learning signals. FastRL combines differential-aware pruning with adaptive rollout sampling to address these costs, and achieves up to 2.14× end-to-end training speedup while outperforming leading GRPO variants.

1 Introduction

  • GRPO avoids PPO’s dedicated critic and reduces memory consumption, but still requires multiple completions per prompt and repeated per-token evaluation.Low-information and homogeneous trajectories can also degrade learning signals and hinder policy optimization.Source: §1 · pp. 1–2
  • Existing variants show an efficiency-performance trade-off: DAPO, GSPO, and GRESO accelerate training but often underperform, while CPPO yields marginal downstream gains.The passage describes these comparisons as motivation for a framework that pursues speed and policy-learning performance together.Source: §1 · p. 2
  • FastRL combines advantage-distribution-based pruning of low-information or homogeneous trajectories with rollout-scale adjustment based on historical pruning statistics.The two mechanisms are presented as complementary ways to reduce redundant computation while preserving informative trajectories.Source: §1 · p. 2
  • Up to 2.14× end-to-end training speedup on visual reasoning benchmarks, while outperforming leading GRPO variants.The reported speedup is an upper bound.Source: §1 · p. 2

2 Related Work

Related work frames GRPO as a lower-memory alternative to PPO, but identifies remaining rollout and policy-update costs alongside risks from homogeneous trajectories. Prior acceleration approaches filter questions or prune trajectories, each with stated efficiency limitations.

Reinforcement Learning

  • GRPO removes PPO’s dedicated critic by estimating baselines from group-wise relative rewards, but repeated rollouts and policy updates remain costly.The passage also links homogeneous trajectories to less informative learning signals, reasoning bias, and eventual mode collapse.Source: §2 · p. 3

GRPO-based Training Acceleration

  • Prior methods filter questions using historical difficulty and reward variance or prune low-absolute-advantage trajectories, but incur overhead or leave rollout latency untouched.Question-level statistics add storage and computation; uniform pruning can cause a “bucket effect” across questions.Source: §2 · p. 3

3 Method

The method builds on GRPO’s group-relative advantages, then prunes redundant trajectories and adapts rollout counts using historical pruning behavior. These steps target both gradient-update computation and rollout sampling.

3.1 Preliminaries

  • GRPO uses relative rewards within each question’s rollout group to form advantages, while low-information and homogeneous trajectories can provide redundant or weak learning signals.Similar trajectories may produce correlated gradients, motivating selection of more diverse representatives.Source: §3.1 · pp. 3–4

3.2 Maximizing Differential Advantage Pruning

  • MDAP groups trajectories by advantage and uses token-level Jaccard similarity as a proxy for gradient similarity when pruning redundant trajectories.Within each advantage group, it prunes a trajectory when its differential score against a previously retained trajectory falls below τ.Source: §3.2 · p. 5

3.3 Adaptive Rollout Sampling

  • ARS adjusts rollout sampling from accumulated pruning statistics: heavy pruning calls for fewer rollouts, while high retention indicates that more trajectories may be needed.The mechanism uses average and maximum retained update counts and adjusts sampling as training progresses.Source: §3.3 · p. 6

4 Experiments

Across two visual-reasoning datasets and model sizes, FastRL improves training speed and average accuracy over the compared baselines. Ablations and further analyses examine component contributions, sampling and pruning behavior, generalization, and exploration.

4.1 Experimental Settings

  • Experiments use Geometry3K and GeoQA8K-R1V in-domain and four out-of-domain visual mathematical reasoning benchmarks, comparing GRPO, DAPO, GSPO, CPPO, and GRESO.The reported setup uses accuracy and training time as metrics, with τ = 0.8 and an initial rollout sampling number of 8.Source: §4.1 · pp. 6–7

4.2 Overall Results and Analysis

  • Up to 2.14× and 2.22× training speedups accompany average-accuracy gains of 1.31% and 1.09% on Geometry3K and GeoQA8K-R1V, respectively.FastRL consistently outperforms baselines in training efficiency and average accuracy across both datasets and both model sizes; CPPO accelerates training but can reduce accuracy.Source: Table 1 · p. 7, §4.2 · pp. 7–8

4.3 Ablation Study

  • Removing MDAP causes the largest ablation decline: Qwen2.5-VL-7B-Instruct accuracy drops by 0.72% and 1.21%, while speedups fall from 2.14× to 1.15× and 2.13× to 1.10×.Removing ARS causes smaller but consistent declines in accuracy and training efficiency.Source: Table 2 · p. 8, §4.3 · p. 8

4.5 Generalization Across Architectures and Datasets

  • On Llama3.1-8B-Instruct trained on text-only MATH, FastRL reaches 73.04% in-domain accuracy and 2.30× speedup, and matches GRPO’s 7.50% AIME average in 9.22 versus 21.16 hours.CPPO and GRESO achieve AIME averages of 4.17% and 3.33%, respectively.Source: Table 3 · p. 9, §4.5 · p. 9

4.4 Analysis of FastRL’s Training Efficiency and Parameter Sensitivity

  • At τ = 0.8, average accuracy peaks with speedups of 2.14× and 2.13×; τ = 1.0 raises speedups to 2.35× and 2.27× but lowers accuracy.FastRL’s rollout scale adapts using historical pruning distributions, while the number of highly distinct trajectories declines later in training; it retains fewer, more representative trajectories than CPPO.Source: §4.4 · pp. 8–9, Figure 3 · p. 9

4.6 Further Analysis

  • FastRL improves avg@16 over GRPO by 3.15 on Geometry3K and 2.14 on GeoQA8K-R1V, and achieves the best Pass@16 on every benchmark.The analysis reports that pruning redundant trajectories sharing identical advantages does not shrink exploration space.Source: §4.6 · pp. 9–10, Table 4 · p. 10
  • On Geometry3K, adding Jaccard-distance pruning to advantage grouping raises average accuracy from 52.09 to 52.63.Random pruning accelerates training but degrades performance as the pruning ratio increases; advantage-only pruning gives a 2.35× speedup with a slight accuracy improvement.Source: Table 5 · p. 10, §4.6 · p. 10

5 Conclusion

FastRL combines structured advantage-diversity pruning with adaptive rollout sampling based on pruning history. The paper reports training acceleration and accuracy improvements for GRPO-based methods.

5 Conclusion

  • FastRL retains informative trajectories with large differences in gradient contribution, reducing redundant forward computation.Its pruning strategy is based on maximizing advantage diversity.Source: §5 · p. 10
  • Adaptive rollout sampling adjusts the sampling scale online according to historical pruning distributions, balancing exploration sufficiency and computational efficiency.Source: §5 · p. 10
  • FastRL consistently delivers both training acceleration and accuracy improvements for GRPO-based methods.Source: §5 · p. 10

A Experimental Datasets.

The experiments use in-domain multimodal and text-based mathematics datasets, alongside out-of-domain reasoning benchmarks. Figure 4 provides examples from three training datasets, and Table 6 is identified as reporting dataset statistics.

A Experimental Datasets.

  • The in-domain datasets include Geometry3K, GeoQA8K-R1V2, and MATH, spanning multimodal geometric reasoning and text-based mathematical problems.The Geometry3K experiments use 2,702 problems from its training and test splits; MATH contains 12,500 problems with step-by-step solutions.Source: A Experimental Datasets · p. 13
  • Out-of-domain evaluation includes visual mathematics benchmarks MathVision, MathVista, MathVerse, and WeMath, plus the text-only AIME3 benchmark.The passage describes their question counts and, for MathVision and WeMath, their coverage.Source: A Experimental Datasets · p. 13
  • Figure 4 shows examples from Geometry3K, GeoQA8k-R1V, and MATH, while Table 6 presents dataset statistics.Source: Figure 4 · p. 13, Table 6 · p. 13

B Additional Method Details and Analysis

The algorithm samples rollouts, prunes trajectories within groups of equal advantage, and uses historical pruning information to adapt rollout sampling. It also aligns the selected set to a divisor for DP alignment.

B.1 Algorithm

  • Adaptive rollout sampling uses an epoch-level buffer and updates its sampling scale at epoch boundaries using historical pruning information.Source: §B.1 · p. 14
  • For each question, FastRL samples trajectories, computes rewards and advantages, then groups trajectories by unique advantage values.Source: §B.1 · p. 14
  • Within each advantage group, trajectories are ordered by absolute advantage, and selected trajectories are added to the retained set.The algorithm also ensures each group has at least one retained trajectory and adds unselected trajectories with the highest absolute advantage round-robin until the set size is divisible by d.Source: §B.1 · p. 14

B.2 Correlation Analysis of Token-Level Jaccard Similarity and Gradient Cosine Similarity

An empirical study on Geometry3K and GeoQA8K-R1V finds a strong positive correlation between token-level Jaccard similarity and gradient cosine similarity for same-advantage trajectory pairs. The results support using token-level Jaccard similarity as a proxy for gradient similarity in the pruning strategy.

B.2 Correlation Analysis of Token-Level Jaccard Similarity and Gradient Cosine Similarity

  • The study sampled 256 questions across training epochs per dataset, yielding 2,048 trajectories, and compared same-advantage trajectory pairs.Source: §B.2 · p. 15
  • The observed relationship supports token-level Jaccard similarity as a proxy for gradient similarity in the proposed pruning strategy.Source: §B.2 · p. 15
  • Both datasets show a strong positive correlation between token-level Jaccard similarity and gradient cosine similarity, with high Pearson and Spearman coefficients.Figure 5 presents the Geometry3K results on the left and GeoQA8K-R1V results on the right.Source: Figure 5 · p. 15, §B.2 · p. 15

C.1 Training Efficiency under Low Policy Lag.

Under low policy lag, FastRL reaches comparable peak accuracy to GRPO and higher accuracy than CPPO while completing 200 steps faster. Two further optimizations reduce redundant work in trajectory updates and log-probability computation.

C.1 Training Efficiency under Low Policy Lag.

  • 8.8 h for 200 steps gives FastRL a 2.9× wall-clock speedup over GRPO, without sacrificing accuracy; its peak accuracy is comparable to GRPO and higher than CPPO.Figure 6 compares step time, validation accuracy by training step, and validation accuracy by cumulative GPU wall-clock time.Source: §C.1 · p. 15, Figure 6 · p. 16
  • The reported 2.9× speedup exceeds the 2.14× Geometry3K result in Table 1 because this subsection uses two further optimizations that remove remaining redundant computation.The pruning criterion and optimization objective are unchanged.Source: §C.1 · p. 16
  • The improved version fixes policy updates per step and splits retained trajectories evenly, making backfilling overhead negligible while keeping optimization-step counts aligned with baselines.This replaces the main-experiment setup, where a fixed mini-batch size could require backfilling and cause update counts to vary with the pruning ratio.Source: §C.1 · p. 16
  • Reordering the pipeline to prune before log-probability computation lets old-policy and reference-model forward passes run only on retained trajectories, removing redundant computation without changing the optimization objective.The pruning criterion uses advantages and token-level Jaccard similarity, while group-relative advantages depend only on rewards.Source: §C.1 · p. 16

C.2 Training Dynamics and Computational Overhead under High Policy Lag.

Under high policy lag, FastRL maintains low token consumption, stable entropy, and competitive or superior accuracy with less training time across Geometry3K and GeoQA8K-R1V. The results indicate a favorable efficiency, stability, and performance trade-off across policy-lag settings.

C.2 Training Dynamics and Computational Overhead under High Policy Lag.

  • FastRL maintains the lowest total token consumption throughout training on both datasets under high policy lag.Figure 7 compares FastRL with GRPO, DAPO, CPPO, GSPO, and GRESO on token consumption, policy entropy, and accuracy.Source: §C.2 · p. 17, Figure 7 · p. 17
  • FastRL maintains stable entropy with moderate fluctuations, while existing GRPO-style methods show entropy surges and unstable or degenerative behavior under high policy lag.The authors attribute FastRL's stability to pruning homogeneous trajectories and preserving diverse advantage patterns and gradient contributions.Source: §C.2 · p. 17
  • FastRL achieves competitive or superior accuracy with substantially less training time, and reaches stable high accuracy earlier than the baselines.Some baselines show later accuracy degradation; others approach comparable performance only at higher computational cost.Source: Figure 7 · p. 17, §C.2 · p. 18
  • Across low- and high-policy-lag regimes, FastRL accelerates training while preserving strong exploration and task performance, and under high lag it further mitigates entropy instability and performance degradation.The high-lag setting uses 16 sequential policy updates per step, compared with four in the low-lag setting.Source: §C.2 · pp. 16–18

C.3 Supplementary Parameter Sensitivity Analysis.

The sensitivity analysis reports reduced training time and improved average accuracy as τ increases, with best average performance at τ = 0.8.

C.3 Supplementary Parameter Sensitivity Analysis.

  • Average accuracy across in-domain and benchmark evaluations is best at τ = 0.8, while training time substantially decreases as τ increases.The analysis covers Geometry3K and GeoQA8K-R1V with Qwen2.5-VL-7B-Instruct and MATH with Llama3.1-8B-Instruct; numerical results are reported in Table 7.Source: §C.3 · p. 18, Table 7 · p. 18

C.4 Extension to Continuous-Reward Settings.

On a grounding task with continuous rewards, FastRL is reported to train faster than GRPO while achieving more competitive English and Chinese general-benchmark performance. The authors present this as an extension beyond discrete mathematical reasoning.

C.4 Extension to Continuous-Reward Settings.

  • On the grounding task under continuous rewards, FastRL achieved a 1.76× training speedup over GRPO with Qwen3.5-9B, while delivering more competitive performance on English and Chinese general benchmarks.Source: §C.4 · p. 19, Table 8 · p. 19
  • For continuous rewards, trajectories are assigned to the same advantage group when their advantage difference is below 10^-4.The grouping adjustment is intended to prevent infinitesimal numerical differences from fragmenting advantage groups.Source: §C.4 · p. 19

D Limitations and Future Work

The paper’s stated scope is primarily discrete-reward reinforcement learning, while continuous rewards require task-specific adaptation. The checklist also records that the main experiments used a single run because of high training cost.

D Limitations and Future Work

  • MDAP is designed primarily for discrete-reward RL; continuous rewards require task-specific representation choices and careful tuning if discretized piecewise.The authors suggest piecewise discretization as one possible approximation.Source: D Limitations and Future Work · p. 19

NeurIPS Paper Checklist

  • The checklist says an anonymous reproduction link is provided, with training and test details referenced to Section 4.1.GPU configurations and training times are also referenced as reported in the experimental tables.Source: NeurIPS Paper Checklist · pp. 21–23
  • The paper reports no error bars or significance information because high training cost restricted main experiments to a single run.Source: NeurIPS Paper Checklist · p. 22

Related papers

Loading 2609.36932v3…