Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Prioritized Replay for RL Post-training

Mehdi Fatemi

TL;DR

This work introduces a non-parametric, problem-level prioritized replay framework for RL post-training of reasoning-focused LLMs, deriving the priority as $\omega = p(1-p)$ from empirical success rates to favor intermediate-difficulty problems under GRPO. A lightweight binary max-heap scheduler selects high-priority problems, updates their statistics online, and reuses them in subsequent training steps, with mechanisms to prevent forgetting (solved/unsolved pools) and to encourage exploration. The approach is demonstrated as a proof-of-concept on a constrained training budget, showing improved early performance over uniform sampling on math reasoning benchmarks while emphasizing the need for careful hyperparameter tuning in larger-scale regimes. Overall, the method provides a principled, scalable alternative to handcrafted curricula that directly aligns data selection with the model's current learning dynamics, without extra predictors or task labels.

Abstract

We introduce a problem-level prioritization framework for RL post-training of large language models. Building on insights from prioritized replay in deep RL, as well as prior observations that rollouts with intermediate success rates tend to produce stronger learning signals under methods such as GRPO, our approach selects problems according to a simple, model-driven priority score derived from empirical success statistics. In contrast to conventional curriculum strategies that emphasize easier tasks early in training, the resulting schedule naturally focuses training on problems that are neither consistently solved nor consistently failed, while deprioritizing those that contribute little gradient information. The method yields a continuously adapting and automatic prioritization process that requires no predefined difficulty tiers, auxiliary predictors, or external labels. We further introduce lightweight mechanisms for practical deployment, including heap-based prioritized sampling and periodic retesting of solved and unsolved problems to mitigate starvation and forgetting. Overall, the approach offers a principled and scalable alternative to manually designed curricula while aligning data selection directly with the dynamics of GRPO-based post-training.

Prioritized Replay for RL Post-training

TL;DR

This work introduces a non-parametric, problem-level prioritized replay framework for RL post-training of reasoning-focused LLMs, deriving the priority as from empirical success rates to favor intermediate-difficulty problems under GRPO. A lightweight binary max-heap scheduler selects high-priority problems, updates their statistics online, and reuses them in subsequent training steps, with mechanisms to prevent forgetting (solved/unsolved pools) and to encourage exploration. The approach is demonstrated as a proof-of-concept on a constrained training budget, showing improved early performance over uniform sampling on math reasoning benchmarks while emphasizing the need for careful hyperparameter tuning in larger-scale regimes. Overall, the method provides a principled, scalable alternative to handcrafted curricula that directly aligns data selection with the model's current learning dynamics, without extra predictors or task labels.

Abstract

We introduce a problem-level prioritization framework for RL post-training of large language models. Building on insights from prioritized replay in deep RL, as well as prior observations that rollouts with intermediate success rates tend to produce stronger learning signals under methods such as GRPO, our approach selects problems according to a simple, model-driven priority score derived from empirical success statistics. In contrast to conventional curriculum strategies that emphasize easier tasks early in training, the resulting schedule naturally focuses training on problems that are neither consistently solved nor consistently failed, while deprioritizing those that contribute little gradient information. The method yields a continuously adapting and automatic prioritization process that requires no predefined difficulty tiers, auxiliary predictors, or external labels. We further introduce lightweight mechanisms for practical deployment, including heap-based prioritized sampling and periodic retesting of solved and unsolved problems to mitigate starvation and forgetting. Overall, the approach offers a principled and scalable alternative to manually designed curricula while aligning data selection directly with the dynamics of GRPO-based post-training.
Paper Structure (12 sections, 5 equations, 3 figures)

This paper contains 12 sections, 5 equations, 3 figures.

Table of Contents

  1. Introduction
  2. Core Method
  3. Computational Deployment
  4. Important Considerations
  5. Forgetting and Starvation
  6. Exploration
  7. Learning Rate
  8. Concise Reasoning
  9. Experiments
  10. Discussion
  11. Code and Reproducibility
  12. Related Work

Figures (3)

  • Figure 1: Success-rate distributions of the training data from the trainable subset of the GURU dataset cheng2025guru and a stratified 1,000-sample sub-subset, shown from the perspective of both models.
  • Figure 2: Evaluation on MATH500 and AIME 2024 datasets during training steps.
  • Figure 3: Dynamics of the prioritized sampling during the training process.