要約

大規模言語モデルを強化学習により適応させる最近の進展は、複雑な推論問題の解決において顕著な性能向上を達成したが、その一方で、オンポリシーのロールアウトコストが高額であり、多様な推論経路の探索が限定的であるという課題を抱えていた。本研究では、系列生成を木構造の探索プロセスとして捉える自己指導型ロールアウトアルゴリズム「TreePO」を提案する。TreePOは、動的木サンプリングポリシーと固定長セグメントデコードから構成され、局所的な不確実性を活用して追加の分岐を生成する。共通の接頭辞に対する計算を共有し、価値が低い経路を早期に pruning することで、更新1回あたりの計算負荷を著しく低減しつつ、探索の多様性を維持または向上させる。主な貢献は以下の通りである：(1) 連続するセグメントを通じてKVキャッシュの負荷を軽減し、早期停止機構を併用して新たな分岐を生成する、セグメント単位のサンプリングアルゴリズム；(2) 全体的および局所的な近接ポリシー最適化を考慮した、木構造に基づくセグメントレベルのアドバンテージ推定；(3) 確率と品質に基づく動的分岐およびフォールバック戦略の有効性に関する分析。実証実験により、TreePOが複数の推論ベンチマーク上で性能向上を示すとともに、訓練済みモデルのサンプリング設計においてGPU時間の22％～43％の削減を実現した。また、既存モデルに対しては経路レベルで最大40％、トークンレベルで最大35％のサンプリング計算負荷削減が確認された。推論効率の「無料のランチ」を提供する一方で、TreePOは、より少ないサンプル数と計算量でRLベースの後学習をスケーリングする実用的な道筋を示している。本研究の公式ページは https://m-a-p.ai/TreePO にて公開されている。