HeaPA: Difficulty-Aware Heap Sampling and On-Policy Query Augmentation for LLM Reinforcement Learning
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , , , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2026
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866911408790175744 |
|---|---|
| author | Wang, Weiqi Liu, Xin Huang, Binxuan Cui, Hejie Zhang, Rongzhi Yu, Changlong Jin, Shuowei Yang, Jingfeng Yin, Qingyu Wang, Zhengyang Li, Zheng Gao, Yifan Nigam, Priyanka Yin, Bing Li, Lihong Song, Yangqiu |
| author_facet | Wang, Weiqi Liu, Xin Huang, Binxuan Cui, Hejie Zhang, Rongzhi Yu, Changlong Jin, Shuowei Yang, Jingfeng Yin, Qingyu Wang, Zhengyang Li, Zheng Gao, Yifan Nigam, Priyanka Yin, Bing Li, Lihong Song, Yangqiu |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_22448 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | HeaPA: Difficulty-Aware Heap Sampling and On-Policy Query Augmentation for LLM Reinforcement Learning Wang, Weiqi Liu, Xin Huang, Binxuan Cui, Hejie Zhang, Rongzhi Yu, Changlong Jin, Shuowei Yang, Jingfeng Yin, Qingyu Wang, Zhengyang Li, Zheng Gao, Yifan Nigam, Priyanka Yin, Bing Li, Lihong Song, Yangqiu Machine Learning Computation and Language 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. |
| title | HeaPA: Difficulty-Aware Heap Sampling and On-Policy Query Augmentation for LLM Reinforcement Learning |
| topic | Machine Learning Computation and Language |
| url | https://arxiv.org/abs/2601.22448 |