Every Rollout Counts: Optimal Resource Allocation for Efficient Test-Time Scaling

📄 arXiv: 2506.15707v2 📥 PDF

作者: Xinglin Wang, Yiwei Li, Shaoxiong Feng, Peiwen Yuan, Yueqi Zhang, Jiayi Shi, Chuyi Tan, Boyuan Pan, Yao Hu, Kan Li

分类: cs.LG, cs.AI

发布日期: 2025-05-30 (更新: 2025-10-20)

备注: Accepted at NeurIPS2025


💡 一句话要点

提出DORA:通过优化资源分配,提升大语言模型在测试时推理的效率和准确率

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 测试时扩展 资源分配 数学推理 方向导向 推理优化 最优策略

📋 核心要点

  1. 现有测试时扩展方法在分配推理预算时效率低下,未能充分利用计算资源。
  2. DORA通过在方向层面进行资源分配,解决了现有方法偏袒候选答案多的推理方向的问题。
  3. 实验表明,DORA在数学推理任务上显著优于现有方法,实现了更高的准确率。

📝 摘要(中文)

本文研究了测试时扩展(TTS)中如何更有效地分配有限的推理预算,以提升大型语言模型(LLM)的性能。作者将测试时搜索建模为一个资源分配问题,并推导出在固定预算下最大化正确解概率的最优分配策略。研究揭示了现有搜索方法的一个核心局限:解级别的分配倾向于偏袒具有更多候选答案的推理方向,导致计算资源的次优利用。为此,作者提出了方向导向资源分配(DORA),这是一种可证明最优的方法,通过将方向质量与候选数量解耦,并在方向级别分配资源来缓解这种偏差。在MATH500、AIME2024和AIME2025等数学推理基准测试中,实验结果表明DORA始终优于强大的基线方法,并在相当的计算成本下实现了最先进的准确率。

🔬 方法详解

问题定义:论文旨在解决测试时扩展(TTS)中,如何在有限的计算资源(rollout budget)下,最大化大型语言模型(LLM)找到正确答案的概率。现有方法,如解级别的资源分配,存在一个根本问题:它们倾向于将更多资源分配给那些拥有更多候选答案的推理方向,即使这些方向的质量可能不高。这种偏见导致计算资源的浪费和次优的性能。

核心思路:论文的核心思路是将测试时搜索过程视为一个资源分配问题,并从理论上推导出最优的资源分配策略。关键在于将资源分配的粒度从解级别提升到方向级别,从而避免因候选答案数量而产生的偏见。通过解耦方向的质量和候选答案的数量,DORA能够更有效地利用计算资源,提升找到正确答案的概率。

技术框架:DORA方法主要包含以下几个阶段:1) 定义搜索空间,即所有可能的推理方向和候选答案;2) 评估每个推理方向的质量,这可能涉及到使用LLM对方向进行打分或排序;3) 根据评估结果和剩余的计算预算,使用最优资源分配策略来决定每个方向应该分配多少资源;4) 执行推理,并根据分配的资源探索每个方向的候选答案;5) 最终选择最有可能正确的答案。

关键创新:DORA最关键的创新在于其方向导向的资源分配策略。与传统的解级别分配不同,DORA将资源分配的重点放在推理方向的质量上,而不是候选答案的数量。这种策略能够更有效地利用计算资源,避免因候选答案数量而产生的偏见,从而提升找到正确答案的概率。

关键设计:DORA的关键设计包括:1) 如何准确评估推理方向的质量,这可能涉及到设计合适的提示工程或使用LLM进行打分;2) 如何推导出最优的资源分配策略,这可能涉及到使用优化算法或动态规划;3) 如何有效地执行推理,并根据分配的资源探索每个方向的候选答案。论文中具体的技术细节(如损失函数、网络结构等)取决于具体的LLM和任务,但核心思想是方向导向的资源分配。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DORA在MATH500、AIME2024和AIME2025等具有挑战性的数学推理基准测试中取得了显著的性能提升。实验结果表明,DORA在与现有方法相当的计算成本下,能够始终优于强大的基线方法,并达到最先进的准确率。这证明了DORA在提升LLM推理效率和准确率方面的有效性。

🎯 应用场景

DORA方法可以广泛应用于需要高精度和可靠性的LLM应用场景,例如:数学推理、代码生成、金融分析、医疗诊断等。通过更有效地利用计算资源,DORA可以提升LLM在这些领域的性能,并降低推理成本。未来,DORA可以与其他测试时优化技术相结合,进一步提升LLM的性能和效率。

📄 摘要(原文)

Test-Time Scaling (TTS) improves the performance of Large Language Models (LLMs) by using additional inference-time computation to explore multiple reasoning paths through search. Yet how to allocate a fixed rollout budget most effectively during search remains underexplored, often resulting in inefficient use of compute at test time. To bridge this gap, we formulate test-time search as a resource allocation problem and derive the optimal allocation strategy that maximizes the probability of obtaining a correct solution under a fixed rollout budget. Within this formulation, we reveal a core limitation of existing search methods: solution-level allocation tends to favor reasoning directions with more candidates, leading to theoretically suboptimal and inefficient use of compute. To address this, we propose Direction-Oriented Resource Allocation (DORA), a provably optimal method that mitigates this bias by decoupling direction quality from candidate count and allocating resources at the direction level. To demonstrate DORA's effectiveness, we conduct extensive experiments on challenging mathematical reasoning benchmarks including MATH500, AIME2024, and AIME2025. The empirical results show that DORA consistently outperforms strong baselines with comparable computational cost, achieving state-of-the-art accuracy. We hope our findings contribute to a broader understanding of optimal TTS for LLMs.