24年12月来自北京交大的论文“o1-coder: an o1 replication for coding”。
O1-CODER,试图复制 OpenAI 的 o1 模型,重点关注编码的任务。它集成强化学习 (RL) 和蒙特卡洛树搜索 (MCTS),以增强模型的系统-2 思维能力。该框架包括训练
测试用例生成器 (TCG)
进行标准化代码测试,使用 MCTS 生成具有推理过程的代码数据,以及迭代微调策略模型以最初生成伪代码,然后生成完整代码。该报告还讨论了在实际应用中部署类似 o1 模型的机遇和挑战,建议过渡到系统-2 思维范式并强调构建世界模型的必要性。
OpenAI 最近推出 o1 模型(OpenAI,2024),该模型展示了系统-2 思维能力。该模型代表人工智能在执行需要高阶认知功能的复杂推理任务方面取得了重大进步。在其发布后,出现了大量分析和复制工作,表明人们对推理模型的兴趣日益浓厚。值得注意的作品包括 g1 (Benjamin Klieger,2024)、OpenO1 (ope,2024)、O1-Journey (GAIR-NLP,2024)、OpenR (Team,2024)、LLaMA-O1 (SimpleBerry,2024)、LLaMA-Berry (Zhang,2024)、Steiner (Ji,2024)、Thinking Claude (Tu,2024)、LLaVA-o1 (Xu,2024),以及几个工业版,如 k0-math、DeepSeek-R1-Lite、Macro-o1 (Zhao,2024)、Skywork o1、QwQ (Qwen Team,2024) 和 InternThinker (上海AI实验室,2024),如图所示:
在 o1 模型之前,大语言模型 (LLM) 主要展现系统-1 的功能,其特点是反应快速、直观。这些模型在主要由问答 (Q, A) 对组成的数据集上进行训练,缺少涉及深思熟虑和分析处理的中间推理步骤。这是因为人类很少在互联网或其他地方记录他们的思维过程。传统上,人们使用诸如思维链 (CoT) 提示之类的技术来指导模型在得出答案之前进行逐步推理。更直接、更有效的方法是创建包含推理序列的数据集,例如 (Q, ..., S/i, ..., A),其中 S/i 表示导致最终答案的单个推理步骤。因此,增强系统-2 推理能力的第二种方法是监督微调 (SFT)。然而,大多数公开可用的数据都以问答形式记录,注释或蒸馏此类数据(尤其是对于复杂任务)既昂贵又具有挑战性。
在探索推理过程数据缺失的情况,还有第三种方法——强化学习 (RL)。
人们普遍认为,o1 通过将强化学习与预训练相结合来解决推理数据不足的问题。强化学习以其在过去十年中探索和发现新策略而非依赖预定义数据的能力而闻名。回顾机器学习的关键发展,深度学习和大规模预训练分别推动了模型架构的转变和对标记数据的要求。相比之下,强化学习解决目标函数转变的不同方面。在缺乏明确指导或明确目标的情况下,RL 利用探索来寻找新知识和解决方案。因此,将预训练与 RL 相结合可以产生强大的学习和搜索协同作用,其中预训练压缩现有的人类知识,而 RL 使模型能够探索新的可能性。
编码是一项典型的任务,需要系统-2 思维,涉及仔细、合乎逻辑和逐步解决问题。此外,编码可以作为解决许多其他复杂问题的基础技能。本文介绍复制 o1 的尝试,重点关注编码方面的任务。该方法集成 RL 和蒙特卡洛树搜索 (MCTS) 以实现自我发挥,使模型能够不断生成推理数据并增强其系统-2 功能。
将自对弈强化学习应用于代码生成有两个主要挑战需要解决。
第一个挑战
是结果评估,即评估生成代码的质量。与围棋或数学等可以根据游戏规则或正确答案直接评估结果的任务不同,评估代码需要在测试环境中运行生成的代码并根据测试用例进行验证。不能假设代码数据集总是会提供足够的测试用例。
第二个挑战
涉及定义思考和搜索行为,即确定状态转换和过程奖励的粒度。对于代码生成,关键问题是如何设计推理过程和策略空间来有效地指导模型的行为。
为了应对第一个挑战,训练一个测试用例生成器 (TCG),它根据问题和真值代码自动生成测试用例。这种方法将有助于构建标准化的代码测试环境,为强化学习提供结果奖励。
对于第二个挑战,可以考虑两种可能的方法。一种是“三思而后行(think before acting)”,即模型先形成完整的思维链,然后一次性生成最终答案。另一种方法是“边思考边行动(think while acting)”(Zelikman,2024),即在推理任务的同时生成部分答案。本研究选择前一种方法。
对于代码生成,这意味着首先思考并写出详细的伪代码,然后使用该伪代码生成最终的可执行代码。这样做有两个优点:一是适应性,因为相同的伪代码可以产生不同的具体代码实现;二是粒度可控,因为调整伪代码中的细节级别可以控制推理/搜索行为的粒度。
如下算法 1 提供概述的框架,该框架包含六个步骤:(1)第一步是训练测试用例生成器 (TCG) γ/TCG,它负责根据问题自动生成测试用例; (2) 第二步,在原始代码数据集上运行 MCTS,生成带有推理过程 D/process 的代码数据集,其中包括一个有效性指标,用于区分正确和错误的推理步骤; (3) 一旦有了包含推理过程的数据,第三步就是微调策略模型 π/θ,训练它以“三思而后行”的方式行事; (4) 推理过程数据还可用于初始化
过程奖励模型 (PRM)
ρ/PRM,该模型用于评估推理步骤的质量; (5) 第 5 步最为关键:利用 PRM ρ/PRM 提供过程奖励、利用 TCG γ/TCG 提供结果奖励,使用强化学习更新策略模型 π/θ; (6) 在第 6 步中,基于更新后的策略模型,可以生成新的推理数据。然后,可以使用这些新数据再次微调 PRM(第 4 步)。
因此,步骤 4-5-6 形成一个迭代循环,其中自我博弈不断推动模型改进。六个步骤之间的流程如图所示。
测试用例生成器训练
目标
测试用例生成器(TCG)是一种用于自动创建输入-输出测试用例的工具,它在支持代码生成任务中的程序验证方面起着关键作用。
在训练阶段,通常使用标准输入-输出测试用例来评估生成代码的正确性。这些测试用例的通过率是评估生成代码质量的关键指标,并作为结果奖励信号来指导策略模型的训练。此奖励信号有助于模型改进其生成策略,从而增强其生成准确且功能性代码的能力。
在推理阶段,当训练后的模型负责代码生成时,通常没有标准测试用例来验证生成代码的正确性。测试用例生成器(TCG)通过为策略模型提供自我验证机制来缓解这一限制,该机制允许策略模型在最终生成之前进行评估。因此,策略模型能够根据验证结果选择最佳输出路径。
训练
训练过程分为两个不同的阶段:监督微调 (SFT) 和直接偏好优化 (DPO) (Rafailov,2024)。将未经微调的生成器表示为 γ/TCG/base。
SFT 阶段的主要目标是确保生成器的输出遵循预定义的格式,从而能够准确解析和提取生成的测试用例。此阶段的训练数据来自 TACO 数据集 (Li,2023),其格式为 {question、solution、test_case}。为了标准化模型的输入和输出,开发一种模板格式,如下图所示:
经过 SFT 后,生成器记为γ/TCG/sft。
DPO 阶段的目标,是引导模型生成符合特定偏好的测试用例,从而提升测试用例生成器(TCG)的性能和可靠性。本研究采用人工构造样本对的 DPO 方法,通过构建偏好数据集,提高模型与期望偏好的一致性。DPO微调依赖于预先构建偏好数据集 D/pref = {x, y/w, y/l},其中 x 为包含指令、问题和代码的提示;y/w 为正例,即与偏好一致的测试用例;y/l 为负例,即与偏好不一致的测试用例。
采用以下规则构建偏好数据:对于y/w,直接使用完全匹配的三个采样测试用例作为正例;对于 y/l,将三个采样测试用例的输出进行混洗,然后将原始输入连接起来,使得三个测试用例的输入输出对不完全匹配,并使用三个不完全匹配的测试用例作为反例。训练目标旨在基于初始 SFT 模型 γ/TCG/sft 去优化 γ/TCG/θ,同时将隐奖励建模与参考模型 γ/TCG/ref 结合起来,参考模型 γ/TCG/ref 代表初始 SFT 模型 γ/TCG/sft。目标函数定义如下:
在 DPO 之后,将生成器记为 γ/TCG/dpo,代表最终生成器 γ/TCG。
实验
使用DeepSeek-1.3B-Instruct(Guo,2024)作为测试用例生成器(TCG)的基模型,然后是 SFT 和 DPO。微调阶段采用 QLoRA 技术(Dettmers,2023),秩参数 r = 1,以调整以下模块:q_proj, o_proj, k_proj, v_proj, gate_proj, u_pproj, down_proj。学习率设置为5×10-4,以平衡训练稳定性和收敛速度。训练数据来自 TACO 训练数据集的一个子集,该子集遵循 ACM 竞赛格式,包含大约 10,000 个样本。同样,测试数据来自 TACO 测试数据集的一个子集,也符合 ICPC 竞赛格式,由314个样本组成。
在 TACO 测试的不同阶段对生成的测试用例的质量进行测试。在 SFT 阶段之后,γ/TCG/sft 在标准代码上生成的测试用例通过率为 80.8%,证明生成器在初步微调后能够高效地生成测试用例。此外,γ/TCG/dpo的性能达到89.2%,与γ/TCG/sft 相比有显著的提高。这表明偏好优化通过改进模型的决策过程,显著提高生成器生成更可靠测试用例的能力。
在实际场景中,生成器的性能通常满足评估代码正确性的要求。展望未来,在推理过程中将测试用例生成器(TCG)作为结果验证器。这种方法旨在通过根据动态生成的测试用例验证生成输出来确保生成输出的正确性,从而实现更强大的代码生成推理时搜索。
此外,考虑将自我对弈纳入 TCG 的训练中。在此设置中,策略模型将生成旨在通过 TCG 生成测试用例的代码,而 TCG 则旨在生成逐渐更具挑战性的测试用例。这种对抗性互动可以促进策略模型和测试用例生成器(TCG)的相互改进。
推理-增强的代码数据合成
基于伪代码的推理过程
推理过程的定义至关重要。探索一种基于伪代码的 MCTS 方法,旨在指导大语言模型对复杂代码任务进行深度推理。伪代码作为自然语言描述和实际代码之间的中间表示,提供一种更抽象、更简洁的方式来表达算法或程序的逻辑流程。为了将伪代码推理集成到步级思维链 (CoT) 中,如图所示,定义三个融入伪代码推理的关键行为动作:
• 动作 1:使用伪代码定义算法结构:在此动作中,模型概述主要功能的结构和接口,而无需深入研究实现细节。目的是使模型能够掌握整体任务结构,包括每个主要功能的输入、输出和核心功能。
• 动作 2:细化伪代码:在此动作中,模型迭代细化动作1中定义的伪代码,逐步明确每个函数的步骤、逻辑和操作,为最终的代码实现做准备。
• 动作 3:从伪代码生成代码:此动作的目标是将伪代码的结构和逻辑准确地转换为可执行代码,确保生成的代码满足任务要求。
这些动作确保模型在推理过程中使用伪代码作为认知工具,增强其对复杂代码生成任务的推理能力。需要注意的是,这三个动作并不意味着推理链仅限于这些步骤。如图所示,模型可能需要在整个推理过程中反复调用动作 2 来迭代细化伪代码,直到它足够成熟以用于最终的代码生成。
为了评估带有伪代码推理的步级 CoT 的有效性,使用 Qwen 系列开源模型(Yang,2024)和 Mostly Basic Python Problems (MBPP) 数据集(Austin,2021)作为基准进行实验。在实验中,采用基于蒙特卡洛树搜索 (MCTS) 的采样策略,并比较常规 CoT 和带有伪代码推理 CoT 的 Pass@1,以及正确推理路径上最后一步的平均采样通过率 (ASPR)。结果表明,在推理正确的情况下,加入伪代码可以显着提高生成代码的质量。
推理过程数据的合成
用蒙特卡洛树搜索 (MCTS) (Kocsis & Szepesva ́ri, 2006; Feng et al., 2023; Qi et al., 2024) 以 D_process = {(Q_i, ··· , S_i^j, v_i^j, ··· , C_i′)} 的形式构建步级过程奖励数据,其中 v_i^j 表示到步骤 S_i^j 为止的推理路径评估,C_i′ 是从最后一步 S_i^m 派生的可执行代码。在此过程中,采用标准 MCTS 推出策略进行路径探索。对于每个问题 Q_i,应用前面定义的伪代码提示策略来指导推理过程。当到达终端节点 S_i^m 时,形成完整的伪代码推理路径 (Q_i , S_i^1 , . . . , S_i^m )。终端节点 S_i^m 的奖励值 v_i^m 基于两个关键指标计算:
• 编译成功率 (compile):此指标确定生成的代码是否可以成功编译。compile 的值是二进制的,compile = 1 表示成功,compile = 0 表示失败。
• 测试用例通过率 (pass):假设编译成功,进一步评估生成的代码是否通过测试用例。通过率的计算公式为 pass = Num_passed/Num_testcase,其中 Num_passed 是通过的测试用例数,Num_testcase 是用于验证的测试用例总数。
终端节点 S_i^m 的奖励值,计算为这两个指标的加权和:v_i^m = α·compile + (1−α)·pass,其中 α 是控制编译成功率和测试通过率相对重要性的超参数。
一旦计算出终端节点的奖励值 v_i^m,就将该值反向传播到路径上的所有前导节点,为每个步骤 (S_i^j,v_i^j) 分配一个奖励值 v_i^j。由于 MCTS 过程中有多次展开,反向传播过程中某个节点 v_i^j 的累计奖励可能超过1,因此对路径上每个节点的奖励值进行归一化,得到最终的步骤有效性值。
在构建推理过程数据集时,对于每个问题Q_i,如果通过搜索找到正确的答案,则保证获得至少一个终端节点(S_i^m, v_i^m),且v_i^m = 1。在完成搜索后,从正确的终端节点 (Q_i, S_i^1, ..., S_i^m, v_i^m)中选择完整的推理路径,其中 v_i^m = 1,形成策略模型的初始化数据集。
策略模型初始化
在推理数据合成任务后,用数据集中每一个完整的推理方案来初始化策略模型 π_θ。此步骤旨在帮助 π_θ 更好地理解任务要求并遵循预期的动作行为,为后续的迭代训练提供最佳起点。
给定问题 Q_i,策略模型 π_θ 在第 j 步生成的具体推理步骤内容,可以表示为 π_θ(S_i^j | Q_i, S_i^1:j−1),其中S_i^j = (w_1, w_2, . . . , w_k)。其中,S_i^j 表示一个推理步骤的内容,以特定的分隔符分隔,w 表示 π_θ 在每个解码步骤生成的 token。S_i^1:j−1 表示由前 i 个推理步骤输出形成的上下文。