arXiv:2505.18065cs.LG2025-05

提升推理模型泛化能力,让大模型用更少计算得更高准确率。

Reward Model Generalization for Compute-Aware Test-Time Reasoning

  • 基于概率框架分析奖励模型泛化误差对计算效率的影响
  • 降低奖励模型误差可减少寻找正确答案所需样本数
  • 提出动态搜索框架CATS,适合追求高效推理的场景

外部测试时推理通过解耦生成与选择来增强大语言模型。推理时,模型生成多条推理路径,由辅助的奖励模型(PRM)评分并选取最优路径。该设置的核心挑战是测试时计算最优性(TCO),即在固定推理预算下最大化答案准确率。本文建立理论框架,分析奖励模型泛化误差如何影响计算效率与推理性能。利用PAC-Bayes理论,推导出泛化界,表明更低的奖励模型泛化误差可减少找到正确答案所需的样本数量。受此启发,提出计算感知树搜索(CATS)——一种动作-评价框架,动态控制搜索行为:动作网络根据奖励分布和稀疏性统计输出采样超参数,评价网络估计其效用以指导预算分配。在MATH和AIME基准上,使用多种LLM与PRM的实验表明,CATS持续优于其他外部测试时搜索方法,验证了理论预测。

原文摘要 · Abstract (English)

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.

大模型推理奖励模型计算优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。