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

AI summary

Overview

Research area: Machine learning / large language model post-training, specifically reinforcement learning with verifiable rewards (RLVR) for reasoning tasks.

Technical level: Intermediate to Advanced. The paper assumes familiarity with group-based RL objectives (PPO-style clipped surrogates, GRPO, DAPO), advantage estimation with within-prompt baselines, and rollout-based training loops.

Scope (one sentence): The paper proposes HeaPA, a query-lifecycle framework that maintains a bounded, evolving prompt pool, samples near the moving "capability frontier" with a dual-heap boundary sampler, and grows the pool through on-policy query augmentation with asynchronous teacher verification, evaluated across two training corpora, two RL recipes, seven math/reasoning benchmarks, and multiple model scales.

What This Paper Is About

When training LLMs on reasoning tasks with reinforcement learning, the cost is dominated by generating rollouts, so the key efficiency question is which prompts get sampled and when. Static or loosely coupled prompt pools miss the policy's moving capability frontier: batches end up full of either very hard prompts (mostly wrong or overlong rollouts) or very easy prompts (mostly correct rollouts), both of which give little learning signal for group-based RLVR updates. HeaPA's goal is to keep a bounded but evolving pool centered on medium-difficulty prompts, to grow that pool using the policy itself plus lightweight teacher annotation of verifiable answers, and to prevent correlated new queries from destabilizing the sampling curriculum.

Key Contributions

1. A bounded, evolving prompt pool with dual-heap boundary sampling. HeaPA splits scored queries into a low partition (harder) and a high partition (easier), defines a boundary between them, and samples candidates that alternate between the "hardest easy" and "easiest hard" items, forming up to 2B candidates from which the required batch items are drawn uniformly. Unscored items sit in a cold queue and are prioritized until their first verified rollout group defines their pool statistic.

2. On-policy query augmentation with asynchronous verification. The current policy generates n_aug candidate augmented queries conditioned on sampled prompts; these are sent to a teacher that annotates a verifiable final answer, while the RL loop continues in parallel. Candidates marked unsolvable, lacking a parsable numeric answer, or violating dataset filters are discarded. Accepted pairs (policy-generated prompt, teacher-annotated answer) enter the cold queue.

3. Lineage-aware pool-statistic re-estimation over an augmentation graph. Each augmented query inherits a parent pointer, forming a directed parent-to-children lineage graph. Pool statistics (not per-rollout training rewards) are periodically refreshed bottom-up, with a policy-estimated relative difficulty factor d_c clamped to [d_min, d_max] (default 1.0 for seeds) attenuating the influence of children estimated to be harder, and a leaf-count-weighted or unweighted aggregation (PathAgg or ChildAgg) blended half-and-half with the parent's previous score.

4. Controlled reinsertion and recycling. Trained records move to an archive; recycling is triggered when the active pool cannot supply a full batch or the archive reaches a size threshold, with refreshed scores and batched reinsertion to amortize overhead and avoid bursts of correlated items dominating the boundary band.

Main Findings

• Consistent average gains across corpora and recipes (Table 1). With Qwen2.5-7B-Instruct, HeaPA (GRPO) + PathAgg reaches 51.1 Avg@16 on the DAPO-Math corpus versus 46.5 for original GRPO (+4.6), and HeaPA (DAPO) + PathAgg reaches 49.3 versus 46.1 for original DAPO (+3.2). On OpenR1-Math, HeaPA (GRPO) + PathAgg reaches 44.0 versus 39.9 (+4.1), and HeaPA (DAPO) + PathAgg reaches 48.5 versus 45.0 (+3.5).

• Sampling or augmentation alone is not sufficient. On DAPO-Math with GRPO, the combined PS+PA baseline reaches 48.5, below HeaPA's 51.1; on OpenR1-Math with DAPO, PS+PA reaches 46.5 versus HeaPA's 48.5. Individual components (prioritized sampling, heap sampling, policy augmentation) produce smaller and less consistent gains.

• Reinforce-Ada is strong but remains below HeaPA on average. Reinforce-Ada reaches 50.3 on DAPO-Math and 43.8 on OpenR1-Math under GRPO, occasionally winning on one or two individual benchmarks.

• A PCL-style fixed-pool curriculum is competitive but not better. Using the updated policy as an online difficulty estimator over a fixed pool reaches 50.3 Avg@16 on DAPO-Math and 43.5 on OpenR1-Math under GRPO, versus HeaPA's 51.1 and 44.0, suggesting the evolving-pool mechanism and lineage-aware reinsertion add benefit beyond fixed-pool curriculum selection.

• Topology-aware aggregation (PathAgg) beats ChildAgg. On DAPO-Math with GRPO, PathAgg improves over ChildAgg by +2.4 average points (51.1 vs. 48.7); on OpenR1-Math with GRPO by +1.7 (44.0 vs. 42.3).

• Transfers across backbone family and scales (Table 2). On Qwen3 backbones, HeaPA improves Avg@16 over GRPO by +1.8 (0.6B), +2.0 (1.7B), +7.2 (4B), and +4.9 (8B), and also exceeds GRPO with prioritized sampling at every scale. The largest gains appear at mid-to-large scales (4B/8B).

• Sampling distribution shifts toward the medium-reward band. With a frozen Qwen2.5-7B-Instruct defining a fixed reward landscape, prioritized sampling concentrates probability mass on the extreme low-reward tail, whereas heap-based boundary sampling shifts mass toward a medium-reward band.

• Query-source dynamics. Augmented queries tend to start with lower reward than seed queries but their reward rises sharply over training and with model scale, with a corresponding shift in the advantage-related signal, indicating augmented queries increasingly form non-saturated rollout groups.

• Efficiency. HeaPA reduces PFLOPs required to first reach the baseline GRPO target performance across all four Qwen3 scales. Averaged across scales, per-step wall-clock time increases by only 2.1%; asynchronous teacher annotation takes about 13.1% of main training time when counted as separate auxiliary service time. Accepted augmented queries have a 1.2% estimated wrong-label rate after filtering, and synthetic label-noise degradation is reported as gradual.

• Formatting details. Evaluation uses Avg@16 (mean correctness over 16 rollouts); correctness requires exact match of the parsed answer to the ground truth, with reward r = 1 for correct, 0 for incorrect, and −1 for invalid rollouts. Dataset statistics for the benchmarks are given in Table 8 of the paper, which is not included in the provided content.

Methodology in Plain English

The researchers treat prompt selection as a supply-chain problem rather than a fixed-list filtering problem. They keep a bounded pool of prompt–answer pairs. Brand-new or freshly added prompts are "cold" (no difficulty estimate yet) and get sampled early so the system can learn how hard they are. Once a prompt has produced a verified group of rollouts, its average shaped reward becomes its pool statistic, and it is placed into one of two heaps: the harder half or the easier half.

Each training step, the sampler fills the batch first with cold items, then with items nearest the split between the two heaps, deliberately targeting a medium-difficulty band rather than the extremes. Rollouts are generated and scored, advantages are computed with the standard within-prompt baseline, and the underlying optimizer (GRPO or DAPO) updates the policy unchanged.

In parallel, the policy generates candidate new queries based on prompts it just trained on; these candidates go to a teacher asynchronously, which supplies a verifiable final answer. Verified candidates enter the pool as cold items, so the pool grows along with the policy's current frontier. Because augmented queries often share templates or are local edits of one another, the team records lineage (parent–child links) and periodically recomputes pool statistics bottom-up through that graph, attenuating children that the policy flagged as harder and blending the aggregate with the parent's prior value. Records that have been trained are moved to an archive and reinserted in small controlled batches when needed, which smooths the curriculum and amortizes overhead.

Why This Matters

Impact on research. HeaPA reframes RLVR efficiency as a query-lifecycle management problem—pool growth, frontier tracking, and stable integration—rather than a one-shot filtering or curriculum problem over a fixed corpus. It shows that a bounded, evolving pool can be coupled to policy progress without dense teacher supervision or token-level distillation, and it provides a diagnostic (pool-statistic distribution over reward) showing where a sampler puts its probability mass.

Real-world applications:

• Training reasoning-capable assistants where verifiable answers exist (mathematics, code execution, structured extraction), reducing rollout compute to reach a target quality level.
• Any setting with expensive inference-time data generation, where deciding which prompts to spend rollouts on dominates the bill.
• Curriculum systems that must keep pace with a model that improves during training, rather than a precomputed easy-to-hard schedule.
• Data-creation pipelines: using an on-policy generator plus lightweight answer annotation to convert an initial seed corpus into a larger, frontier-aligned training set.

Industry relevance. Because rollout generation dominates cost in RLVR, reducing PFLOPs-to-target and adding only 2.1% per-step wall-clock overhead directly translates into training cost and turnaround. The asynchronous teacher design keeps annotation off the critical path, and the framework is presented as plug-and-play across optimizers (GRPO, DAPO) and backbones (Qwen2.5-7B and Qwen3 at 0.6B/1.7B/4B/8B), which matters for teams that swap base models or recipes. Training code is released at https://github.com/horizon-llm/HeaPA.

Future Directions

• Extending beyond verifiable-answer domains. The current design relies on a deterministic verifier and a teacher that can annotate a parsable numeric answer; how the lifecycle would work for open-ended tasks with no exact-match reward is not addressed.
• Reducing or removing teacher dependence. Teacher annotation costs about 13.1% of main training time as auxiliary service time and the paper reports a 1.2% estimated wrong-label rate after filtering; how performance changes with weaker teachers, larger label-noise rates, or entirely teacher-free answer-preserving augmentation (the SvS-style direction the paper mentions as complementary) remains an open question.
• Tuning the pool and boundary hyperparameters. The target fraction of scored items in the low partition (α), pool capacity N, the number of augmentations n_aug, the difficulty clamp range [d_min, d_max], and the optional mixing of very easy items are all design choices; the paper does not report a systematic sensitivity study of these in the provided content.
• Scaling analysis at larger sizes. Gains are reported as more pronounced at mid-to-large scales; whether the trend continues at scales beyond 8B, and how pool capacity should scale with model size, is left open.

Target Audience

Researchers and engineers working on LLM post-training and reinforcement learning with verifiable rewards, especially those optimizing rollout cost and prompt selection. Practitioners running GRPO- or DAPO-style pipelines at scale will benefit from the pool-management and asynchronous-verification design. Academics studying data selection, curriculum learning, self-play augmentation, and distillation-based training will find the PCL-style and SvS-style comparisons and the lineage-aware re-estimation mechanism most relevant. Readers need comfort with group-based advantage estimation and clipped surrogate objectives to follow the preliminaries; the framework sections themselves are described algorithmically and are accessible to an engineer implementing a training loop.

Authors’ abstract

RLVR has become a standard recipe for training LLMs on reasoning tasks with verifiable outcomes, but when rollout generation dominates the cost, efficiency hinges on which prompts are sampled and when. In practice, prompt pools are often static or only weakly coupled to policy progress, so uniform sampling fails to track the moving capability frontier and wastes rollouts on regions that are already solved or still unreachable. Prior methods improve efficiency via filtering, curricula, adaptive rollout allocation, or teacher guidance, but they often assume a fixed pool, which does not support stable on-policy pool growth, or they introduce additional teacher cost and latency. In this work, we propose HeaPA (Heap Sampling and On-Policy Query Augmentation), which maintains a bounded, evolving pool, tracks the frontier with heap-based boundary sampling, grows the pool via on-policy augmentation under lightweight asynchronous validation, and stabilizes correlated queries via topology-aware pool statistics re-estimation and controlled reinsertion. Across two training corpora, two training recipes, and seven benchmarks, HeaPA consistently improves accuracy and reaches target performance with fewer computations at comparable wall-clock time. Analyses attribute the gains to frontier-focused sampling and on-policy pool growth, with more pronounced improvements at mid-to-large model scales. Our training code is publicly available at https://github.com/horizon-llm/HeaPA.