Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

HeaPA: Difficulty-Aware Heap Sampling and On-Policy Query Augmentation for LLM Reinforcement Learning

Weiqi Wang, Xin Liu, Binxuan Huang, Hejie Cui, Rongzhi Zhang, Changlong Yu, Shuowei Jin, Jingfeng Yang, Qingyu Yin, Zhengyang Wang, Zheng Li, Yifan Gao, Priyanka Nigam, Bing Yin, Lihong Li, Yangqiu Song

TL;DR

HeaPA addresses RLVR inefficiency from moving capability frontiers by merging frontier-aware heap sampling with on-policy query augmentation. The framework maintains a bounded, evolving prompt pool and grows it through asynchronously verified, policy-generated queries, while lineage-aware statistics stabilize sampling. Key innovations include the dual-heap pool with boundary sampling (PathAgg over ChildAgg for stable frontier tracking), asynchronous verification, and controlled reinsertion to prevent correlated insertions from destabilizing curricula. Across two math corpora, multiple backbones, and seven benchmarks, HeaPA improves accuracy and reduces rollout compute, with larger gains as model scale increases, enabling more efficient, scalable RLVR training.

Abstract

RLVR is now a standard way to train LLMs on reasoning tasks with verifiable outcomes, but when rollout generation dominates the cost, efficiency depends heavily on which prompts you sample and when. In practice, prompt pools are often static or only loosely tied to the model's learning progress, so uniform sampling can't keep up with the shifting capability frontier and ends up wasting rollouts on prompts that are already solved or still out of reach. Existing approaches improve efficiency through filtering, curricula, adaptive rollout allocation, or teacher guidance, but they typically assume a fixed pool-which makes it hard to support stable on-policy pool growth-or they add extra teacher cost and latency. We introduce HeaPA (Heap Sampling and On-Policy Query Augmentation), which maintains a bounded, evolving pool, tracks the frontier using heap-based boundary sampling, expands the pool via on-policy augmentation with lightweight asynchronous validation, and stabilizes correlated queries through topology-aware re-estimation of pool statistics and controlled reinsertion. Across two training corpora, two training recipes, and seven benchmarks, HeaPA consistently improves accuracy and reaches target performance with fewer computations while keeping wall-clock time comparable. Our analyses suggest these gains come from frontier-focused sampling and on-policy pool growth, with the benefits becoming larger as model scale increases. Our code is available at https://github.com/horizon-rl/HeaPA.

HeaPA: Difficulty-Aware Heap Sampling and On-Policy Query Augmentation for LLM Reinforcement Learning

TL;DR

HeaPA addresses RLVR inefficiency from moving capability frontiers by merging frontier-aware heap sampling with on-policy query augmentation. The framework maintains a bounded, evolving prompt pool and grows it through asynchronously verified, policy-generated queries, while lineage-aware statistics stabilize sampling. Key innovations include the dual-heap pool with boundary sampling (PathAgg over ChildAgg for stable frontier tracking), asynchronous verification, and controlled reinsertion to prevent correlated insertions from destabilizing curricula. Across two math corpora, multiple backbones, and seven benchmarks, HeaPA improves accuracy and reduces rollout compute, with larger gains as model scale increases, enabling more efficient, scalable RLVR training.

Abstract

RLVR is now a standard way to train LLMs on reasoning tasks with verifiable outcomes, but when rollout generation dominates the cost, efficiency depends heavily on which prompts you sample and when. In practice, prompt pools are often static or only loosely tied to the model's learning progress, so uniform sampling can't keep up with the shifting capability frontier and ends up wasting rollouts on prompts that are already solved or still out of reach. Existing approaches improve efficiency through filtering, curricula, adaptive rollout allocation, or teacher guidance, but they typically assume a fixed pool-which makes it hard to support stable on-policy pool growth-or they add extra teacher cost and latency. We introduce HeaPA (Heap Sampling and On-Policy Query Augmentation), which maintains a bounded, evolving pool, tracks the frontier using heap-based boundary sampling, expands the pool via on-policy augmentation with lightweight asynchronous validation, and stabilizes correlated queries through topology-aware re-estimation of pool statistics and controlled reinsertion. Across two training corpora, two training recipes, and seven benchmarks, HeaPA consistently improves accuracy and reaches target performance with fewer computations while keeping wall-clock time comparable. Our analyses suggest these gains come from frontier-focused sampling and on-policy pool growth, with the benefits becoming larger as model scale increases. Our code is available at https://github.com/horizon-rl/HeaPA.
Paper Structure (38 sections, 2 theorems, 18 equations, 6 figures, 4 tables)

This paper contains 38 sections, 2 theorems, 18 equations, 6 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Data Selection and Sampling in LLM Post-Training
  4. Teacher-Augmented Training and Distillation
  5. Reinforcement Learning Recipes
  6. Preliminaries
  7. The HeaPA Framework
  8. Bounded Dual-Heap Prompt Pool
  9. RLVR Update with Group-Based Advantages
  10. On-Policy Query Augmentation with Asynchronous Verification
  11. Lineage-Aware Pool Statistics Re-estimation
  12. Controlled Reinsertion and Recycling
  13. Main Evaluations of HeaPA
  14. Experiment Setup
  15. Evaluation Results
  16. ...and 23 more sections

Key Result

Proposition 3.1

Assume the lineage structure under a node $q$ is a rooted tree. For each leaf $\ell$ under $q$, define its (difficulty-adjusted) leaf signal as $z_\ell \in \mathbb{R}$ (e.g., $z_\ell = \phi(r_\ell,d_\ell)$ or a path-product variant). Let $\mathcal{L}(q)$ be the set of leaves under $q$ and define the Partition leaves by the first edge from $q$: for each child $c\in\mathcal{C}(q)$, let $\mathcal{L}(

Figures (6)

  • Figure 1: Overview of HeaPA: a bounded prompt pool with a cold queue and difficulty-based heaps enables boundary sampling, while on-policy augmentation with asynchronous teacher verification grows the pool and controlled reinsertion recycles archived records.
  • Figure 2: Comparison of seed and augmented query statistics across model scales and training steps.
  • Figure 3: Comparison of probability mass over reward for two sampling strategies. Rewards are computed by averaging rollouts with a frozen Qwen2.5-7B-Instruct model.
  • Figure 4: Training compute (PFLOPs) required to first reach the baseline GRPO target performance (final Avg@16 of the baseline run) for Qwen3 backbones at four scales. Lower is better.
  • Figure 5: Visualization of reward-axis partitioning under multi-heap prompt pools. Each row shows the contiguous reward intervals assigned to bins (heaps) for a given $H$ (we linearly rescale the pool statistic $\tilde{r}\in[-1,1]$ to $[-1,1]$ for visualization, where larger values indicate easier prompts). The numbers inside each segment denote the corresponding interval endpoints. Labels of the form "$k$H @ $-1$" indicate that $k$ bins are anchored and concentrated near the hardest end of the reward axis (starting at $-1$), while "$k$H @ $1$" indicates bins anchored near the easiest end (ending at $1$). As $H$ increases, more boundaries fall into the extreme-hard region, which increases the chance that boundary-based sampling selects too-difficult prompts.
  • ...and 1 more figures

Theorems & Definitions (4)

  • Proposition 3.1: Unbiased estimation of the uniform-descendant objective
  • proof
  • Proposition 3.2: Minimum-variance weighting under a simple noise model
  • proof