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
分类: cs.LG, cs.CL
发布日期: 2026-01-30
🔗 代码/项目: GITHUB
💡 一句话要点
提出HeaPA，通过堆采样和在线查询增强提升LLM强化学习效率。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 大型语言模型 强化学习 prompt采样 在线学习 堆数据结构
📋 核心要点
- 现有方法在LLM强化学习中，prompt采样效率低下，无法有效追踪模型能力边界，导致计算资源浪费。
- HeaPA通过堆采样聚焦能力边界附近的prompt，并利用在线策略增强动态扩展prompt池，提升采样效率。
- 实验表明，HeaPA在多个基准测试中，能够以更少的计算资源达到更高的准确率，尤其在大模型上效果更显著。
📝 摘要（中文）
本文提出了一种名为HeaPA（Heap Sampling and On-Policy Query Augmentation）的方法，旨在提高大型语言模型（LLM）在推理任务上的强化学习效率。HeaPA维护一个有界的、动态变化的prompt池，利用基于堆的边界采样来追踪能力边界，并通过轻量级的异步验证进行在线prompt增强来扩展prompt池。此外，HeaPA还通过拓扑感知的池统计重估和受控的重插入来稳定相关的查询。在两个训练语料库、两种训练方法和七个基准测试中，HeaPA始终提高了准确性，并以更少的计算量达到目标性能，同时保持了相当的实际运行时间。分析表明，这些增益来自于以能力边界为中心的采样和在线prompt池增长，并且随着模型规模的增加，收益变得更大。
🔬 方法详解
问题定义：现有基于RLVR（Reinforcement Learning with Verifiable Rewards）的LLM训练方法，在推理任务中面临prompt采样效率低下的问题。静态或与模型学习进度关联不紧密的prompt池，以及均匀采样策略，无法有效追踪模型不断变化的能力边界，导致大量rollout被浪费在模型已经解决或无法解决的prompt上。现有方法虽然尝试通过过滤、课程学习、自适应rollout分配或教师指导来提高效率，但通常假设prompt池是固定的，或者增加了额外的教师成本和延迟。
核心思路：HeaPA的核心思路是维护一个动态变化的prompt池，并采用一种能够有效追踪模型能力边界的采样策略。通过堆数据结构来维护prompt池，并利用堆的特性进行边界采样，从而将rollout资源集中在模型能力边界附近的prompt上。同时，利用在线策略增强来动态扩展prompt池，使其能够适应模型学习进度的变化。
技术框架：HeaPA的整体框架包含以下几个主要模块：1) Prompt池维护：使用有界堆数据结构来存储prompt，并根据prompt的难度进行排序。2) 堆采样：从堆的边界附近采样prompt，以聚焦于模型能力边界。3) 在线策略增强：利用当前策略生成新的prompt，并使用轻量级的异步验证来筛选高质量的prompt加入prompt池。4) 池统计重估：通过拓扑感知的池统计重估和受控的重插入来稳定相关的查询，避免prompt池中出现大量相似的prompt。
关键创新：HeaPA的关键创新在于：1) 堆采样：利用堆数据结构来维护prompt池，并进行边界采样，从而能够更有效地追踪模型的能力边界。2) 在线策略增强：利用在线策略生成新的prompt，并将其加入prompt池，从而能够动态地扩展prompt池，使其能够适应模型学习进度的变化。3) 拓扑感知的池统计重估：通过拓扑感知的池统计重估和受控的重插入来稳定相关的查询，避免prompt池中出现大量相似的prompt。
关键设计：HeaPA的关键设计包括：1) 堆的排序标准：prompt的难度，可以通过模型在prompt上的表现来估计。2) 边界采样的策略：可以采用不同的采样策略，例如均匀采样、重要性采样等。3) 在线策略增强的策略：可以采用不同的策略来生成新的prompt，例如基于模型的生成、基于规则的生成等。4) 异步验证的策略：可以采用不同的策略来验证新生成的prompt的质量，例如基于模型的验证、基于人工的验证等。
🖼️ 关键图片
📊 实验亮点
实验结果表明，HeaPA在两个训练语料库、两种训练方法和七个基准测试中，始终提高了准确性，并以更少的计算量达到目标性能，同时保持了相当的实际运行时间。尤其是在模型规模较大时，HeaPA的优势更加明显，证明了其在提升LLM强化学习效率方面的有效性。
🎯 应用场景
HeaPA可应用于各种需要利用LLM进行推理的强化学习任务，例如代码生成、数学问题求解、策略游戏等。通过提高训练效率，HeaPA能够降低训练成本，加速LLM在这些领域的应用。此外，HeaPA的在线prompt增强策略也为LLM的持续学习和能力提升提供了新的思路。
📄 摘要（原文）
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.