Reward Model Generalization for Compute-Aware Test-Time Reasoning

要約

外部のテスト時間推論は、生成と選択を切り離すことにより、大規模な言語モデル（LLM）を強化します。
推論時に、モデルは複数の推論パスを生成し、補助プロセス報酬モデル（PRM）を使用して、最高のスコアと選択を選択します。
この設定の中心的な課題は、テスト時間計算最適性（TCO）です。つまり、固定推論予算の下で回答の精度を最大化する方法です。
この作業では、PRMの一般化エラーが計算効率と推論パフォーマンスにどのように影響するかを分析するための理論的枠組みを確立します。
PAC-Bayes理論を活用すると、一般化の境界を導き出し、PRMの一般化エラーが低いと、正解を見つけるのに必要なサンプルが少なくなることが示されます。
この分析に動機付けられているため、検索動作を動的に制御する俳優criticフレームワークであるCompute-Aware Tree Search（CATS）を提案します。
アクターは、報酬分布とスパース統計に基づいてサンプリングハイパーパラメーターを出力しますが、批評家は予算の割り当てを導くためにその有用性を推定します。
さまざまなLLMSおよびPRMSを使用した数学とAIMEのベンチマークに関する実験は、CATが他の外部TTSメソッドを常に上回ることを示しており、理論的予測を検証しています。

要約(オリジナル)

External test-time reasoning enhances large language models (LLMs) by decoupling generation and selection. At inference time, the model generates multiple reasoning paths, and an auxiliary process reward model (PRM) is used to score and select the best one. A central challenge in this setting is test-time compute optimality (TCO), i.e., how to maximize answer accuracy under a fixed inference budget. In this work, we establish a theoretical framework to analyze how the generalization error of the PRM affects compute efficiency and reasoning performance. Leveraging PAC-Bayes theory, we derive generalization bounds and show that a lower generalization error of PRM leads to fewer samples required to find correct answers. Motivated by this analysis, we propose Compute-Aware Tree Search (CATS), an actor-critic framework that dynamically controls search behavior. The actor outputs sampling hyperparameters based on reward distributions and sparsity statistics, while the critic estimates their utility to guide budget allocation. Experiments on the MATH and AIME benchmarks with various LLMs and PRMs demonstrate that CATS consistently outperforms other external TTS methods, validating our theoretical predictions.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google