Work done by Weiqi Wang¹², Xin Liu²

¹ HKUST ² Amazon

📄 arXiv:2601.22448 💻 GitHub: horizon-rl/HeaPA

📄 Abstract

Reinforcement learning from verifiable rewards (RLVR) has emerged as a powerful paradigm for improving LLM reasoning, yet its efficiency is bottlenecked by rollout generation on a static, uniformly-sampled prompt pool. As training progresses, prompts that were once informative become either trivially solved or persistently unsolvable — both yielding near-zero gradient signal. We present HeaPA (Heap-based Prioritized sampling and on-policy Augmentation), a data curriculum framework that maintains a bounded, dynamically-evolving prompt pool via heap-based boundary sampling, on-policy pool expansion with lightweight asynchronous teacher verification, and topology-aware reward re-estimation along an augmentation lineage graph. Across GRPO and DAPO optimizers, two training corpora (DAPO-Math-17k and OpenR1-Math), and seven benchmarks, HeaPA delivers consistent +3–5 accuracy points while requiring 12–16% fewer FLOPs to reach baseline performance, at only ~2.1% additional wall-clock overhead on average — with gains that increase with model scale.

🔍 TL;DR

Problem

• RL is now the standard way to boost LLM reasoning (GRPO, DAPO, etc.). But current pipelines are data-inefficient:
  • They sample queries uniformly from a static dataset.
  • They never adapt the dataset as the model improves → the data becomes off-policy and misaligned with the model’s current ability.

Idea: HeaPA. Turn a static dataset into a living curriculum by:

1. Dual-heap query pool
   • Maintains a hard heap and an easy heap based on reward (difficulty proxy).
   • Always samples from the boundary region → medium-difficulty problems where the model sometimes succeeds and sometimes fails.
2. On-policy query augmentation
   • Use the current policy to generate new math problems by tweaking numbers in existing ones.
   • These new problems are calibrated to the model’s current capability.
3. Teacher-guided verification
   • A stronger “teacher” model checks whether augmented problems are solvable and provides ground-truth answers before they enter the pool.
4. Reward propagation via a lineage graph
   • Parents and children share information: rewards flow along the augmentation tree to keep difficulty estimates fresh.

Results (Qwen2.5-7B, math RL)

• Across DAPO-Math & OpenR1-Math, GRPO & DAPO:
  • +3–5 points average accuracy over the baseline.
  • On competition benchmarks:
    • AIME24 (GRPO, DAPO-Math): 21.4% vs 17.3%
    • AMC23 (GRPO, DAPO-Math): 82.4% vs 75.7%
• Compute efficiency:
  • Up to 12–16% less compute to reach the same performance vs original sampler.
  • Additional wall-clock overhead from HeaPA is only ~1–3%, and <1% for larger models.

Plug-and-play

• HeaPA works on top of GRPO and DAPO without changing the core RL algorithm.