在人工智能领域,大型语言模型(LLMs)的规模通常被认为是其性能的关键因素。然而,随着模型规模的不断增大,计算资源的消耗也呈指数级增长。那么,是否存在一种方法,能够让规模较小的模型通过优化计算策略,在推理阶段超越那些规模庞大的模型呢?近日,来自上海AI LAB等研究机构的研究者,提出了计算最优的测试时扩展(Test-Time Scaling, TTS)策略,揭示了小模型在特定条件下超越大模型的潜力。本文将详细介绍这一研究的背景、方法、实验结果及其意义。
论文地址:
https://arxiv.org/abs/2502.06703
网站: Can 1B LLM Surpass 405B LLM? Rethinking Compute-Optimal Test-Time Scaling
1. 引言
大型语言模型(LLMs)在多个领域展现了显著的性能提升,尤其是在推理任务中。然而,随着模型规模的增大,计算资源的消耗也急剧增加。测试时扩展(TTS:Test-Time Scaling)是一种通过在推理阶段增加计算量来提升模型性能的方法。尽管TTS已被证明有效,但当前的研究并未系统分析策略模型、过程奖励模型(PRMs)和问题难度对TTS的影响。这种分析的缺乏限制了TTS方法的理解和实际应用。
本文聚焦于两个核心问题:(1)在不同策略模型、PRMs和问题难度下,如何最优地扩展测试时计算?(2)通过扩展计算,LLMs在复杂任务上的性能能提升到什么程度,小模型是否可以通过这种方法超越大模型?
通过在MATH-500和AIME24任务上的综合实验,本文得出以下观察:(1)计算最优的TTS策略高度依赖于策略模型、PRM和问题难度的选择。(2)通过计算最优的TTS策略,极小的策略模型可以超越大模型。
例如,1B的LLM在MATH-500上可以超越405B的LLM。此外,在MATH-500和AIME24上,0.5B的LLM超越了GPT-4o,3B的LLM超越了405B的LLM,7B的LLM超越了o1和DeepSeek-R1,同时具有更高的推理效率。这些发现表明,根据任务和模型的特定特征调整TTS策略的重要性,并表明TTS是提升LLMs推理能力的有前途的方法。
2. 设置与预备知识
2.1 问题公式化
本文将推理问题公式化为马尔可夫决策过程(MDP),定义为元组(S, A, P, R, γ),其中S是状态空间,A是动作空间,P是转移函数,R是奖励函数,γ是折扣因子。给定一个提示x ~ X,策略模型生成初始动作a1 ~ πθ(·|s1),其中s1 = x是初始状态。策略接收奖励R(s1, a1),状态转移到s2 = [s1, a1],其中[·,·]表示两个字符串的连接。这个过程持续到达到最大步数或生成
2.2 测试时扩展方法
本文考虑了三种TTS方法:Best-of-N(BoN)、束搜索(Beam Search)和多样化验证树搜索(DVTS)。BoN方法生成N个响应,然后应用评分和投票方法选择最终答案。束搜索方法在给定束宽度N和束大小M的情况下,首先生成N步,验证器选择前N/M步进行后续搜索。DVTS通过将搜索过程分为N/M个子树,每个子树独立使用束搜索进行探索。
2.3 计算最优的测试时扩展
为了最大化TTS的性能,Snell等人提出了计算最优的扩展策略,选择对应于给定测试时策略的超参数以最大化特定提示的性能收益。给定提示x,Target(θ, N, x)表示由策略模型θ在计算预算N下生成的输出分布。
3. 重新思考计算最优的测试时扩展
3.1 计算最优扩展策略应考虑奖励
计算最优的TTS旨在为每个问题分配最优的计算资源。以往的研究使用单一PRM作为验证器,但这种方法存在分布外(OOD)问题。本文提出将奖励函数R整合到计算最优的TTS策略中,确保计算最优的扩展适应策略模型、提示和奖励函数。
3.2 绝对问题难度标准比分位数更有效
Snell等人根据Pass@1准确率的分位数将问题分为五个难度级别。然而,本文发现使用MATH的难度级别或基于Pass@1准确率分位数的oracle标签并不有效,因为不同策略模型具有不同的推理能力。因此,本文使用绝对阈值来测量问题难度,将问题分为简单(50%~100%)、中等(10%~50%)和困难(0%~10%)三个级别。
4. 如何最优地扩展测试时计算?
4.1 实验设置
本文在MATH-500和AIME24数据集上进行实验,使用Llama 3和Qwen2.5系列的策略模型,以及多个开源的PRMs进行评价。评分方法包括PRM-Min、PRM-Last和PRM-Avg,投票方法包括多数投票、PRM-Max和PRM-Vote。
4.2 不同策略模型和PRMs对TTS的提升
PRMs在不同策略模型和任务上的泛化能力较差。实验结果表明,使用Skywork和Qwen2.5-Math PRMs的搜索方法在更大计算预算下性能显著提升,而使用Math-Shepherd和RLHFlow PRMs的搜索方法性能较差。BoN方法在使用Math-Shepherd和RLHFlow PRMs时表现最佳,而搜索方法在使用Skywork和Qwen2.5-Math PRMs时表现更好。
4.3 不同难度级别对TTS的提升
对于小策略模型,BoN方法在简单问题上表现更好,而束搜索在困难问题上表现更好。对于中等规模的策略模型,DVTS在简单和中等问题上表现良好,而束搜索在困难问题上表现更好。对于72B的策略模型,BoN在所有难度级别上表现最佳。
4.4 PRMs对响应长度和投票方法的敏感性
PRMs对步骤长度有偏见,不同PRMs的推理令牌数量差异显著。PRMs对投票方法敏感,Skywork-PRM-7B在使用PRM-Vote时表现更好,而Qwen2.5-Math-PRM-7B对投票方法不敏感。
5. 计算最优测试时扩展的结果
5.1 小策略模型能否超越大模型
通过计算最优的TTS策略,小策略模型可以超越大模型。例如,Llama-3.2-3B-Instruct在MATH-500和AIME24上超越了Llama-3.1-405B-Instruct。Qwen2.5-0.5B-Instruct和Llama-3.2-3B-Instruct超越了GPT-4o,
DeepSeek-R1-Distill-Qwen-1.5B超越了o1-preview和o1-mini,
DeepSeek-R1-Distill-Qwen-7B超越了o1和DeepSeek-R1。
5.2 计算最优TTS与CoT和多数投票的比较
计算最优的TTS比多数投票高效256倍,比CoT提升了154.6%的推理性能。随着策略模型参数的增加,TTS的提升逐渐减小。
5.3 TTS是否比长CoT方法更有效
TTS在MATH-500和AIME24上优于rStar-Math、Eurus-2、SimpleRL和Satori,但在AIME24上表现不如
DeepSeek-R1-Distill-Qwen-7B。TTS在简单任务上比复杂任务更有效。
6. 相关工作
本文回顾了LLM测试时扩展、提升数学推理能力和过程奖励模型的相关工作。LLM测试时扩展通过多数投票、搜索方法和细化等方法提升性能。提升数学推理能力的方法包括大规模数学语料库预训练和监督微调。过程奖励模型通过自动数据收集和高效MCTS等方法提升性能。
7. 结论与讨论
本文通过全面的实验分析,揭示了计算最优的TTS策略对策略模型、PRMs和问题难度的依赖性,验证了小模型在应用计算最优TTS时可以超越大模型。本文还展示了7B PRM通过监督更强大的72B策略模型实现强TTS结果的重要性,表明需要开发更高效的监督方法以提升小模型在复杂任务上的性能。
——完——
@北方的郎 · 专注模型与代码
喜欢的朋友,欢迎赞同、关注、分享三连 ^O^