Hierarchical Budget Policy Optimization for Adaptive Reasoning

著者 Shangke Lyu、Linjuan Wu、Yuchen Yan、Xingyu Wu、Hao Li他5名
所属 Zhejiang University、SF Technology
投稿日 2025年07月22日
カテゴリ cs.AI, cs.LG

Hierarchical Budget Policy Optimization for Adaptive Reasoning

基本情報

  • arXiv ID: 2507.15844v1 (https://arxiv.org/abs/2507.15844)
  • 著者: Shangke Lyu、Linjuan Wu、Yuchen Yan、Xingyu Wu、Hao Li他5名
  • 所属: Zhejiang University、SF Technology
  • 投稿日: 2025年07月22日
  • カテゴリ: cs.AI、cs.LG

簡単に説明すると

この論文は、大規模推論モデルが問題の複雑さに関係なく一様に長い思考連鎖を生成してしまう計算効率の問題を解決する手法を提案しています。「HBPO(Hierarchical Budget Policy Optimization)」という強化学習フレームワークで、問題の難易度に応じて推論の深さを自動的に調整します。最大60.6%のトークン削減と3.14%の精度向上を同時に達成しました。プロジェクトページ: https://zju-real.github.io/hbpo/ 、GitHub: https://github.com/zju-real/hbpo

1. 研究概要

1.1 背景と動機

大規模推論モデルは思考連鎖(Chain-of-Thought)手法により複雑な推論タスクで優れた性能を発揮しますが、簡単な問題でも不必要に長い推論連鎖を生成してしまう根本的な非効率性があります。例えば、基本的な算術問題に数千トークンを消費することもあります。

最近の研究では以下のことが明らかになっています。

  • 中間ステップなしでも競争力のある精度を維持できる場合がある
  • 簡単なタスクでは短い推論パスの方が同等またはより良い性能を示すことがある
  • 最適な推論長はタスクによって大きく異なる(GSM8Kで約1,100トークン、OlympiadBenchで3,000トークン以上)

既存の効率化手法は以下の2つのカテゴリに分類されます。

  1. 長さ制御手法:明示的なメカニズムで生成を直接制約
  2. 報酬ベース手法:訓練目的に長さペナルティを組み込む

しかし、これらの手法は効率を優先するあまり精度を犠牲にしています。モデルがいつ長い推論を必要とし、いつ短い推論で十分かを自律的に判断する機構も欠けています。

1.2 主要な貢献

  • 階層的バジェット最適化(HBPO)という新しい強化学習フレームワークの提案
  • 一様な効率制約が探索空間を崩壊させ推論能力を低下させることの実証
  • 問題特性に基づいて自動的に推論深度を調整する創発的な適応推論の実現
  • 数学的推論ベンチマークで最大60.6%のトークン削減と3.14%の精度向上を達成

2. 提案手法

2.1 手法の概要

HBPOは、探索空間を複数のバジェット制約されたサブグループに分割し、推論の多様性を保ちながら問題特性と必要な計算努力の自然な対応関係を発見できるようにする手法です。

具体的には、ロールアウトサンプルを異なるトークンバジェット(512、1024、2048、2560トークン)を持つサブグループに分割します。推論タグの後に長さプロンプト(例:「I will answer the question within n tokens」)を挿入することで実装します。

2.2 技術的詳細

階層的バジェット探索

  • クエリに対してn個のロールアウトサンプルを生成
  • k個のサブグループ{G1, G2, ..., Gk}に分割
  • 各サブグループGiはトークンバジェットbiに関連付けられる
  • バジェット値は昇順(b1 < b2 < ... < bk)で構成

バジェット対応報酬設計
区分的な報酬関数を使用。

R(n_gen | b) = 
  f1(n_gen, b)  if correct, n_gen > b, and n_gen ≤ L_max
  f2(b)         if correct, n_gen ≤ b, and n_gen ≤ L_max  
  0             otherwise

ここで。

  • f1(n_gen, b) = β × cos(πn_gen/2L_max) - α|n_gen - b|
  • f2(b) = β × cos(πb/2L_max)

アドバンテージ計算
階層構造を活用して2つの補完的なコンポーネントに分解。

  1. サブグループ内アドバンテージ:バジェット期待値に対する相対的パフォーマンスを測定
  2. サブグループ間アドバンテージ:異なるバジェット間での比較学習を可能にする

2.3 新規性

  • 探索空間の階層的分割により、効率訓練で一般的な探索空間崩壊を防ぐ
  • バジェットレベル全体でのパフォーマンス比較により、グローバル最適化を使わず各問題に適した計算量を発見
  • 古典的報酬とコサイン型報酬の利点を組み合わせた区分的な報酬関数
  • 明示的な複雑性ラベルや外部ガイダンスなしに計算リソースを問題要件に合わせる学習

3. 実験結果

3.1 実験設定

データセットとモデル

  • 訓練:DeepScaleRデータセット(AIME、AMC、Omni-Math、STILLから40K問題)
  • ベースモデル:DeepSeek-R1-Distill-Qwen-1.5BとDeepScaleR-Preview-1.5B

実装詳細

  • VeRLフレームワークを使用
  • コンテキストウィンドウ:4,096トークン
  • 学習率:10^-6、バッチサイズ:64
  • クエリあたり16ロールアウト、4サブグループに均等分割
  • バジェット制約:{512, 1024, 2048, 2560}トークン

評価プロトコル

  • 4つの数学的推論ベンチマーク:GSM8K、Math500、OlympiadBench、AIME25
  • 2つの評価設定:自然推論と効率プロンプティング

3.2 主要な結果

自然推論の条件下での性能

  • DeepSeek-R1-Distill-Qwen-1.5Bベース:平均精度56.3%→59.4%、トークン使用量60.6%削減(7,921→3,120)
  • DeepScaleRベース:精度63.7%を維持しつつトークン50.2%削減(4,744→2,364)
  • AIME25で31.1%の精度達成(ベースラインおよび全効率手法を上回る)

効率プロンプティング設定での性能

  • ベースラインモデルは強制的にトークンを最小化すると10%以上の精度低下
  • HBPOは堅牢な性能を維持(DeepScaleRで59.4%精度、947トークン)
  • L1-Max(1024)と同等の精度を32%少ないトークンで達成

3.3 既存手法との比較

適応的行動の創発

  • L1-Max:問題の難易度に関わらず一様なトークン使用(MATH500で3,260、AIME25で3,163)
  • HBPO:問題の複雑さに応じたトークン使用(MATH500で1,818、AIME25で3,988)
  • 2.2倍の変動が問題の複雑さと直接相関

階層構造の分析

  • 単一バジェット訓練:59.8%精度(一様探索の限界)
  • デュアルバジェット:61.7%精度(1.9%改善)
  • 4バジェット構成:63.7%精度(最適性能)
  • 6-8バジェット:性能がやや低下(サブグループあたりのサンプル数減少)

4. 実用性評価

4.1 実装の容易性

  • 標準的なGRPOフレームワークの拡張として実装可能
  • バジェット固有のプロンプトを推論タグ後に挿入するだけの簡単な実装
  • コードとモデルが公開されており再現性が高い

4.2 計算効率

  • 訓練は1エポック(629ステップ)で完了
  • 推論時は自動的にトークン使用量を調整
  • 簡単な問題では数百トークン、複雑な問題では数千トークンを使用

4.3 応用可能性

  • 数学的な推論以外のタスクへの拡張可能性
  • より大規模なモデルへの適用可能性
  • 計算リソースが制限された環境での実用性

5. まとめと所感

5.1 論文の意義

この研究は、推論効率と能力が本質的に対立するものではなく、適切に構造化された階層的訓練により同時に最適化できることを示しました。特に重要なのは、一様な制約や離散的なモード選択ではなく、問題の複雑さに基づいて自動的に推論深度を調整する創発的な適応行動を実現した点です。

推論パターンの分析では、HBPOが問題の難易度に応じて異なる推論戦略を開発することが示されました。GSM8Kでは思考内容の割合が81%、AIME25では89%と単調増加し、反省キーワード(wait、alternatively、but等)も6回から30回へと増加しています。

5.2 今後の展望

  • より大規模なモデルでの検証と性能評価
  • 数学以外のドメイン(コーディング、論理推論等)への適用
  • バジェット階層の自動調整メカニズムの開発
  • リアルタイム推論コスト最適化への応用
  • 人間の問題解決プロセスとの類似性の研究