[論文レビュー] HeaPA: Difficulty-Aware Heap Sampling and On-Policy Query Augmentation for LLM Reinforcement Learning
HeaPAは有界で進化するクエリプールを提案し、ヒープベースの境界サンプリングとオンポリシー増強を用いてLLM推論タスクにおけるRLVRの効率と精度を改善し、モデルとデータセット全体で一貫した利益を示す。
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.
研究の動機と目的
- 動的な能力フロンティアへプロンプトを整合させることで RLVR のサンプル効率を向上させることを動機づける。
- 難易度とフロンティアの進捗を追跡する有界な進化的クエリプールを開発する。
- プールを成長させるための非同期検証を伴うオンポリシーのクエリ増強を導入する。
- 幾何的に意義を考慮したプール統計の再推定と制御再挿入によって増強効果を安定化させる。
- 既存の RLVR レシピとバックボーンとのプラグアンドプレイ互換性を示す。
提案手法
- 各レコード統計 tilde{r}_{k} を用いて低難度(難易度低)と高難度(難易度高)のアイテムに分割された有界なデュアルヒーププロンプトプールを維持し、サンプリングをガイドする。
- 境界サンプリングを用いてフロンティア付近の中難易度帯を優先し、安定性のために易しいアイテムを混ぜることを選択可能とする。
- 現在のポリシーが候補プロンプトを生成し、教師と非同期に検証してペアをプールに挿入することでオンポリシー増強を実行する。
- 難易度補正統計をトポロジー的に伝播させる増強系系統図を構築し、相関する増強の中でもサンプリング優先度を安定化させる。
- 学習済みレコードをアーカイブし、制御挿入機構によって再利用してプールの進化とスループットを平滑化する。
- 基盤となるポリシー更新ルールを変更することなく GRPO および DAPO RLVR レシピに組み込む。
実験結果
リサーチクエスチョン
- RQ1フロンティア焦点のヒープベースサンプリングはLLMsのRLVRにおけるサンプル効率と最終精度を改善するか。
- RQ2非同期検証を伴うオンポリシー増強は、教師コストや待機遅延を過度に増大させることなくクエリプールを安全に成長させることができるか。
- RQ3トポロジー認識プール統計の再推定は、増強が相関していても学習を安定化させるか。
- RQ4HeaPAはモデルサイズと学習コーパス間でベースラインと比較してどのように拡張するか。
- RQ5HeaPAは既存のRLVRパイプライン（GRPO/DAPO）のドロップインコンポーネントとしてバックボーン間で互換性があるか。
主な発見
| Recipe / Variant | AIME24 | AIME25 | AMC23 | GPQA | MATH500 | MinervaMath | Olym.Bench | Avg. |
|---|---|---|---|---|---|---|---|---|
| GRPO (原典; Shao ら, 2024) | 17.3 | 17.0 | 75.7 | 39.2 | 80.6 | 46.7 | 49.2 | 46.5 |
| GRPO + 優先サンプリング (PS) | 16.7 | 16.5 | 76.8 | 38.8 | 81.6 | 47.8 | 51.2 | 47.1 |
| GRPO + ヒープサンプリング | 18.3 | 16.3 | 74.8 | 39.3 | 82.3 | 46.3 | 49.9 | 46.7 |
| GRPO + ポリシー増強 (PA) | 17.7 | 15.4 | 73.7 | 38.0 | 84.8 | 48.3 | 46.6 | 46.4 |
| GRPO + PS + PA | 19.9 | 18.4 | 76.3 | 40.0 | 85.3 | 49.5 | 50.3 | 48.5 |
| Reinforce-Ada (Xiong ら, 2025b) | 22.0 | 19.8 | 81.0 | 41.9 | 84.9 | 52.0 | 51.8 | 50.3 |
| HeaPA (GRPO) + ChildAgg | 16.4 | 22.4 | 77.8 | 41.8 | 85.2 | 45.0 | 52.2 | 48.7 |
| HeaPA (GRPO) + PathAgg | 21.4 | 21.0 | 82.4 | 42.7 | 85.6 | 51.4 | 52.9 | 51.1 |
| DAPO (原典; Yu ら, 2025) | 19.6 | 18.0 | 76.5 | 38.1 | 76.7 | 47.9 | 45.9 | 46.1 |
| DAPO + 優先サンプリング (PS) | 19.4 | 18.3 | 74.1 | 37.0 | 76.5 | 48.8 | 44.2 | 45.5 |
| DAPO + ヒープサンプリング | 19.0 | 16.5 | 75.3 | 36.9 | 75.1 | 46.3 | 44.2 | 44.8 |
| DAPO + ポリシー増強 (PA) | 19.0 | 16.4 | 76.5 | 38.9 | 75.9 | 47.5 | 45.3 | 45.6 |
| DAPO + PS + PA | 19.6 | 18.7 | 78.4 | 39.4 | 75.3 | 48.6 | 46.1 | 46.6 |
| HeaPA (DAPO) + ChildAgg | 19.4 | 20.4 | 83.4 | 40.9 | 76.9 | 48.6 | 46.6 | 48.0 |
| HeaPA (DAPO) + PathAgg | 23.3 | 21.7 | 80.8 | 42.0 | 78.0 | 50.6 | 48.8 | 49.3 |
| OpenR1-Math (Hugging Face, 2025) | ||||||||
| GRPO (原典; Shao ら, 2024) | 16.5 | 11.5 | 61.7 | 37.0 | 76.4 | 36.0 | 40.2 | 39.9 |
| GRPO + 優先サンプリング (PS) | 18.0 | 13.0 | 63.8 | 37.6 | 77.4 | 37.1 | 44.2 | 41.6 |
| GRPO + ヒープサンプリング | 18.1 | 12.0 | 64.3 | 36.1 | 76.7 | 35.8 | 45.5 | 41.2 |
| GRPO + ポリシー増強 (PA) | 16.9 | 11.9 | 63.7 | 36.8 | 77.6 | 35.6 | 44.6 | 41.0 |
| GRPO + PS + PA | 17.1 | 14.9 | 66.3 | 36.8 | 78.1 | 35.8 | 43.3 | 41.8 |
| Reinforce-Ada (Xiong ら, 2025b) | 19.8 | 16.4 | 67.5 | 38.0 | 81.6 | 37.2 | 46.0 | 43.8 |
| HeaPA (GRPO) + ChildAgg | 17.7 | 15.6 | 63.7 | 36.9 | 79.2 | 36.5 | 46.2 | 42.3 |
| HeaPA (GRPO) + PathAgg | 20.6 | 15.5 | 68.4 | 38.5 | 81.3 | 37.7 | 45.9 | 44.0 |
| DAPO (原典; Yu ら, 2025) | 19.5 | 15.4 | 65.4 | 37.8 | 80.6 | 51.1 | 45.1 | 45.0 |
| DAPO + 優先サンプリング (PS) | 20.3 | 16.7 | 64.0 | 37.7 | 80.4 | 53.0 | 43.4 | 45.1 |
| DAPO + ヒープサンプリング | 16.0 | 15.3 | 64.2 | 36.9 | 83.0 | 51.8 | 44.6 | 44.5 |
| DAPO + ポリシー増強 (PA) | 19.9 | 14.8 | 66.4 | 36.6 | 84.8 | 53.7 | 44.5 | 45.8 |
| DAPO + PS + PA | 20.5 | 16.1 | 68.3 | 38.1 | 82.2 | 54.8 | 45.3 | 46.5 |
| HeaPA (DAPO) + ChildAgg | 25.6 | 17.1 | 68.3 | 37.0 | 80.4 | 52.7 | 47.5 | 46.9 |
| HeaPA (DAPO) + PathAgg | 25.7 | 17.1 | 69.4 | 39.8 | 84.2 | 54.8 | 48.5 | 48.5 |
- HeaPAは2つの学習コーパス、2つの学習レシピ、7つのベンチマークで平均性能を一貫して改善する。
- DAPO-Math において、HeaPAはそれぞれ 51.1 (GRPO) および 49.3 (DAPO) に達し、ベースラインの 46.5 および 46.1 を上回る。
- OpenR1-Math において、HeaPAはそれぞれ 44.0 (GRPO) および 48.5 (DAPO) に達し、ベースラインの 39.9 および 45.0 を上回る。
- アブレーションは、 frontier-focused sampling とオンポリシーのプール成長を組み合わせると、いずれか一方の成分だけを用いるよりも大きな利益を生むことを示す。
- Topology-aware PathAgg は ChildAgg と同等かそれ以上の効果を示し、プール統計の安定化の恩恵を示唆する。
- HeaPA は 目標性能に到達するまでのトレーニング計算量 (PFLOPs) を削減し、非同期検証による1ステップあたりの実時間にわずか約 2.1% の追加コストしか発生させない。
より良い研究を、今すぐ始めましょう
論文設計から論文執筆まで、研究時間を劇的に削減しましょう。
クレジットカード登録不要
このレビューはAIが作成し、人間の編集者が確認しました。