【AI学习】Huggingface复刻Test-time Compute Scaling技术

OpenAI ChatGPT o1 背后的关键技术Test-time Compute Scaling，Huggingface实现并开源了！

Hugging Face 团队发布了一篇关于“开源模型中的推理阶段计算扩展”（Test-time Compute Scaling）的研究文章。Hugging Face 团队通过复现 DeepMind 等研究成果，通过测试时计算扩展（test-time compute scaling）来提升小模型的性能，使其在某些情况下超越更大规模的模型。这种方法不依赖于更大的预训练预算，而是通过动态推理策略，让模型在更难的问题上“思考更长时间”。

. 关键技术：

计算最优扩展（compute-optimal scaling）：通过复刻DeepMind的技巧来提升测试时开放模型的数学能力。
多样性验证器树搜索 (DVTS)：提高了多样性并提供了更好的性能，特别是在测试时计算预算较大的情况下。
搜索和学习：一个轻量级工具包，用于使用LLM实现搜索策略，并使用vLLM实现速度提升。

实验结果：在MATH-500基准测试中，1B和3B的Llama Instruct模型在给定足够的“思考时间”后，性能超越了更大的8B、70B模型。

基于验证器的搜索方法

直接看下面的图，快速理解：
在这里插入图片描述
这张图示意了三种搜索方法：
1、Best-of-N：让LLM生成N个候选答案，用频率的方法，就是选择出现次数最多的答案。使用奖励模型，则是强调答案质量而非频率。通过奖励模型，为每个候选答案分配分数，然后选择奖励最高的答案。文章中提到的一种变体是：汇总所有相同响应的得分，通过对相同响应的回答加权求和选择总得分最高的回答。
2、Beam Search：一种探索解决方案空间的系统搜索方法，通常与过程奖励模型 (PRM) 结合使用，以优化问题解决中间步骤的采样和评估。与对最终答案产生单一分数的传统奖励模型不同，PRM 会提供一系列分数，其中推理过程的每个步骤都有一个分数。这种细粒度反馈能力使得 PRM 成为 LLM 搜索方法的自然选择。
3、多样性验证器树搜索 (DVTS)：HuggingFace 开发的Beam Search扩展，将初始beam拆分为独立的子树，然后使用 PRM 贪婪地扩展这些子树。这种方法提高了解决方案的多样性和整体性能，尤其是在测试时计算预算较大的情况下。

问题回答过程

在这里插入图片描述
问题回答过程：
1、首先给 LLM 提供一个数学问题，让其生成 N 个部分解，例如，推导过程中的中间步骤。
2、每个 step 都由 PRM 评分，PRM 估计每个步骤最终达到正确答案的概率。
3、一旦搜索策略结束，最终候选解决方案将由 PRM 排序以产生最终答案。

DVTS：通过多样性提升性能

为什么需要DVTS？
虽然总体上很明显，Beam Search是一种比 Best-of-N 或多数投票更好的搜索策略，但 DeepMind 的论文表明，每种策略都有权衡，这取决于问题的难度和测试时计算预算。
Beam Search在中等难度和困难难度问题中取得了持续的进展，但在较简单问题上，尤其是在计算预算较大的情况下，它的表现往往比 Best-of-N（甚至多数投票）更差。
通过观察Beam Search生成的结果树，HuggingFace 意识到，如果单个步骤被分配了高奖励，那么整棵树就在该轨迹上崩溃，从而影响多样性。这促使他们探索一种最大化多样性的集束搜索扩展。

为了解决这个问题，HuggingFace 开发了一个扩展，称之为「多样性验证器树搜索」（DVTS），旨在最大限度地提高 N 较大时的多样性。

DVTS 的工作方式与Beam Search类似，但有以下修改：

对于给定的 N 和 M，将初始集束扩展为 N/M 个独立子树。
对于每个子树，选择具有最高 PRM 分数的步骤。
从步骤 (2) 中选择的节点生成 M 个新步骤，并选择具有最高 PRM 分数的步骤。
重复步骤 (3)，直到达到 EOS token 或最大树深度。

下图是将 DVTS 应用于 Llama 1B 的结果：
在这里插入图片描述
试验中N的取值为 4、16、64、256
可以看到，DVTS 为集束搜索提供了一种补充策略：在 N 较小时，集束搜索更有效地找到正确的解决方案；但在 N 较大时，DVTS 候选的多样性开始发挥作用，可以获得更好的性能。

有了各种各样的搜索策略，一个自然的问题是哪一个是最好的？在 DeepMind 的论文中（可参考《Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters 》），提出了一种计算 - 最优扩展策略，该策略可以选择搜索方法和超参数 θ，以便在给定的计算预算 N 下达到最佳性能。

例如：对于较简单的问题和较低的计算预算，最好使用 Best-of-N 等策略，而对于较难的问题，Beam Search是更好的选择。下图为计算 - 最优曲线！

在这里插入图片描述

DeepMind 论文中的测试时搜索方法（test-time search methods）

DeepMind 论文：https://arxiv.org/pdf/2408.03314

在这里插入图片描述

文章中提供了三种不同的测试时搜索方法（test-time search methods）的比较，这些方法都是针对一个给定的问题（例如数学问题）如何在大型语言模型（LLM）中进行搜索以找到最佳答案的策略。这三种方法分别是：

√ Best-of-N：这种方法涉及从基础语言模型中独立地采样N个完整的解决方案，然后根据过程验证器（PRM）的最终答案判断来选择最佳答案。
√ Beam Search：这种方法通过在每一步中采样N个候选答案，并根据PRM预测的步骤奖励估计来评分，然后选择得分最高的M个步骤继续搜索。这个过程会重复进行，直到找到解决方案或达到最大搜索轮数。
√ Lookahead Search：这种方法修改了Beam Search评估单个步骤的方式。它使用预先查看（lookahead）模拟来改进PRM在搜索过程中每一步的价值估计的准确性。具体来说，在Beam Search的每一步中，不是使用当前步骤的PRM分数来选择顶部候选，而是进行模拟，向前滚动多达k步，如果达到解决方案的末尾则提前停止。然后使用这个模拟结束时的PRM预测来评分当前步骤。

图中的每个部分都用流程图的形式展示了这些方法的步骤，包括：