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

2026-01-30

AI Summary

This work addresses the inefficiency and computational waste in reinforcement learning for large language models caused by static or loosely associated prompt pools that fail to adapt to the dynamic frontier of model capabilities. To overcome this, the authors propose a continuously evolving bounded prompt pool mechanism that integrates difficulty-aware heap-structured sampling, lightweight asynchronous validation-driven online policy enhancement, topology-aware statistical re-estimation, and a controlled reinsertion strategy. This approach focuses sampling on capability-boundary examples and enables stable prompt pool expansion without additional teacher supervision. Evaluated across two corpora, two training protocols, and seven benchmarks, the method achieves higher accuracy with less computation and comparable wall-clock time, with performance gains amplifying as model scale increases.

Technology Category

• Natural Language Processing: Prompt Engineering / Prompting
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• LLM reinforcement learning
• prompt sampling
• reasoning tasks
• training efficiency
• on-policy learning

Innovation

• heap-based sampling
• on-policy augmentation
• difficulty-aware prompting
• dynamic prompt pool
• frontier-focused learning