PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Data Efficiency for LLM Reinforcement Fine-tuning Through Difficulty-targeted Online Data Selection and Rollout Replay

Yifan Sun, Jingyan Shen, Yibin Wang, Tianyu Chen, Zhendong Wang, Mingyuan Zhou, Huan Zhang

arXiv:2506.05316v4cs.LGcs.AIcs.CL

TL;DR

LLM RL fine-tuning is computationally expensive, and its data-efficiency problem has received limited attention. The paper combines adaptive-difficulty online selection with rollout replay, reducing fine-tuning time while matching original GRPO performance across six LLM–dataset combinations.

  • Problem

    LLM RL fine-tuning is computationally expensive, while relatively little work has addressed improving its data efficiency.

  • Method

    The method selects questions with moderate adaptive difficulty using a small reference set and similarity-based prediction, then reuses recent rollouts through replay.

  • Results

    23% to 62%: the method reduces RL fine-tuning time across six LLM–dataset combinations while achieving the same performance as original GRPO.

  • Takeaways & Limitations

    Difficulty-targeted selection reduces required training steps, while rollout replay reduces per-step rollout cost without sacrificing performance.

  • Takeaways & Limitations

    Experiments use the same type and number of GPU devices across comparisons, and the theoretical gradient result relies on mild assumptions about rewards and likelihood gradients.

Abstract

from arXiv · show

Reinforcement learning (RL) has become an effective approach for fine-tuning large language models (LLMs), particularly to enhance their reasoning capabilities. However, RL fine-tuning remains highly resource-intensive, and existing work has largely overlooked the problem of data efficiency. In this paper, we propose two techniques to improve data efficiency in LLM RL fine-tuning: difficulty-targeted online data selection and rollout replay. We introduce the notion of adaptive difficulty to guide online data selection, prioritizing questions of moderate difficulty that are more likely to yield informative learning signals. To estimate adaptive difficulty efficiently, we develop an attention-based framework that requires rollouts for only a small reference set of questions. The adaptive difficulty of the remaining questions is then estimated based on their similarity to this set. To further reduce rollout cost, we introduce a rollout replay mechanism inspired by experience replay in traditional RL. This technique reuses recent rollouts, lowering per-step computation while maintaining stable updates. Experiments across 6 LLM-dataset combinations show that our method reduces RL fine-tuning time by 23% to 62% while reaching the same level of performance as the original GRPO algorithm. Our code is available at https://github.com/ASTRAL-Group/data-efficient-llm-rl.

1 Introduction

LLM RL fine-tuning is computationally expensive, while data efficiency has received comparatively little attention. The paper introduces difficulty-targeted online selection and rollout replay to reduce training steps and per-step computation.

  • Motivation: LLM RL fine-tuning remains computationally expensive despite growing algorithmic progress.A 1.5B-parameter model reportedly required over 3,800 A100 GPU hours for 40K samples.
  • Difficulty-targeted Online Data Selection: Difficulty-targeted online data selection prioritizes questions with moderate adaptive difficulty because they are more likely to provide meaningful learning signals.Adaptive difficulty reflects how likely the current policy is to fail on a question and changes during training.
  • Difficulty Prediction: An attention-based framework estimates adaptive difficulty using rollouts from only a small reference set and similarity-based attention for remaining questions.This avoids generating full rollouts for every question at each step.
  • Rollout Replay: Rollout replay reuses recent rollouts to reduce per-step computation while maintaining stable updates.The method combines newly generated and retrieved recent rollouts.
  • Results: 23% to 62%: the method reduces RL fine-tuning time across six LLM–dataset combinations while achieving the same performance as original GRPO.The reported contribution combines fewer required steps with lower per-step rollout cost.

2 Related Work

Prior work studies online data selection and experience replay, but LLM RL methods still face substantial inefficiency from repeatedly collecting fresh rollouts. The paper builds on these ideas for adaptive selection and rollout reuse.

  • Online Data Selection: Static data selection assumes sample importance remains fixed, whereas online methods repeatedly adapt selection to the model’s evolving state.This adaptability is especially relevant to RL because policy updates and environment dynamics are non-stationary.
  • Experience Replay: PPO and GRPO rely on freshly collected rollouts for policy updates, creating substantial data inefficiency and computational overhead.Experience replay addresses this by retaining recent transitions for multiple passes.

3 Problem Setup

The problem setup focuses on improving GRPO’s data efficiency by selecting questions that provide stronger learning signals and adapting selection as the policy changes.

  • GRPO: GRPO samples response groups from the old policy and computes normalized group-level advantages from their rewards.The objective uses these advantages to update the policy.
  • GRPO Objective: The GRPO objective uses a clipped policy update based on the probability ratio between current and old policies, together with a KL penalty relative to a reference policy.The cited formulation describes both the clipped update and reference-policy regularization.
  • Problem Setup: Standard GRPO uniformly samples questions, although questions too easy or too hard for the current policy may produce weak gradient signals.This can slow convergence because not all questions contribute equally to learning progress.
  • Online Selection: Online selection chooses a fixed-size batch according to the current policy and repeats this selection throughout training.More frequent selection improves adaptation but increases computational overhead.

4 Method

The method combines adaptive-difficulty prediction and targeted online selection with rollout replay. It estimates difficulty from a small reference set, selects questions near 0.5, and reuses informative recent rollouts with importance correction.

  • Adaptive Difficulty: Adaptive difficulty measures how challenging each question is for the current policy and changes as training progresses.It is defined from the policy’s rollout failure behavior.
  • Difficulty Prediction: The predictor generates rollouts only for a small reference subset, then estimates remaining questions’ difficulty using embedding similarity and attention-weighted averaging.Reference questions provide the difficulty values used to estimate unlabeled questions.
  • Difficulty-targeted Selection: Questions with predicted adaptive difficulty closest to 0.5 are prioritized because they are neither too easy nor too hard for the current policy.The paper links this target to avoiding degenerate zero-advantage cases in GRPO.
  • Theoretical Justification: Under the stated Bernoulli-reward assumptions, the expected squared gradient norm is maximized when p = 0.5.This provides theoretical support for targeting moderate difficulty.
  • Dynamic Coverage: The selection mechanism promotes broader coverage because repeatedly trained questions move away from the target difficulty and become less likely to be sampled again.This allows under-explored questions to enter the selection pool over time.
  • Rollout Replay: Rollout replay generates new rollouts for only a fraction δB of each batch and fills the remainder from a FIFO buffer of recent rollouts.The buffer has capacity C and stores selected rollout information for future updates.
  • Replay Correction: Importance sampling with the behavior policy corrects the bias introduced when past rollouts are reused under the current policy.The modified GRPO loss is designed to preserve training stability while reducing rollout generation.

5 Experiments

Experiments evaluate DOTS and rollout replay across multiple models, datasets, benchmarks, and a non-mathematical domain. The combined method reaches comparable performance with fewer steps and lower per-step and total training costs, while its adaptive-difficulty predictor tracks policy behavior and selects more effective questions.

  • Experimental Setup: Experiments use three Qwen model scales, four mathematical training datasets, and four standard mathematical reasoning benchmarks.The models are Qwen2.5-Math-1.5B, Qwen2.5-3B, and Qwen2.5-Math-7B; training uses MATH, DeepScaleR-40K, ORZ, and DeepMath-103K.
  • Main Results: 13.33% to 56.67% fewer training steps are required for DOTS+RR to match original GRPO’s final performance at 60 steps.Results are reported across all LLM–dataset combinations and attribute faster convergence to prioritizing informative training samples.
  • Main Results: 11%–13% lower per-step training time is achieved by reducing the number of fresh rollouts through rollout replay.Rollout generation accounts for approximately 47%, 46%, and 54% of per-step time for the 1.5B, 3B, and 7B models, respectively.
  • Main Results: 40.7% average total training-time reduction is achieved across six settings, with a maximum improvement of 61.65% on Qwen2.5-3B trained on DeepMath.DOTS+RR maintains higher accuracy at almost every training step than original GRPO.
  • Effectiveness of Adaptive Difficulty Prediction Framework: ρ > 0.7 correlation is achieved between predicted and ground-truth adaptive difficulties across settings and training dynamics.The framework is intended to track policy behavior as training evolves.
  • Effectiveness of Adaptive Difficulty Prediction Framework: 25.4% more effective questions are selected on average by DOTS than original GRPO across LLM–dataset combinations.Effective questions have adaptive difficulty strictly between 0 and 1, whereas difficulty values of 0 or 1 yield zero group-normalized advantage and no gradient signal.
  • Effectiveness of Adaptive Difficulty Prediction Framework: 10K samples are processed in 1.71 seconds at deployment using cached question embeddings and a lightweight encoder.The predictor is designed to add minimal computational overhead and scale to large datasets.
  • Analysis and Discussion: 20% lower training time is obtained by adding rollout replay to DOTS while preserving DOTS performance.DOTS alone produces a steeper learning curve than original GRPO, while rollout replay reduces per-step cost.

6 Conclusion

The paper proposes two techniques to improve data efficiency in LLM RL fine-tuning: Difficulty-targeted Online Data Selection and Rollout Replay.

  • The paper proposes Difficulty-targeted Online Data Selection and Rollout Replay to improve data efficiency in LLM RL fine-tuning.

A.1 Limitations and Future Work

The paper identifies limitations involving reference-set construction, replay strategy, input representations, and rollout-generation overhead, while outlining concrete extensions.

  • Randomly sampling the reference set can affect prediction quality, motivating more diverse or fixed reference questions with sufficient training-set coverage.
  • The current replay strategy randomly reuses eligible rollouts, leaving more principled approaches such as prioritized experience replay for future work.
  • Reference solutions may slightly improve difficulty prediction but are unavailable for some datasets, including DeepScaler and ORZ.
  • Generating reference-set rollouts can add nontrivial overhead, although reusing near-0.5 predicted-difficulty rollouts reduces rollout cost by 4–12% per step while maintaining final performance.
  • Compared with static selection, DOTS is described as more efficient and applicable in an online data-selection setting.

B Proofs

The proofs analyze gradient-signal strength under binary and multi-component rewards, showing that group-normalized policy-gradient signal is greatest near a 50% success rate.

  • Theorem 1 analyzes a single question with independently sampled policy responses and binary Bernoulli rewards.
  • Under mild assumptions, the expected squared gradient norm is proportional to p(1 −p) · (1 −1/G) and is maximized when p = 0.5.
  • The proof focuses on the leading term T1 after assuming the correction term T2 is negligible.
  • Empirical ratios of 0.081 ± 0.0065 and 0.081 ± 0.0051 support the weak-dependence assumption on two LLM-dataset combinations.
  • The same insight extends to independent correctness and format rewards, with the expression maximized when both α = 0.5 and β = 0.5.

C.1 Design and Implementation Details

The framework predicts adaptive difficulty efficiently by combining a frozen language-model backbone with lightweight trainable components, reference-set labels, and similarity-based attention. It is trained and illustrated using cross-model mathematical question data and attention patterns.

  • Prediction architecture: The predictor freezes Qwen2.5-Math-1.5B-Instruct and adds a lightweight adapter plus calibration head.The adapter and calibration head provide the trainable components for adaptive-difficulty prediction.
  • Prediction architecture: The adapter is a three-hidden-layer GELU MLP with 896 units per layer, dropout 0.1, and LayerNorm on its projection output.The calibration head uses reference difficulty statistics to produce a scale parameter and bounded bias term.
  • Training data: Each training instance pairs a query question with a reference set of scored questions and a ground-truth difficulty label.Data are collected from models disjoint from the policy models and assembled into Dpred-train.
  • Training objective: The adapter and calibration head are trained with binary cross-entropy using the calibrated predicted difficulty.The calibrated prediction is denoted by ˆdq,cal.
  • Attention behavior: Attention assigns higher weights to reference questions sharing mathematical topics and structures with the target, while down-weighting unrelated questions.A qualitative example highlights rhombus and incircle geometry among highly attended references.

D.1 Training Datasets and Models

The experiments use three Qwen model sizes and four open-source mathematical reasoning datasets, with curated or filtered training subsets and standardized RL recipes. The datasets span competition, reasoning, and high-difficulty mathematics.

  • Models: Experiments use Qwen2.5-Math-1.5B, Qwen2.5-3B, and Qwen2.5-Math-7B as policy models.These three model sizes define the evaluated model configurations.
  • Training datasets: MATH contributes Level 3–5 competition problems after merging splits and excluding overlap with MATH500.The retained questions cover seven mathematical subjects and five original difficulty levels.
  • Training datasets: ORZ supplies 57,000 cleaned reasoning problems, from which 8,192 are sampled for training.Rule-based and LLM-based filters target evaluability and difficulty balance.
  • Training datasets: DeepMath-103K is a decontaminated collection focused on high-difficulty mathematical problems, with 8,192 training problems sampled.The construction supports reliable benchmark evaluation.
  • RL recipe: GRPO uses a shared configuration with correctness-only rewards, no KL regularization terms, and DOTS selection every two RL fine-tuning steps.Exact answer matches receive reward 1 and otherwise receive 0.

D.4 Evaluation Details

Evaluation averages accuracy across four mathematical reasoning benchmarks to reduce benchmark-specific variance. Very small benchmarks are excluded because their scores fluctuate substantially and yield unreliable comparisons.

  • Evaluation protocol: Performance is evaluated with average accuracy across four commonly used mathematical reasoning benchmarks.The aggregation is intended to mitigate benchmark-specific variance.
  • Evaluation benchmarks: The evaluation suite includes GSM8K, a 1,319-question benchmark requiring multi-step arithmetic reasoning.GSM8K represents grade-school mathematical word problems.
  • Evaluation benchmarks: MATH500 is used as an evaluation subset of the MATH test split and is excluded from the MATH training data.This separation prevents overlap between training and evaluation questions.
  • Evaluation benchmarks: Minerva Math contains 272 undergraduate-level science and mathematics questions from MIT OpenCourseWare.It provides an undergraduate-level evaluation setting.
  • Evaluation benchmarks: OlympiadBench contains 675 problems from international mathematics olympiads and physics contests.It evaluates performance on olympiad-style mathematics and physics problems.
  • Evaluation limitations: AIME 24 and AMC 23 are excluded because their 30 and 40 questions produce high evaluation variance and unreliable comparisons for smaller models.Repeated AIME 24 evaluation shows substantial accuracy fluctuations without a clear training-step trend.

E.1 Ablation Study on the Adaptive Difficulty Prediction Framework

Ablations show that trained adapters and calibration improve adaptive-difficulty prediction, while DOTS remains effective with smaller reference sets and across extended, step-based, and wall-clock evaluations. The combined results support robust performance and reduced training requirements relative to original GRPO.

  • Difficulty prediction ablation: Off-the-shelf Qwen2.5-Math-1.5B-Instruct embeddings fail to capture difficulty structure, while adapters and calibration yield consistently higher Pearson correlation.The comparison is reported across all tested settings.
  • Reference-set robustness: A reference set of 128 achieves RL performance comparable to a reference set of 256 for DOTS on DeepScaleR.The comparison covers Qwen2.5-Math-1.5B and Qwen2.5-3B.
  • Extended training: After extending training to 100 training steps, DOTS-based training continues to outperform original GRPO in two model–dataset settings.The settings are Qwen2.5-Math-1.5B with DeepScaleR and Qwen2.5-3B with DeepMath.
  • Evaluation views: DOTS and DOTS+RR maintain strong performance when results are plotted against either training steps or wall-clock time.Alternate plots provide the two evaluation views.
  • Efficiency mechanism: DOTS without RR incurs reference-rollout and difficulty-prediction overhead but often reaches the original GRPO final accuracy in fewer training steps.The reduced number of steps can offset the per-step overhead and reduce overall training time.
Loading 2506.05316v4…