前段日子OpenAI推出的o1模型，以其提升显著的逻辑推理能力，引发了人们对它背后训练方法的热烈讨论。关于o1的介绍和输出结果demo，这里就不再赘述，大家可以去openai的官网上阅读（很短，读起来很快，因为秘密都藏好了）。我相信最近的一段时间里，当大家在网上探索o1是如何训练时，肯定会看到以下几个热点词：
Test/Inference-Time scaling law，通过增加推理阶段的算力提升模型的推理能力
Post Training，通过后训练提升模型的推理能力
强化学习、self-play（自我博弈）与MCTS（蒙特卡洛搜索树算法
）
等等。
当这些词单个出现在我们面前时，我们似乎很难把他们串在一起。不仅如此，我们也不知道单个词背后的原理，比如“什么是test/inference-time scaling law"？什么叫把算力花在推理阶段？为什么把算力花在推理阶段就有更好的结果？它和post training又是什么关系？诸如此类，令人很难在脑海里想象出完整的流程图。
在我对o1的探索期间，
我参考了这个github仓库（https://github.com/hijkzzz/Awesome-LLM-Strawberry）
，里面收集了内外网可能和o1背后的算法有关的各种资料，包括相关论文、代码和twitter post等等。其中不乏有o1核心contributor在近几年的研究成果。
在阅读这些资料后，我发现他们可以被拆成2点：
框架类研究成果
：可以理解成可能是o1训练算法背后的广义上的框架。它引入了一些诸如Test/Inference-Time Scaling law的基本思想，并给出了一些简单通用的实践方案。
细节变种类研究结果
：可以理解成是套在框架的大壳里，或者受到框架的启发，在实操细节上做的算法变种（但发表时间上却不一定晚于框架）。总结起来就是，各家有各家的做法。
而本篇文章，就是侧重对框架类研究成果进行细致的解读和展开，
我筛选出了2篇我认为可以被视作框架的文章
，分别是：
Let's Verify Step by Step （OpenAI）
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, Karl Cobbe
Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters (Google Deepmind)
Charlie Snell, Jaehoon Lee, Kelvin Xu, Aviral Kumar
本文将以第二篇文章为基础，同时穿插第一篇文章的核心内容进行讲解。后续还会单独再开一篇lets verify step by step的解读。而在更后面的文章中，将会挑选出我认为比较有意思的细节变种类研究结果，例如可能和o1具有某些相似性的算法，配合源码（只要他们有）进行深入分析。
最后，这两篇文章，都属于大致读来没有难度，但是细节上又较难把握的类型，所以我尽量按照自己的经验在细节上做了解读，力求还原训练的方方面面，同时重拆了原始论文，尽量从一个【框架】的角度更好做阐述
。另外在阅读本文前，建议大家先看openai o1技术报告中对于其输出结果的demo，更好了解什么是“模型的思考过程”
一、什么是Test/Inference-time Scaling Law
设想一下，当我们手里有一个基础模型（我们称其为generator），但是这个模型的逻辑推理能力（比如解数学题的能力）较差时，我们该怎么改进它？
再说的具体点，不考虑数据集相关的成本，假设我手头的gpu算力（FLOPs）是有限的，我该怎么利用它，能让我的模型最终能推理出更好的结果？
一个比较直接的想法是：把算力花在它的pretain阶段，给模型注入更多数理逻辑的预训练知识
。例如用更好、更多的代码数学等数据，或者扩展模型的参数规模。这个做法启发自大家都很熟悉的scaling law（更具体地说是pretrain-time scaling law）。
但是，
当我们研读openai o1的技术报告时，我们会发现，它把这个算力更多地用在了2个地方：
用在了rlhf的训练上（post training）
用在了模型的推理阶段上（Test/Inferece）
详情可以参见o1报告中绘制的这幅图：
曾经，为了提升模型的逻辑推理能力，我们把算力都花在pretrain阶段，由此诞生了pretrain scaling law
现在，已经有现成的产品证明，算力如果花在post training和inference上，模型的推理能力将得到更大提升，也就是存在一个Test/Inferece scaling law。
正如pretrain scaling law受到模型参数和训练数据的影响一样，Test/Inferece scaling law也必然受某些因素影响，而这些因素是什么，又是怎么影响的，就是本文要探索的内容。
不过等等，此时你肯定想问：
问题1：一般来说，一个模型的效果是由它的训练阶段决定的，所以如果这里说通过pretrain或者post training来提升模型的推理能力，我都能理解。
但是inference阶段是怎么提升模型的推理能力的？你说的把算力用在inference阶段到底是什么意思？
问题2：
post training和inferece是两种独立的提升模型推理能力的方法吗？它们可以结合在一起使用吗？
下面我们循序渐进地来回答这2个问题。
把算力用在inference阶段，也就是说，在不动我们已经训练好的模型的情况下，只通过优化推理方法，来提升模型最后的生成效果。这里又分成两种情况。
1.1 优化推理输入：prompt
这个方法大家应该非常熟悉了。例如，原来你的模型吃一个问题，直接吐给你回答。但是现在为了让模型能更好模拟人类的思考方式，你希望【模型在步步思考后再给出回答，也就是模型的生成结果里包含思考步骤+答案】，那么你可以选择在prompt中给模型相应的例子，或者在多轮对话中引导模型think step by step，来实现这个目标。相关论文可以参见：
Chain-of-Thought Prompting Elicits Reasoning in Large Language Models
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, Denny Zhou
你的prompt给的越细节，你的多轮引导给的越多，模型或许就能产出更好的结果。而更多的token意味着推理阶段需要花费更多的算力，所以这就是我们所说的【把算力花在推理阶段上可以提升模型效果】的具体内容之一。
1.2 优化推理输出：revise output distribution
可是，优化推理输入的方法还是不够直接。难道对于每一个问题，我需要精心设计prompt，或者手动诱导模型think step by step才行。
所以能不能让模型吃下一个问题后，自动化地去做CoT的过程呢？
也就是说，现在我们希望模型在吃下一个问题后，能自主产生以下输出：
attempt1 -> attempt2 -> attempt3 -> ...-> attempti -> answer
其中，每个attempt包含“多个中间步骤steps + 最终答案”，也就是它在模拟人类的思考过程：先做一次attempt，然后发现问题，在此基础上在做别的attempt，直到找到最终答案。
那么我要怎么让模型做到这点呢，一个直观的方法就是，如果我有：
problem -> attempt1 -> ... -> attempti -> answer
这种带标签的数据，那我不就能直接训练了？训练的方法也有很多，例如：
解法1：我直接做sft，把最正确的attempt放在输入序列最后，当作label进行训练即可
解法2：我用类似rlhf的方法，先有一个奖励模型，它能对每一个思考步骤做评估，然后利用这个评估结果，指引模型步步搜索，每一步都找到最佳的思考步骤，最后不就能找到答案了？
这两种解法，
仅从训练方法上来说，都可以算成是post-training
，也就是我们通过把算力花在post-training上来提升模型的逻辑推理能力。
可是，本文的标题不是【把算力花在inference上】吗？inference在哪里呢？我们再重新端详这2种解法
：
假设我们使用解法1或者解法2 post training好了模型，现在我们拿它做推理。模型吃一个问题，产出一系列中间结果和答案，
但是你能保证，这些中间结果和答案一定是最好的吗？
而此时，如果我们能有一个可以用来评估中间步骤好坏的verifier（比如解法2中的奖励模型），那么我们是不是就能在在使用这些post-training过的模型做推理时，更好指引模型步步产出更好的结果
？例如，我们对一个问题采样多个attempts链，从中找最好的。或者在单个attempts中找到最好的attempt，诸如此类。
或者说，假设我们在post-training阶段，使用这个verifier来指导模型
自动化
产生高质量的训练数据（这是个inference步骤），基于这些数据我们再做对齐。那么或许我们就能直接信任post-training的结果了。
所以，【优化推理输出】这一部分，你可以把算力全部花在post-training上，也可以花在post-training+inference上，从o1的技术报告上看，它应该选择了后者，同时post-training选择了某种基于强化学习的方法（
其实o1在pretrain阶段应该也有变动，具体的分析我们在后文中会通过实验数据给出猜想
）。至此，我们就把问题1和问题2都回答清楚了。
在理解了这些的基础上，我们就大概知道
【框架】
长什么样，具体而言：
首先，我们需要引导模型从“只产生结果”变成“同时产生中间步骤和结果”，这个过程称为post-training，这个过程里，你既可以使用基于强化学习的技术，也可以只做sft微调。你既可以只关注模型是否遵循了格式（即只关注是否产出了中间结果，而不关注中间结果的质量），也可以同时关注格式+中间结果质量
其次，我们需要训练一个能够评估中间结果的verifier，这其实也算post-training的一部分。这个verifier既可以被用在基于强化学习的post-training中（但它不一定是唯一的价值评估模型），也可以被用在第3步的inference阶段中引导搜索。
接着，我们需要设计一种搜索方法。它能够根据verifier返回的中间结果评估分数，在推理阶段更好指引模型做步步搜索，以达到思考的每一步都能最优化。（如果在post-training阶段，你只关注格式，那么inference阶段应用这种搜索方法就能更好指引搜索结果；如果你已经关注了格式+质量，这种方法也依然能达到更上一层楼的效果）。同时注意，如果在post-training阶段我们已经使用了基于强化学习的方案，那么这种搜索方法同样也能用在模型生产“经验数据”的阶段，这样可以
在自产自消的基础上，动化筛选高质量的数据集再做对齐
（避开人类标注，不是必须，但如果做得好的话，inference可能只需要在post-training阶段使用，目的是筛选高质量数据，而不需要在post-training后再做使用了）。
有了对框架的理解，接下来我们就来看deepmind基于这个框架做的两种方案：
方法一：利用PRM（Process-supervised Reward Model）指引搜索
post-training：只引导模型对齐格式 + 训练基于过程的verifier（PRM）
方法二：通过sft直接改变模型的输出分布（Revise proposal distribution）
post-training：通过sft方式引导对齐格式 + 保障中间结果基本质量 + 训练基于过程的verifier（PRM），这里直接用了的方案一中的PRM。
其中，方案一着重探讨训练PRM的一般方法和搜索过程的设计。方案二则给出sft型post training的一个例子。如果只看原始论文，我们很容易把这两种方案理解成Test/Inference scaling law的两方面。但是在我们总结出【框架】的基础上，就可以发现其实它们讲的是同一件事情，只是描述的侧重点不一样。
【注⚠️⚠️：本节内容打破了原始论文的逻辑结构，含有大量笔者的主观理解，大家选择性阅读】
二、方法一：利用PRM指引搜索
这种方法的主要目的，是通过训练一个
能够评估过程数据的奖励模型（Process-supervised Reward Model，PRM。与之相对应的是基于结果做评估的ORM，outcome-supervised Reward Model）
，来引导模型在推理阶段更好搜索出最佳答案。具体步骤如下：
格式训练
：先对模型做sft，使其能够产出带有过程数据的结果
训练PRM
：训练一个能够评估过程的verifier，我们称其为PRM
使用PRM指导搜索过程
：利用训练好的PRM，引导模型搜索出最佳答案。
我们来细看这三点。
2.1 格式训练
对于一个不做任何处理的基础模型，我们喂给它一个问题时，它一般会直接把答案吐给我们。
现在，我们希望引导模型在给出答案前，花更多的时间进行“思考”，也就是说，我们希望模型按照“思考步骤 + 回答”的这种格式，返回给我们response。
所以在这里我们需要先对模型做格式finetune。具体的方法是：
自生产数据
：在prompt中添加格式例子，引导模型按我们想要的格式产出结果。例如，我们可以按照一个思考步骤（step）一行的方式要求模型（正如lets verify step by step这篇论文说的一样）。我们把这些结果添加进sft数据集
sft
：使用这批自动化构建的sft数据集，微调模型，使其在之后回答问题时，都能按照“思考步骤+回答”的方式生成数据。
注意，这个sft过程仅仅是侧重于格式微调，并不关注思考步骤的质量
。思考步骤的质量是我们在后续过程中要考虑的事情。
2.2 训练PRM
现在，我们的模型已经能在生成结果里产出“思考步骤”数据了。我们需要训练一个能够评估这些steps的奖励模型，也就是
PRM（Process-supervised Reward Model）
。
有“supervised”，必然需要带有label的数据，也就是对于每个step，我们需要给它一个真值评分
。那么根据你是【超级有钱人】【有钱人】【一般有钱人】（能训得起来的都不算穷），我们有不同的构造方法。
（1）【超级有钱人】
直接调用格式微调后的模型，喂它吃一波问题，产出一波“steps + 回答”的数据（数据量级庞大）
让人工对steps打label（例如positive/negative/neural）。
这种方法除了昂贵和耗时，似乎没有别的缺点。
（2）【有钱人】
比起粗暴地将一堆数据打包给人工做标注，我们可不可以通过某种方式做细节筛选，只送那些我们认为有价值的数据给人工做标注呢？
直接调用格式微调后的模型，喂它吃一波问题，产出一波“steps + 回答”的数据（数据量级不大）
过滤掉一些含有“无效回答”的数据。例如，无法被解析的回答（无法解析的latex公式等等）。注意，“无效!=错误”
让人工对steps打label（例如positive/negative/neural）
决定如何使用PRM结果对“steps +回答”做一个整体打分，例如：
连乘式（prod）
：PRM会对每个step进行打分（离散标签下，这个得分表示概率）。我们将所有steps的得分相乘，用于表示整体的分数
最小式（min）
：取所有steps中最小的得分作为整体得分。这是因为从直觉上，steps有一步出错，整个逻辑链就很可能出问题。所以我们取最坏的情况作为整体得分。
最后一步式（last step）：当你训练PRM时，你肯定不是把单个step喂给PRM让它去做打分，而是把全部steps一起喂给它，让他们拥有上下文关系（PRM需要根据前面的steps去判断当前step的得分）
。所以理论上，你取last step的得分，也可以反映出整体得分。
prod和min都是openAI在《lets verify step by step》中探索过的方法，last step则是deepmind在本文中使用的方法。这些方法间没有绝对好坏之分，取决于你是如何设计整个训练过程的（我们在后文还会再来看这点）。
我们在这里只强调“需要制定出一种规则，能够将PRM的单步打分映射成整体打分。”
再次调用模型，产出一波数据。调用当前版本的PRM，利用上述所说的整体打分规则，对一个问题的“steps”整体进行打分。
我们特别筛选出convincing wrong answer数据（整体得分很高，但是最后的答案却是错误的），只把这些数据送给人工做标注
。之所以选择这种带有bias的数据，是因为它们对当前PRM具有迷惑性。
总结起来，这种方法也需要人工做标注，但通过对数据集的筛选，把钱花在了刀刃上。
（3）【一般有钱人】（重点关注）
在这种情况下，我们完全不需要人工标注数据，也就是说
每个step价值的真值label也是我们通过某种自动化的方法估计出来的
，那么具体怎么做呢？
直接调用格式微调后的模型，喂它吃一波问题，产出一波“steps + 回答”的数据（数据量级看需要）
更具体来说，对于每个问题，我们采样N个samples
（在deepmind的设置中N=16，一个sample就是一个“steps + 回答”）
过滤掉一些含有“无效回答”的数据。例如，无法被解析的回答（无法解析的latex公式等等）。注意，“无效!=错误”
我们通过N Monte-Carlo rollouts来预估每个step的价值
。deepmind的文章中没有给出具体的操作方法，这里笔者根据经验猜一下可能的实现思路：
假设某个sample为"step1 -> step2 -> step3 -> answer”，现在我们想分别估计step1~step3的价值
以step1为例，我们以它为起点（起点的意思是，问题 + step1），接下来继续采样N个samples。它模拟了N种“从step1出发到产出回答”的方案。
我们计算能够得出正确answer的方案比例。这个比例就作为step1的预估价值label。
由于这个label只是标志“从当前step出发，能找到正确答案的可能”，并不直接表示当前step是postive/negative/neural，所以这个label又被称为“soft label”。
对于step2和step3，我们也是同理估算其价值
有了这个被我们自动化模拟出来的step soft label，我们就可以用它来训练PRM模型了
。在deepmind的设置中，PRM依然是一个二分类模型，只是计算loss时真值标签非离散。
2.3 使用PRM指导搜索过程
到这一步为止，我们已经有以下模型：
一个能按照格式，产出中间思考步骤的模型（generator）
，但中间思考步骤质量得不到保证
一个能对中间思考步骤进行评估的奖励模型PRM（verifier）
。再次强调，当我们把数据喂给PRM做评估时，我们不是只喂单个step，而是喂完整的“问题 + steps ”数据，因为逻辑思考步骤间存在上下文关系。
而现在我们想做的事情是：
如何在不对generator继续做任何训练的情况下，使用verfier，来引导generator搜索出最佳的“steps + answer”？
也就是说，我们不想把额外的算力（FLOPs）花费在generator本身的training上了。
我们希望把算力花费在推理阶段，用更长更复杂的推理流程（表现在generator上就是“思考”的时间变长），让genertaor找到最佳答案。
所以接下来，我们就来详细介绍和分析3种常用的搜索方案
。下图给出了这3种方案的概览。接下来，我们会详细介绍这3种搜索方案的实现细节，并分析它们的效果。
这里再额外提一句，当PRM训练好之后，我们同样需要制定一个规则函数，将PRM的单步打分映射成整体得分
，也就是对一个问题下的某个“steps + answer”有个整体评分。在这个基础上我们再来做搜索。详情前文中已提过，大家可以翻回去阅读（
Deepmind使用的是last step打分方法，所以接下来的介绍我们都默认PRM的整体打分方式是这种
）。
（1）Best-of-N
Best-of-N方案因为简单直观，所以其搜索出的答案的准确率经常被当作baseline，和其他方案做比较
。Best-of-N的具体方案如下：
对于一个问题，我们采样N个samples（单个sample = steps + answer）
调用PRM，对这N个samples的中间步骤进行整体打分
。常用的打分方法有prod式、min式、last step式（前文提过）
选取整体得分最高的那组steps + answer作为输出
还有一种稍微改良的、加权版Best-of-N weighted方式
（Deepmind给出的介绍比较模糊，所以下面介绍中包含了笔者的一些猜想）：
对于一个问题，我们采样N个samples（单个sample = steps + answer）
检查这N个samples的answer。假设这些samples一共给出了x，y，z这3类答案，每个答案下分别有a，b，c个answers（N = a+b+c）
那么对于答案是x的sample，它的最终得分 = (a/N) * PRM给出的整体得分。a/N就被称为权重。
(2）Beam Search
对于一个问题，我们先并行采样N个step1（例如上图中N=4）
将每个“问题+step1”送入PRM，获取step1的得分。
对这些被筛选出来的step1继续做生成，每个step1下再产出N个step2，重复上面步骤（
评估step2的分数时要传入问题 + step1 + step2
），以此类推，直到达到设定的搜索停止条件为止
（3）Lookahead Search
对于一个问题，我们先并行采样N个step1（例如上图中N=4）
此时，我们不急着直接用PRM对step1们打分。
对每个step1，我们让它继续往下生成K个step
。然后将最后一个steps们送进PRM进行打分，并筛选出分数最高的top M个结果，返回它们对应的step1。现在这些step1才是我们最终筛选出来的结果。
从这些step1出发，重复上面的步骤，直到达到设定的搜索停止条件为止。
不难看出，lookahead search的核心在于，对每一步做筛选时，都会先“向前看K步”，用K步后的收益去评估当前步骤的结果。所以，beam seach可以看成是lookahead seach的K=1版本。
值得一提的是，如果你了解过
MCTS（Monte Carlo tree search，蒙特卡洛搜索树算法）
，你会发现lookahead seach和它在广义上非常相似。在MCTS中，同样是搜索最佳路径（策略），我们会在搜索的过程中使用一些随机因子，目的是通过多
探索（explore）
的方式去预估每一步的价值，进而学习价值函数。
而在lookahead search + PRM中，这个训练好的、固定的PRM则替代了MCTS中的explore过程
，这时lookahead search只要负责
利用（exploit）
就可以了。
所以，lookahead search + PRM，从广义上看就是一种MCTS的方法。
关于MCTS，我们会在后面的文章中来详细讲解它。
2.4 如何选择最佳搜索方法
有这么多的搜索方法，我们应该如何做选择，才能保证能得到最好的搜索效果呢？（也就是让模型的推理效果最优）
我们可以先从直觉上大概猜想一下，什么要素会影响模型的搜索效果：
搜索预算（generation budget）
：这就是指采样数量N。即对于一个问题，你在搜索时想做多少条并行采样。我们考虑N的原因是GPU的算力是有限的。
问题难度（difficulty）
：对于简单的问题，模型或许利用预训练好的知识就能做回答，不需要引入复杂的搜索方法。而对于较难的问题则可能更需要设计好搜索方法引导模型步步思考。
根据这两个直觉性的思考，作者做了如下两组实验：
（1）搜索预算对搜索效果的影响
我们先看左图，它探究了搜索预算对模型推理效果的影响
。橘色代表best-of-N weighted，是我们用于比较的基线。从图中可以看出：
当搜索预算较小时（即对一个问题的并行采样N较小时）
，beam seach > best-of-N > lookahead
当搜索预算较大时（即对一个问题的并行采样N较大时）
，best-of-N > beam seach > lookahead，同时在beam search的M超参选择合理的情况下，其效果约等于best-of-N
最复杂的lookahead搜索方法，表现基本处于垫底趋势。
我们可以从直觉上给出上述现象的一些解释
（这里论文说得也很模糊，所以依然是笔者根据论文部分内容做了主观理解）：
当N较小时，在对一个问题并行采样出的N种结果里，可能很难直接命中最佳答案
。这时通过更为复杂的搜索策略（例如beam search），从一次性采样转变为步步精心选择，可能会取得更好的结果。当N较大时同理可推。
在
PRM训练得足够好的情况下，在搜索的每一步，我们可以更关注利用当前（expolit）而不是探索未来（exploration）
，正如前文所说，探索的本质是去寻找更好的价值函数，但它的风险是随机性较强。所以复杂的lookahead方法在我们的场景下表现并不好。
（2）问题难度对搜索效果的影响
我
们再来看右图，它探究了问题难度对搜索效果的影响。bin1～5分别表示5种难度（从易到难），每个bin下的4根柱子分别表示不同的搜索预算（4，16，64，256）。也许是考虑到先前实验中lookahead表现不佳，以及beam search本身也能代表一种lookahead搜索的原因，这里作者直接把lookahead从实验对象中剔除了。从这个实验中我们可以发现：
在中上难度的问题上（难度3，4），beam search > best-of-N；
在简单的问题上（难度1，2），best-of-N > beam search（或者beam search略略好一些）。
当然，上面这些趋势只是概括性的，具体的波动还收到搜索预算的限制，大家可以自行细看实验图。
我们从直觉上对以上现象做一些解释
（同样包含笔者的主观解读）：
在简单的问题上，模型在pretrain阶段获取的知识已能使其大概率作出正确回答
。也就是说，你对一个简单问题随机sample出N个结果，这个N里大概率有正确答案。这时就没有必要引入复杂的搜索方法（引入PRM），做步步评估了，毕竟评估过程是链式的，有一个step评估错误，就可能有连锁负面效果。对复杂的问题同理可推。
当问题特别困难时，“PRM + 某种搜索方法”的后处理模式可能已经不能单独适用，这时或许提升模型在pretrain阶段的预训练知识更为必要
。这使我想起了openAI o1来，它在一些复杂的数学推理上也能取得很好的效果，所以虽然官方说得是在inference阶段增加了算力（包含强化学习的后处理方式），但是或许它对base模型也做过了处理，例如增加了代码、数理逻辑知识的比例。或者可能改变这个比例直接重新训练了base模型。
（3）对搜索方法的小结
搜索的效果（也即模型的推理效果）受到搜索预算和问题难度的限制，我们是在这些限制中去选择合适的搜索方法。
当搜索预算较小，问题较难时，更适合beam search方法，但要注意超参的调整
