OpenAI o1 技术系列1:整体框架,利用Test-Time Scaling Law提升逻辑推理能力

文摘   2024-09-29 17:03   中国香港  

前段日子OpenAI推出的o1模型,以其提升显著的逻辑推理能力,引发了人们对它背后训练方法的热烈讨论。关于o1的介绍和输出结果demo,这里就不再赘述,大家可以去openai的官网上阅读(很短,读起来很快,因为秘密都藏好了)。我相信最近的一段时间里,当大家在网上探索o1是如何训练时,肯定会看到以下几个热点词:

  • Test/Inference-Time scaling law,通过增加推理阶段的算力提升模型的推理能力
  • Post Training,通过后训练提升模型的推理能力
  • PRM/ORM:基于过程/结果的奖励模型
  • CoT:思维链
  • 强化学习、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都回答清楚了。

在理解了这些的基础上,我们就大概知道【框架】长什么样,具体而言:

  1. 首先,我们需要引导模型从“只产生结果”变成“同时产生中间步骤和结果”,这个过程称为post-training,这个过程里,你既可以使用基于强化学习的技术,也可以只做sft微调。你既可以只关注模型是否遵循了格式(即只关注是否产出了中间结果,而不关注中间结果的质量),也可以同时关注格式+中间结果质量
  1. 其次,我们需要训练一个能够评估中间结果的verifier,这其实也算post-training的一部分。这个verifier既可以被用在基于强化学习的post-training中(但它不一定是唯一的价值评估模型),也可以被用在第3步的inference阶段中引导搜索。
  1. 接着,我们需要设计一种搜索方法。它能够根据verifier返回的中间结果评估分数,在推理阶段更好指引模型做步步搜索,以达到思考的每一步都能最优化。(如果在post-training阶段,你只关注格式,那么inference阶段应用这种搜索方法就能更好指引搜索结果;如果你已经关注了格式+质量,这种方法也依然能达到更上一层楼的效果)。同时注意,如果在post-training阶段我们已经使用了基于强化学习的方案,那么这种搜索方法同样也能用在模型生产“经验数据”的阶段,这样可以在自产自消的基础上,动化筛选高质量的数据集再做对齐(避开人类标注,不是必须,但如果做得好的话,inference可能只需要在post-training阶段使用,目的是筛选高质量数据,而不需要在post-training后再做使用了)。

有了对框架的理解,接下来我们就来看deepmind基于这个框架做的两种方案:

  • 方法一:利用PRM(Process-supervised Reward Model)指引搜索

    • post-training:只引导模型对齐格式 + 训练基于过程的verifier(PRM)
    • inference:使用PRM指引搜索
  • 方法二:通过sft直接改变模型的输出分布(Revise proposal distribution)

    • post-training:通过sft方式引导对齐格式 + 保障中间结果基本质量 + 训练基于过程的verifier(PRM),这里直接用了的方案一中的PRM。
    • inferece:使用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

  • 决定如何使用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具有迷惑性。

  • 重复上面过程,迭代式地训练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的得分。

  • 取出得分为top M的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略略好一些)。
  • 在最难的问题上(难度5),两者的表现都不好。
  • 当然,上面这些趋势只是概括性的,具体的波动还收到搜索预算的限制,大家可以自行细看实验图。

我们从直觉上对以上现象做一些解释(同样包含笔者的主观解读):

  • 在简单的问题上,模型在pretrain阶段获取的知识已能使其大概率作出正确回答。也就是说,你对一个简单问题随机sample出N个结果,这个N里大概率有正确答案。这时就没有必要引入复杂的搜索方法(引入PRM),做步步评估了,毕竟评估过程是链式的,有一个step评估错误,就可能有连锁负面效果。对复杂的问题同理可推。

  • 当问题特别困难时,“PRM + 某种搜索方法”的后处理模式可能已经不能单独适用,这时或许提升模型在pretrain阶段的预训练知识更为必要。这使我想起了openAI o1来,它在一些复杂的数学推理上也能取得很好的效果,所以虽然官方说得是在inference阶段增加了算力(包含强化学习的后处理方式),但是或许它对base模型也做过了处理,例如增加了代码、数理逻辑知识的比例。或者可能改变这个比例直接重新训练了base模型。

(3)对搜索方法的小结

  1. 搜索的效果(也即模型的推理效果)受到搜索预算和问题难度的限制,我们是在这些限制中去选择合适的搜索方法。
  1. 当搜索预算较小,问题较难时,更适合beam search方法,但要注意超参的调整
  1. 当搜索预算较大,问题较简单时,更适合best-of-N方法
  1. 当PRM训练得足够好时,更复杂的搜索方法(例如lookahead search)的表现可能并不好
  1. 当问题特别困难时,test-time scaling law的作用可能有限,这时可能有必要重新审视pretrain阶段,通过增加/调整数据比例,扩张模型规模等等方式,给pretrain模型注入更多的相关知识。

三、方法二:直接改变模型的输出分布

在上文中,我们介绍了使用“PRM + 搜索方法”的模式,使用PRM来引导模型的步步推理过程,找到最佳答案。现在,我们再来看另外一种提升模型推理准确率的方法:直接改变模型的输出分布(refining the proposal distribution)。

在“PRM + 搜索方法”的模式中,我们曾提到过,先通过sft的方式训练模型能够按给定的格式(步步思考的方式)生成结果,这一个步骤只起格式对齐作用,不对思考过程质量负责。那么此时你肯定会想:如果我能用具有高质量思考步骤的数据,直接sft模型,那么不是也能达到同样的效果吗?而这样的模型一旦训成,在之后的推理阶段,它就会在吃进一个问题后,【自主地】开始步步思考了,不需要什么搜索方法,也不需要PRM,只需要利用它的生成能力,就像下图这样:

如上图,模型在吃进问题后,经过了多次attemp,步步逼近正确结果。

而训练这样一个模型,sft的过程没什么好说的,关键是在“高质量sft数据如何收集”这一方面,所以接下来我们就详细来看这一点。除了本文提供的这种训练方法外,我们最近在openAI o1热搜里高频看到的STaR、Quiet-STaR这两个模型,也可以大致归为这一类(Quiet-STaR有一些特殊,涵盖了更多巧思),我们在后面的文章里,会结合源码再来看看这两个有趣的模型。

3.1 收集sft数据

如何才能使用最少的成本,生成高质量的涵盖中间思考步骤的sft数据?我们直接来看作者们的做法:

  • 首先,我们有“问题 + 答案”这种有监督的数据,只是我们缺少中间过程。我们希望模型能模拟人类思考方式,通过尝试各个attempt,逐步逼近正确答案。

  • 然后,我们用同样的方法,先对模型做格式微调,使其能够产出“中间结果 + 答案”(注意,不是上图中的多个attempt,而是上图中的单个attempt)。这一步同样只负责格式,不负责质量

  • 接着,利用模型,对每个问题,并行采样出N个attempt(单个attempt = 中间结果 + 答案,N=64),同样过滤掉一些涵盖无效答案的attempt(再复习下,无效!=错误)

  • 由于每个attempt中都包含了答案,答案是有标签的,所以我们可以知道哪些attempt给出了正确答案。我们说这些attempts是正确。

  • 现在,我们希望为每个正确attempts匹配上若干错误attempts,作为一条训练数据。也就是我们的训练数据是“问题 + 若干错误attempts + 正确attempt”的形式,这一步是让模型模拟人类思考的模式,从步步错误的attempts中推出正确的。具体做法如下

    • 从0~4中均匀采样出一个数字,记为x。现在对于每个正确回答,我们都要为其匹配x个错误回答

    • 先从所有的错误回答中,根据编辑距离,找到和这个正确回答最相似的错误回答(编辑距离可能不是精确衡量相似度的好方法,但是对于这个场景够用了)

    • 然后再从剩下的不正确的回答中,随机采样x-1条。

    • 对于一条正确回答,现在我们可以按照(问题,随机不正确回答x-1,最相似的不正确回答,正确回答)来组成一条经验数据,这条数据被称为trajectory,将被我们用来对模型做sft。其中,正确回答是模型要预测的部分,之所以要选择一条比较相似的错误回答放在这里,是为了能让模型知道自己是在不断纠正错误中【学得更好】

  • 同时注意,当你在某次采样时,采样出的x < 4的时候,你需要把采样分布调为0~x,然后进行下一次采样。这样做是为了确保我们的trajectories尽可能涵盖各个不同数量的非正确-正确答案组合。

当你审视上面这个收集sft数据的方法时,你可能会觉得它有很多可以改进之处,例如:只是答案正确,就认为这个attempt正确了吗?它的中间步骤中可能是有错的,这可能是个false positive。所以,一种可以想见的改进办法是,依然引入训练好的PRM,对中间步骤再做评估,更有利于选择尽可能正确的attempt。举这个例子是为了说明,本文只是提供了一种可能的数据收集方法,而不是数据收集标准,在我们实际的研究中,还有很多可以探索的地方,不过它们的核心都是【自动化】收集supervised数据。

3.2 模型训练的停止时间

我们再来看一个有趣的问题:当我们使用上面的数据,进入sft训练后,我们该怎么判断sft训练可以停止的时间?

  • 模型在评估集上的loss判断模型的训练程度。我们当然期望评估集loss逐步变小,当我们观察到评估集loss上升时,这往往意味着模型已经开始出现过拟合了,我们一般会在这个趋势开始前对模型训练做早停。

  • 但是,在这个实验里,作者观察到在评估集loss上升很久以后,模型依然表现得非常好。这是因为评估集数据是off-policy的,模型在不断更新,评估集的数据分布已经赶不上模型了,所以模型在它上面的loss并不能真实反馈出模型训练的情况

  • 所以,最终作者选择在发现过拟合现象的一段时间后,再停止训练模型

3.3 如何选择最佳生成方法

正常来说,如果这样一个sft模型训练得足够好,那么我们喂给它一个问题,只需要让它源源不断生成attempt直到<end>,然后取最后一个attempt就可以了。但是事实上,可能会存在以下问题:

  • 我们并不能保证最后一个attempt就一定是正确的。很有可能中间产出了正确attempt,模型又把它修错了。

  • 我们并不能保证对于一个问题,只生成一个回答(即一个attempts链),这其中就一定有正确的attempt。

所以,尽管理论上我们可以使用sft模型。但是为了取得更好的效果,实际上我们依然会配合“verifier + 搜索方法”的方式做推理。之所以这里写成verifier而不是PRM,是因为verifier不一定是基于过程的,也可以是基于结果的(orm),我们来看一张架构图:

  • 【左侧-parallel sampling】:这里只是告诉我们parallel sampling是如何做的,并没有告诉我们该选择哪个结果。在我们的场景下,parallel sampling表示对一个答案生成多个attempts链(每个链里涵盖若干attempt)

  • 【左侧-sequential revisions】:这是最naive的推理方式,也是生成单个attempt链,其中包含多个attempt。同样只阐述过程,不讲选择方案

  • 【右侧-parallel best-of-N】:右侧的图开始讲选择方案。这里引入了verifier做选择。对于一个问题,模型可以并行采样多个结果,每个结果都是一串attempts链条,最后借助verifier选择得分最高的attempts链条。

  • 【右侧-sequential revisions】:同样也是引入verifier,我们让模型按顺序生成多个attempts,不过我们并不能完全相信最后一个attempt一定是正确的。所以我们用verifier帮我们判断在这个链条中应该选择哪个attempt。

  • 【右侧-combing sequential/parallel】:同时生成多个链,每个链里若干个attempt。先在每条链内部通过verifier找到最佳attempt,然后所有的最佳attempt放在一起,通过verifier找到最优的

作者在这里同样通过实验的方式,在搜索预算(采样数量N)和问题难度的限制下,决定自己应该使用哪种生成方法。实验内容我们就不放了,大家可以自行阅读文章,这里我们直接给出结论:

  1. 在简单的问题上,通过sequential的生成方式(单个attempts链中找最佳attempt)效果更好
  1. 在复杂的问题上,通过sequential + parallel合并的方式效果更好,但要选择好合适的超参

四、Pretrain还是Inference?

在上面的过程中,我们介绍了两种把算力集中在inference/post training阶段来提升模型最终推理效果的方法(但正如本文前言所说,这其实算两种框架,其中的具体做法是指例子,实践起来我们可以做很多变种)。在了解这些方法的前提下,你一定很想从量化实验结果中知道:在inference scaling law指引下的结果,能比pretraining scaling law指引下的结果好多少?换句话说,我给inference和pretrain分配相同的算力,他们的表现又会怎么样?

在论文中,作者简单给出了一些建模指标和说明。笔者曾尝试从更细节的角度去理解它们,但是很遗憾,论文给出的描述有许多有歧义和模糊的地方,以笔者目前的能力无法给出一个自洽的解释。所以在这一节,笔者打算直接放出实验效果图,从直观上大概解释一下pretrain VS inference的问题:

左右两幅图分别是在【利用PRM指引搜索】和【直接改变模型输出分布】这两种方法上的效果,两者的实验方法一致,所以我们只需看其中一幅图就好。在这个实验中:

  • 横轴表示算力(FLOPs)
  • 纵轴表示模型的推理效果。
  • 不同颜色的曲线表示不同的的问题难度
  • 带圆点的曲线表示将这部分算力全部用于inference时模型的效果
  • 每个虚线表示算力在pretrain和inference中的交换比例。星号表示当把这部分算力以一定比例给pretrain时,模型最终的推理效果(不能完全换给pretrain,因为pretrain阶段做继续训练后,本身也需要算力做推理)

从中,我们可以发现,对于较简单的问题(难度在1~3),星号普遍在圆点之下,也就是这一部分固定的算力全部用于inference时表现较好;而对于较复杂的问题(难度在4~5),星号普遍在圆点之上,这意味着复杂问题还是需要依赖预训练知识的更新。


备注:进群,进入大模型技术群

id:baobaogpt,记得备注呦

包包算法笔记
大模型技术和行业认知
 最新文章