Hugging Face 复现 o1 策略:Llama 3B 模型通过扩展测试时计算超越 80B 模型
OpenAI 的 o1 系列模型展示了推理能力的新高度,其核心在于扩展了测试时的计算资源。尽管 OpenAI 未公开具体实现细节,但 Hugging Face 团队基于 DeepMind 的研究成果,成功逆向工程并开源了相关的扩展测试时计算方法。实验表明,将这一方法应用于轻量级模型(如 Llama 1B/3B),其数学推理性能可超越参数量大得多的模型,甚至达到计算机科学博士生的平均水平。
背景:Scaling Law 的新范式
传统的 Scaling Laws(缩放定律)主要关注模型参数规模、数据量与计算预算之间的关系。然而,o1 及其后续研究揭示了一种新的范式:在推理阶段增加计算量(Test-time Compute)可以显著提升模型表现,而无需增加训练时的参数规模。
这种"扩展测试时计算"的方法允许模型在生成答案的过程中进行自我反思、搜索验证或迭代优化。Hugging Face 团队在此基础上进行了改进,提出了更适用于开源生态的实现方案。
核心策略:多样化验证器树搜索
目前扩展测试时计算主要有两种策略:自我优化和搜索。
1. 自我优化 (Self-Optimization)
模型识别并纠正后续迭代中的错误来优化输出。虽然对某些任务有效,但这通常要求模型具备内置的自我优化机制,限制了其在通用场景下的适用性。
2. 搜索方法 (Search Methods)
侧重于生成多个候选答案并使用验证器选择最佳答案。这种方法更灵活,能适应不同难度的问题。Hugging Face 的研究主要聚焦于搜索方法,特别是结合可学习的验证器(Verifier)。
研究中重点介绍了三种搜索策略:
Best-of-N
为每个问题生成 N 个响应,使用奖励模型为每个候选答案分配分数,选择分数最高的答案。该方法强调答案质量而非频率。分为原版(Vanilla)和加权版(Weighted)。加权版汇总所有结果相同的答案,选择总分数最高的,在算力预算充足时表现更佳。
Beam Search
一种系统探索解决方案空间的方法,通常与过程奖励模型(PRM)结合。传统奖励模型仅在最终答案上产生单个分数,而 PRM 为推理过程的每个中间步骤分配分数,提供精细反馈。这使 PRM 非常适合大模型的推理优化。
多样化验证器树搜索 (DVTS)
这是新开发的 Beam Search 变体。它将初始 Beam 拆分为独立的子树,然后使用 PRM 做贪婪扩展。具体流程如下:
- 对于给定的 Beam 宽度 M 和生成数量 N,初始 Beam 集设定为 N/M 个独立子树。
- 对于每个子树,选择 PRM 分数最高的步骤。
- 生成 M 个新的下一步,继续选择分数最高的。
- 重复此过程,直到生成 EOS token 或达到最大深度。
这种方法显著提高了解决方案的多样性,特别是在测试时算力预算较大的情况下。
实验设置
为了验证上述策略的有效性,研究使用了以下配置:
- 语言模型:Llama-3.2-1B-Instruct。选择轻量级模型是为了快速迭代,且其在数学基准测试中尚未饱和,便于观察提升效果。
- 流程奖励模型 (PRM):Llama3.1-8B-PRM-Deepseek-Data。与语言模型同属一个系列,且在测试中给出了更好的结果。
- 数据集:MATH-500。由 OpenAI 发布的数学问题子集,横跨 7 个科目,对人类和大多数模型都具有挑战性。
实验结果与分析
动态分配策略最优
实验对比了多种搜索策略的表现:
- 多数投票策略:相比贪婪解码基线有显著改进,收益在 N=64 后趋于稳定。限制在于难以解决需要细致推理的问题,或当多个答案都错向同一方向时。
- Best-of-N:加入奖励模型后表现提高。加权版始终优于原版,确保了频率较低但质量较高的答案也能被选中。
- Beam Search:让 1B 模型表现开始高于 8B 模型。但在简单问题上不如 Best-of-N。研究发现,如果一个中间步骤获得高分,整个树可能坍塌到这一步,影响多样性。
- DVTS:改进了答案多样性。在 N 较大时增强了对简单/中等难度问题的性能。而 Beam Search 在 N 较小时表现最好。
最终,基于问题难度动态分配策略的方法取得了最佳成绩。这意味着根据题目复杂度自动调整搜索深度和广度是未来的关键方向。
性能突破
- 在 Llama 1B 上应用该方法,数学分数超过 8 倍大的模型,也超过了计算机科学博士生的平均分数(40%)。
- 在 Llama 3B 上,进步幅度更大,甚至能和 20 几倍大的 70B 模型媲美。
未来展望与挑战
团队指出,未来这项技术还有更多值得探索的地方:
- 更强的验证器:提高验证器的稳健性和泛化能力至关重要。最终目标是实现自我验证,但目前实践中仍难以实现,需要更细致的策略。
- 结构化推理:在生成过程中加入明确的中间步骤或'想法',通过将结构化推理整合到搜索过程中,可以在复杂任务中获得更好的性能。
- 合成数据:搜索方法可用于合成数据,创建高质量的训练数据集,从而进一步微调模型。
- 开源贡献:开放的流程奖励模型目前数量较少,是开源社区可以做出重大贡献的领域。
- 领域扩展:目前的方法在数学和代码等可验证领域表现出色。如何将这些技术扩展到结构性较差或评判标准主观的任务,仍是一个重大挑战。
成本与部署考量
值得注意的是,评论区有观点指出,这种方法更适合本地部署,而不是 API 调用。因为调用 256 次 3B 模型和过程奖励模型,通常会比调用一次 70B 模型更贵。这表明在实际应用中,需要在推理精度和计算成本之间寻找平衡点。
此外,也有建议提出在 Qwen 系列模型上尝试类似方法,以及指路天工 Skywork 发布了两个基于 Qwen 的 PRM 模型,显示了该技术在多模型生态中的潜力。
开源资源
相关开源代码已发布在 GitHub 上,供社区参考和复现:
https://github.com/huggingface/search-and-learn
通过复现和扩展这些技术,开源社区有望缩小与大厂闭源模型之间的推理能力差距,推动 AI 技术的普惠化发展。