KIMI K1.5: SCALING REINFORCEMENT LEARNING WITH LLMS

作者/机构: Kimi Team

A1 主要贡献

本文介绍了Kimi k1.5的训练方法，这是一个利用强化学习（RL）进行训练的多模态大语言模型（LLM），旨在探索超越现有静态数据集限制的持续扩展新路径。其核心设计与训练要素包括：

该模型的长CoT版本在多个基准和模态上达到了最先进的推理性能（例如，AIME 77.5分，MATH 500 96.2分，Codeforces 94百分位，MathVista 74.9分），与OpenAI的o1相当。其短CoT版本同样取得了领先的推理结果（例如，AIME 60.8分，MATH 500 94.6分，LiveCodeBench 47.3分），显著优于现有的短CoT模型如GPT-4o和Claude Sonnet 3.5。

图1：Kimi k1.5 long-CoT 结果
图1：Kimi k1.5 long-CoT 结果

图2：Kimi k1.5 short-CoT 结果
图2：Kimi k1.5 short-CoT 结果

A2 方法细节

Kimi k1.5的开发包括预训练、常规监督微调（SFT）、长CoT监督微调和强化学习（RL）等阶段。本报告重点关注RL阶段，首先概述RL提示集构建（2.1节）和长CoT监督微调（2.2节），然后在2.3节深入讨论RL训练策略。

2.1 RL提示集构建

构建高质量RL提示集的关键属性。初步实验表明，RL提示集的质量和多样性对确保强化学习的有效性至关重要。一个精心构建的提示集不仅能引导模型进行稳健的推理，还能降低奖励操纵（reward hacking）和对表面模式过拟合的风险。一个高质量的RL提示集应具备三个关键属性：
* 多样化覆盖：提示应涵盖科学、技术、工程和数学（STEM）、编码和通用推理等广泛领域，以增强模型的适应性和跨领域的广泛适用性。
* 难度均衡：提示集应包含分布均匀的简单、中等和困难问题，以促进渐进式学习并防止对特定复杂性水平的过拟合。
* 可准确评估：提示应允许验证者进行客观可靠的评估，确保模型性能的衡量基于正确的推理过程，而非表面模式或随机猜测。

确保提示集多样化的方法。为了实现提示集的多样化覆盖，我们采用自动过滤器来选择需要丰富推理且易于评估的问题。我们的数据集包括来自STEM领域、竞赛和通用推理任务等不同领域的问题，并整合了纯文本和图文问答数据。此外，我们开发了一个标签系统，按领域和学科对提示进行分类，以确保不同主题领域的均衡表示【索引25，M. Li et al.，From quantity to quality: Boosting llm performance with self-guided data selection for instruction tuning，2023】；【索引28，W. Liu et al.，What makes good data for alignment? a comprehensive study of automatic data selection in instruction tuning，2023】。

基于模型的难度评估方法。我们采用一种基于模型的方法，利用模型自身的能力来自适应地评估每个提示的难度。具体来说，对于每个提示，一个SFT模型使用相对较高的采样温度生成十次答案。然后计算通过率，并将其作为提示难度的代理指标——通过率越低，难度越高。这种方法使难度评估与模型的内在能力保持一致，使其在RL训练中非常有效。通过这种方法，我们可以预先过滤掉大多数琐碎的案例，并在RL训练期间轻松探索不同的采样策略。

防止奖励操纵的策略。为了避免潜在的奖励操纵【索引9，Everitt et al.，Reward Tampering Problems and Solutions in Reinforcement Learning: A Causal Influence Diagram Perspective，2021】；【索引37，Pan et al.，The Effects of Reward Misspecification: Mapping and Mitigating Misaligned Models，2022】，我们需要确保每个提示的推理过程和最终答案都能被准确验证。经验观察表明，一些复杂的推理问题可能答案相对简单且容易猜测，导致假阳性验证——即模型通过错误的推理过程得出了正确答案。为解决此问题，我们排除了易于出现此类错误的题目，例如多项选择题、判断题和证明题。此外，对于一般问答任务，我们提出了一种简单而有效的方法来识别和移除容易被操纵的提示。具体地，我们让模型在没有任何CoT推理步骤的情况下猜测可能的答案。如果模型在N次尝试内预测出正确答案，则认为该提示过于容易被操纵并被移除。我们发现设置N=8可以移除大多数容易被操纵的提示。开发更先进的验证模型仍是未来研究的一个开放方向。

2.2 长CoT监督微调

使用长CoT预热数据集进行SFT。在拥有了精炼的RL提示集后，我们采用提示工程构建了一个小而高质量的长CoT预热数据集，其中包含针对文本和图像输入的经过精确验证的推理路径。这种方法类似于拒绝采样（RS），但侧重于通过提示工程生成长CoT推理路径。由此产生的预热数据集旨在封装对类人推理至关重要的关键认知过程，例如：规划，模型在执行前系统地概述步骤；评估，涉及对中间步骤的批判性评估；反思，使模型能够重新考虑和完善其方法；以及探索，鼓励考虑替代解决方案。通过在这个预热数据集上进行轻量级的SFT，我们有效地引导模型内化这些推理策略。结果是，经过微调的长CoT模型在生成更详细、逻辑更连贯的响应方面表现出更强的能力，从而提升了其在各种推理任务中的性能。

2.3 强化学习

2.3.1 问题设定

基于思维链（CoT）的RL问题设定。给定一个由问题 $x_i$ 和对应真实答案 $y_i^*$ 组成的训练数据集 $D = \{(x_i, y_i^*)\}_{i=1}^n$，我们的目标是训练一个策略模型 $\pi_\theta$ 来准确解决测试问题。在复杂推理的背景下，从问题x到解决方案y的映射并非易事。为了应对这一挑战，思维链（CoT）方法提出使用一系列中间步骤 $z = (z_1, z_2, \ldots, z_m)$ 来连接x和y，其中每个 $z_i$ 是一个连贯的标记序列，作为解决问题的重要中间步骤【索引61，J. Wei et al.，Chain-of-thought prompting elicits reasoning in large language models，2022】。在解决问题x时，首先自回归地采样出思维 $z_t \sim \pi_\theta(\cdot|x, z_1, \ldots, z_{t-1})$，然后是最终答案 $y \sim \pi_\theta(\cdot|x, z_1, \ldots, z_m)$。我们使用 $y, z \sim \pi_\theta$ 来表示这个采样过程。注意，思维和最终答案都是作为语言序列被采样的。

规划算法在CoT中的作用。为了进一步增强模型的推理能力，规划算法被用来探索各种思维过程，在推理时生成改进的CoT【索引66，Yao et al.，Tree of thoughts: Deliberate problem solving with large language models，2024】；【索引62，Y. Wu et al.，Inference scaling laws: An empirical analysis of compute-optimal inference for problem-solving with language models，2024】；【索引47，Snell et al.，Scaling llm test-time compute optimally can be more effective than scaling model parameters，2024】。这些方法的核心思想是显式地构建一个由价值估计引导的思维搜索树。这使得模型能够探索一个思维过程的多种延续，或者在遇到死胡同时回溯以探究新的方向。具体来说，设 $T$ 是一个搜索树，其中每个节点代表一个部分解 $s = (x, z_{1:|s|})$。这里 $s$ 由问题 $x$ 和一个思维序列 $z_{1:|s|} = (z_1, \ldots, z_{|s|})$ 组成，其中 $|s|$ 表示序列中思维的数量。规划算法使用一个评论家模型 $v$ 来提供反馈 $v(x, z_{1:|s|})$，这有助于评估当前解决问题的进展并识别现有部分解中的任何错误。我们注意到，反馈可以由判别性得分或语言序列提供【索引68，L. Zhang et al.，Generative verifiers: Reward modeling as next-token prediction, 2024，2024】。在所有 $s \in T$ 的反馈引导下，规划算法选择最有希望的节点进行扩展，从而生长搜索树。上述过程迭代重复，直到推导出完整的解决方案。

将规划算法视为对推理步骤序列的操作。我们也可以从算法的角度来看待规划算法。给定在第 $t$ 次迭代时可用的过去搜索历史 $(s_1, v(s_1), \ldots, s_{t-1}, v(s_{t-1}))$，规划算法 $A$ 迭代地确定下一个搜索方向 $A(s_t|s_1, v(s_1), \ldots, s_{t-1}, v(s_{t-1}))$ 并为当前搜索进展提供反馈 $A(v(s_t)|s_1, v(s_1), \ldots, s_t)$。由于思维和反馈都可以被看作是中间推理步骤，并且这些组成部分都可以表示为语言标记序列，我们用 $z$ 替换 $s$ 和 $v$ 来简化符号。因此，我们将规划算法看作是直接作用于一系列推理步骤的映射 $A(\cdot|z_1, z_2, \ldots)$。在这个框架中，规划算法使用的搜索树中存储的所有信息都被展平为提供给算法的完整上下文。这为生成高质量CoT提供了一个有趣的视角：我们或许可以训练一个模型来近似这个过程，而不是显式地构建搜索树和实现规划算法。在这里，思维的数量（即语言标记）可以类比于传统上分配给规划算法的计算预算。长上下文窗口的最新进展促进了训练和测试阶段的无缝扩展。如果可行，这种方法使模型能够直接通过自回归预测在推理空间中进行隐式搜索。因此，模型不仅学会了解决一组训练问题，还培养了有效解决单个问题的能力，从而提高了对未见测试问题的泛化能力。

RL优化目标。因此，我们考虑使用强化学习（RL）来训练模型生成CoT【索引35，OpenAI，“Learning to reason with LLMs”，2024】。设 $r$ 是一个奖励模型，它根据真实答案 $y^*$ 来判断给定问题 $x$ 的提议答案 $y$ 的正确性，通过分配一个值 $r(x, y, y^*) \in \{0, 1\}$。对于可验证的问题，奖励直接由预定义的标准或规则决定。例如，在编码问题中，我们评估答案是否通过了测试用例。对于具有自由形式真实答案的问题，我们训练一个奖励模型 $r(x, y, y^*)$ 来预测答案是否与真实答案匹配。给定一个问题 $x$，模型 $\pi_\theta$ 通过采样过程 $z \sim \pi_\theta(\cdot|x)$，$y \sim \pi_\theta(\cdot|x, z)$ 生成一个CoT和最终答案。生成的CoT的质量通过它是否能导向一个正确的最终答案来评估。总的来说，我们考虑优化以下策略目标：

RL训练的目标与优势。通过扩大RL训练规模，我们的目标是训练一个模型，它能融合简单基于提示的CoT和规划增强的CoT的优点。该模型在推理过程中仍然自回归地采样语言序列，从而避免了在部署期间高级规划算法所需的复杂并行化。然而，与简单的基于提示的方法的一个关键区别是，模型不应仅仅遵循一系列推理步骤。相反，它还应该学习关键的规划技能，包括错误识别、回溯和解决方案精炼，通过利用整个探索过的思维集合作为上下文信息。

2.3.2 策略优化

在线策略镜像下降算法。我们应用在线策略镜像下降的一种变体作为我们的训练算法【索引1，Abbasi-Yadkori et al.，“Politex: Regret bounds for policy iteration using expert prediction”，2019】；【索引32，Mei et al.，“On principled entropy exploration in policy optimization”，2019】；【索引54，Tomar et al.，“Mirror descent policy optimization”，2020】。该算法迭代执行。在第 $i$ 次迭代中，我们使用当前模型 $\pi_{\theta_i}$ 作为参考模型，并优化以下相对熵正则化的策略优化问题：

其中 $\tau > 0$ 是控制正则化程度的参数。该目标有一个闭式解：

这里的 $Z = \sum_{y',z'} \pi_{\theta_i}(y', z'|x) \exp(r(x, y', y^*)/\tau)$ 是归一化因子。对两边取对数，我们得到对于任何 $(y, z)$，以下约束都成立，这使我们能够在优化过程中利用离策略数据：

代理损失函数及梯度。这启发了以下代理损失函数：

为了近似 $\tau \log Z$，我们使用从 $\pi_{\theta_i}$ 采样的样本 $(y_1, z_1), \ldots, (y_k, z_k)$：$\tau \log Z \approx \tau \log \frac{1}{k} \sum_{j=1}^k \exp(r(x, y_j, y^*)/\tau)$。我们还发现，使用采样奖励的经验均值 $\bar{r} = \text{mean}(r(x, y_1, y^*), \ldots, r(x, y_k, y^*))$ 在实践中能产生有效的结果。这是合理的，因为当 $\tau \to \infty$ 时，$\tau \log Z$ 接近于在 $\pi_{\theta_i}$ 下的期望奖励。最后，我们通过对代理损失求梯度来完成我们的学习算法。对于每个问题 $x$，使用参考策略 $\pi_{\theta_i}$ 采样 $k$ 个响应，梯度由下式给出：

与策略梯度方法的比较。对于熟悉策略梯度方法的人来说，这个梯度类似于使用采样奖励均值作为基线的公式(2)的策略梯度【索引20，Kool et al.，“Buy 4 reinforce samples, get a baseline for free!”，2019】；【索引2，Ahmadian et al.，“Back to basics: Revisiting reinforce style optimization for learning from human feedback in llms”，2024】。主要区别在于响应是从 $\pi_{\theta_i}$ 而不是在策略（on-policy）中采样的，并且应用了 $l_2$ 正则化。因此，我们可以将其看作是常规的在策略正则化策略梯度算法到离策略（off-policy）情况的自然扩展【索引34，Nachum et al.，“Bridging the gap between value and policy based reinforcement learning”，2017】。我们从 $D$ 中采样一批问题并更新参数到 $\theta_{i+1}$，然后将其用作下一次迭代的参考策略。由于每次迭代因参考策略的改变而考虑不同的优化问题，我们也在每次迭代开始时重置优化器。

排除价值网络的原因。我们在训练系统中排除了价值网络，这在以前的研究中也曾被利用【索引2，Ahmadian et al.，“Back to basics: Revisiting reinforce style optimization for learning from human feedback in llms”，2024】。虽然这一设计选择显著提高了训练效率，但我们也假设，在经典RL中用于信用分配的传统价值函数可能不适用于我们的情境。考虑一个场景，模型已经生成了一个部分CoT $(z_1, z_2, \ldots, z_t)$，并且有两个潜在的下一个推理步骤：$z_{t+1}$ 和 $z'_{t+1}$。假设 $z_{t+1}$ 直接导向正确答案，而 $z'_{t+1}$ 包含一些错误。如果一个神谕般的价值函数可用，它会指出 $z_{t+1}$ 比 $z'_{t+1}$ 保留了更高的价值。根据标准的信用分配原则，选择 $z'_{t+1}$ 会受到惩罚，因为它相对于当前策略具有负优势。然而，探索 $z'_{t+1}$ 对于训练模型生成长CoT非常有价值。通过使用从长CoT得出的最终答案的正确性作为奖励信号，只要模型成功恢复并达到正确答案，它就可以从选择 $z'_{t+1}$ 中学习到试错的模式。这个例子的关键启示是，我们应该鼓励模型探索多样的推理路径，以增强其解决复杂问题的能力。这种探索性方法产生了丰富的经验，支持了关键规划技能的发展。我们的主要目标不仅仅是在训练问题上获得高准确率，而是侧重于为模型装备有效的问题解决策略，最终提高其在测试问题上的性能。

2.3.3 长度惩罚

引入长度奖励以应对“过度思考”。我们观察到一种“过度思考”现象，即在RL训练期间，模型的响应长度显著增加。尽管这带来了更好的性能，但过长的推理过程在训练和推理时成本高昂，并且人类通常不倾向于过度思考。为了解决这个问题，我们引入了一个长度奖励来抑制 token 长度的快速增长，从而提高模型的 token 效率。给定问题 $x$ 的 $k$ 个采样响应 $(y_1, z_1), \ldots, (y_k, z_k)$ 和真实答案 $y^*$，令 $len(i)$ 为 $(y_i, z_i)$ 的长度，$\text{min\_len} = \min_i \text{len}(i)$ 且 $\text{max\_len} = \max_i \text{len}(i)$。如果 $\text{max\_len} = \text{min\_len}$，我们为所有响应设置长度奖励为零，因为它们的长度相同。否则，长度奖励由下式给出：

长度奖励机制。实质上，我们在正确响应中鼓励较短的响应并惩罚较长的响应，同时明确惩罚答案错误的较长响应。这个基于长度的奖励然后与原始奖励加权相加。

长度惩罚的预热策略。在我们的初步实验中，长度惩罚可能会在初始阶段减慢训练速度。为了缓解这个问题，我们提出在训练过程中逐步预热长度惩罚。具体来说，我们首先采用标准的无长度惩罚的策略优化，然后在训练的剩余时间里使用恒定的长度惩罚。

2.3.4 采样策略

改进采样策略以提升效率。尽管RL算法本身具有相对较好的采样特性（更难的问题提供更大的梯度），但它们的训练效率有限。因此，一些定义明确的先验采样方法可能会带来更大的性能提升。我们利用多种信号来进一步改进采样策略。首先，我们收集的RL训练数据自然带有不同的难度标签。例如，数学竞赛问题比小学数学问题更难。其次，由于RL训练过程会多次采样同一个问题，我们也可以跟踪每个问题的成功率作为难度的度量。我们提出了两种采样方法来利用这些先验知识提高训练效率。

2.3.5 训练配方更多细节

编码问题的测试用例生成。由于网络上的许多编码问题没有可用的测试用例，我们设计了一种方法来自动生成测试用例，作为训练我们模型的RL奖励。我们的重点主要放在不需要特殊判题器的问题上。我们还假设这些问题有可用的标准解答，以便我们能利用这些解答来生成更高质量的测试用例。我们利用广泛认可的测试用例生成库CYaRon来增强我们的方法。我们使用我们的基础Kimi k1.5模型根据问题陈述生成测试用例。CYaRon的使用说明和问题描述作为生成器的输入。对于每个问题，我们首先使用生成器产生50个测试用D例，并为每个测试用例随机采样10个标准提交。我们用测试用例对这些提交进行测试。如果一个测试用例在10个提交中至少有7个产生匹配的结果，则该测试用例被认为是有效的。经过这一轮筛选，我们得到一组选定的测试用例。如果一个问题及其相关的选定测试用例集能通过10个提交中的至少9个，那么该问题及其测试用例将被添加到我们的训练集中。从1000个在线竞赛问题的样本中，约有614个不需要特殊判题器。我们开发了463个测试用例生成器，产生了至少40个有效测试用例，最终将323个问题纳入我们的训练集。

数学问题的奖励建模。评估数学解题的一个挑战是，不同的书写形式可能代表相同的底层答案。例如，$a^2 - 4$ 和 $(a + 2)(a - 2)$ 都可能是同一问题的有效解。我们采用了两种方法来提高奖励模型的评分准确性：
1. 经典RM：借鉴InstructGPT【索引36，Ouyang et al.，“Training language models to follow instructions with human feedback”，2022】的方法，我们实现了一个基于价值头的奖励模型，并收集了约80万个数据点进行微调。该模型最终以“问题”、“参考答案”和“响应”作为输入，并输出一个标量，指示响应是否正确。
2. 思维链RM：最近的研究【索引3，Ankner et al.，Critique-out-Loud Reward Models，2024】；【索引31，McAleese et al.，LLM Critics Help Catch LLM Bugs，2024】表明，增强了思维链（CoT）推理的奖励模型在性能上可以显著超越经典方法，尤其是在正确性标准较为微妙的任务上，如数学。因此，我们收集了同样大规模的约80万个带有CoT标签的样本来微调Kimi模型。基于与经典RM相同的输入，思维链方法在以JSON格式提供最终正确性判断之前，明确地生成一个逐步的推理过程，从而实现更稳健和可解释的奖励信号。
在人工抽查中，经典RM的准确率约为84.4%，而思维链RM达到了98.5%的准确率。在RL训练过程中，我们采用了思维链RM以确保更准确的反馈。

视觉数据。为了提高模型在现实世界图像推理方面的能力，并实现视觉输入与大型语言模型（LLM）之间更有效的对齐，我们的视觉强化学习（Vision RL）数据主要来源于三个不同的类别：
1. 现实世界数据：包括需要图形理解和推理的各年级科学问题、需要视觉感知和推断的地点猜测任务，以及涉及理解复杂图表的数据分析等。这些数据集提升了模型在现实世界场景中进行视觉推理的能力。
2. 合成视觉推理数据：人工生成的数据，包括程序化创建的图像和场景，旨在提高特定的视觉推理技能，如理解空间关系、几何图案和对象交互。这些合成数据集为测试模型的视觉推理能力提供了一个受控的环境，并提供了源源不断的训练样本。
3. 文本渲染数据：通过将文本内容转换为视觉格式创建，使模型在处理跨不同模态的基于文本的查询时能够保持一致性。通过将文本文档、代码片段和结构化数据转换为图像，我们确保模型无论输入是纯文本还是渲染为图像的文本（如截图或照片），都能提供一致的响应。这也有助于增强模型处理富含文本的图像的能力。
每种类型的数据对于构建一个全面的视觉语言模型都至关重要，该模型能够有效管理广泛的现实世界应用，同时确保在各种输入模态下性能的一致性。

2.4 Long2short：短CoT模型的上下文压缩

Long2short问题的多种方法。虽然长CoT模型性能强大，但与标准的短CoT LLM相比，它在测试时消耗更多的token。然而，可以将长CoT模型的思维先验知识迁移到短CoT模型中，从而在有限的测试时token预算下也能提升性能。我们提出了几种解决这个long2short问题的方法，包括模型合并【索引65，Yang et al.，“Model merging in llms, mllms, and beyond: Methods, theories, applications and opportunities”，2024】、最短拒绝采样、DPO【索引41，Rafailov et al.，“Direct preference optimization: Your language model is secretly a reward model”，2024】和long2short RL。这些方法的详细描述如下：
* 模型合并：模型合并被发现在维持泛化能力方面很有用。我们也发现它在合并长CoT模型和短CoT模型时能有效提高token效率。这种方法将一个长CoT模型与一个较短的模型结合，无需训练即可获得一个新模型。具体来说，我们通过简单地平均它们的权重来合并两个模型。
* 最短拒绝采样：我们观察到我们的模型对同一问题生成的响应长度差异很大。基于此，我们设计了最短拒绝采样方法。该方法对同一问题采样n次（在我们的实验中，n=8），并选择最短的正确响应进行监督微调。
* DPO：与最短拒绝采样类似，我们利用长CoT模型生成多个响应样本。选择最短的正确解作为正样本，而较长的响应被视为负样本，包括错误的较长响应和正确的较长响应（比所选正样本长1.5倍）。这些正负样本对构成了用于DPO训练的成对偏好数据。
* Long2short RL：在标准的RL训练阶段之后，我们选择一个在性能和token效率之间达到最佳平衡的模型作为基础模型，并进行一个独立的long2short RL训练阶段。在第二阶段中，我们应用2.3.3节介绍的长度惩罚，并显著减少最大展开长度，以进一步惩罚那些可能正确但超过期望长度的响应。

2.5 其他训练细节

2.5.1 预训练

Kimi k1.5预训练概述。Kimi k1.5基础模型在一个多样化、高质量的多模态语料库上进行训练。语言数据涵盖五个领域：英语、中文、代码、数学推理和知识。多模态数据包括字幕、图文交错、OCR、知识和问答数据集，使我们的模型能够获得视觉-语言能力。严格的质量控制确保了整个预训练数据集的相关性、多样性和平衡性。我们的预训练分为三个阶段：（1）视觉-语言预训练，首先建立强大的语言基础，然后逐步进行多模态整合；（2）冷却（Cooldown），使用精选和合成数据巩固能力，特别是推理和知识型任务；（3）长上下文激活，将序列处理扩展到131,072个token。关于我们预训练工作的更多细节可以在附录B中找到。

2.5.2 常规监督微调

SFT数据集构建。我们创建了覆盖多个领域的常规SFT语料库。对于非推理任务，包括问答、写作和文本处理，我们首先通过人工标注构建一个种子数据集。这个种子数据集用于训练一个种子模型。随后，我们收集了多样化的提示，并使用种子模型为每个提示生成多个响应。标注员然后对这些响应进行排序，并优化排名最高的响应以产生最终版本。对于数学和编码问题等推理任务，基于规则和奖励模型的验证比人工判断更准确高效，我们利用拒绝采样来扩充SFT数据集。

SFT数据集构成与训练细节。我们的常规SFT数据集包含约100万个文本样本。具体来说，50万个样本用于通用问答，20万用于编码，20万用于数学和科学，5千用于创意写作，2万用于长上下文任务，如摘要、文档问答、翻译和写作。此外，我们构建了100万个图文样本，涵盖图表解读、OCR、图像驱动的对话、视觉编码、视觉推理以及带视觉辅助的数学/科学问题等多种类别。我们首先在32k token的序列长度下训练模型1个epoch，然后在128k token的序列长度下再训练1个epoch。在第一阶段（32k），学习率从 $2 \times 10^{-5}$ 衰减到 $2 \times 10^{-6}$，然后在第二阶段（128k）重新预热到 $1 \times 10^{-5}$，最终衰减到 $1 \times 10^{-6}$。为了提高训练效率，我们将多个训练样本打包到每个单独的训练序列中。

2.6 RL基础设施

图3：LLM的大规模强化学习训练系统
图3：LLM的大规模强化学习训练系统

2.6.1 LLM的大规模强化学习训练系统

Kimi k1.5的同步RL框架。在人工智能领域，强化学习（RL）已成为大型语言模型（LLM）的关键训练方法【索引36，Ouyang et al.，“Training language models to follow instructions with human feedback”，2022】；【索引16，Jaech et al.，“Openai o1 system card”，2024】，其灵感来源于AlphaGo【索引46，Silver et al.，“Mastering the game of go without human knowledge”，2017】、AlphaStar【索引59，Vinyals et al.，“Grandmaster level in StarCraft II using multi-agent reinforcement learning”，2019】和OpenAI Dota Five【索引4，Berner et al.，“Dota 2 with large scale deep reinforcement learning”，2019】等系统在掌握Go、星际争霸II和Dota 2等复杂游戏中的成功。Kimi k1.5系统沿袭这一传统，采用了一个迭代同步的RL框架，旨在通过持续学习和适应来增强模型的推理能力。该系统的一项关键创新是引入了部分展开（Partial Rollout）技术，旨在优化对复杂推理轨迹的处理。

RL系统迭代流程。如图3a所示，RL训练系统通过迭代同步的方式运作，每次迭代包括一个展开阶段和一个训练阶段。在展开阶段，由中央主节点协调的展开工作节点通过与模型交互生成展开轨迹，即对各种输入的响应序列。这些轨迹随后被存储在回放缓冲区中，通过打乱时间相关性来确保训练数据集的多样性和无偏性。在接下来的训练阶段，训练工作节点访问这些经验来更新模型的权重。这个循环过程使模型能够不断从其行动中学习，并随时间调整其策略以提高性能。

系统各组件的角色。中央主节点作为中心指挥者，管理展开工作节点、训练工作节点、奖励模型评估和回放缓冲区之间的数据流和通信。它确保系统和谐运作，平衡负载并促进高效的数据处理。训练工作节点访问这些展开轨迹（无论是在一次迭代中完成还是分多次迭代完成），以计算梯度更新，从而优化模型参数并提升性能。此过程由奖励模型监督，该模型评估模型输出的质量并提供关键反馈来指导训练过程。奖励模型的评估在确定模型策略的有效性并将模型引向最佳性能方面尤为关键。

代码执行服务。此外，该系统集成了一个代码执行服务，专门用于处理与代码相关的问题，并且是奖励模型的组成部分。该服务在实际编码场景中评估模型的输出，确保模型的学习与真实世界的编程挑战紧密对齐。通过对照实际代码执行来验证模型的解决方案，这个反馈循环对于优化模型的策略和提升其在代码相关任务中的性能至关重要。

2.6.2 用于长CoT RL的部分展开（Partial Rollouts）

部分展开技术。我们工作的一个主要思想是扩展长上下文RL训练。部分展开是一项关键技术，它通过管理长短轨迹的展开，有效解决了处理长CoT特征的挑战。该技术设定了一个固定的输出token预算，限制了每个展开轨迹的长度。如果一个轨迹在展开阶段超过了token限制，未完成的部分将被保存到回放缓冲区，并在下一次迭代中继续。这确保了没有任何一个冗长的轨迹会独占系统资源。此外，由于展开工作节点是异步操作的，当一些节点在处理长轨迹时，其他节点可以独立处理新的、较短的展开任务。这种异步操作通过确保所有展开工作节点都在积极为训练过程做出贡献，从而最大化了计算效率，优化了系统的整体性能。

部分展开系统的工作原理。如图3b所示，部分展开系统通过在多次迭代中（从iter n-m到iter n）将长响应分解为多个片段来工作。回放缓冲区作为中央存储机制，维护这些响应片段，其中只有当前迭代（iter n）需要进行在策略（on-policy）计算。之前的片段（iter n-m到n-1）可以从缓冲区中高效地重用，从而无需重复展开。这种分段方法显著降低了计算开销：系统不是一次性展开整个响应，而是增量地处理和存储片段，从而能够在保持快速迭代时间的同时生成更长的响应。在训练期间，可以从损失计算中排除某些片段以进一步优化学习过程，使得整个系统既高效又可扩展。

重复检测功能。部分展开的实现还提供了重复检测功能。系统能识别生成内容中的重复序列并提前终止它们，减少了不必要的计算，同时保持了输出质量。检测到的重复内容可以被赋予额外的惩罚，从而有效地抑制了在提示集中生成冗余内容。

2.6.3 训练与推理的混合部署

混合部署流程。RL训练过程包括以下阶段：

图4：混合部署框架
图4：混合部署框架

训练阶段：开始时，Megatron【索引45，Shoeybi et al.，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism，2020】和vLLM【索引21，Kwon et al.，“Efficient Memory Management for Large Language Model Serving with PagedAttention”，2023】在不同的容器中执行，由一个名为checkpoint-engine的垫片进程封装（见2.6.3节）。Megatron开始训练过程。训练完成后，Megatron卸载GPU内存，并准备将当前权重转移给vLLM。
* 推理阶段：Megatron卸载后，vLLM以虚拟模型权重启动，并用从Megatron通过Mooncake【索引40，Qin et al.，Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving，2024】传输过来的最新权重进行更新。展开完成后，checkpoint-engine停止所有vLLM进程。
* 后续训练阶段*：一旦分配给vLLM的内存被释放，Megatron重新加载内存并开始新一轮的训练。

现有工作的挑战。我们发现现有工作难以同时支持以下所有特性：
* 复杂的并行策略：Megatron可能与vLLM有不同的并行策略。在Megatron中分布在多个节点上的训练权重可能难以与vLLM共享。
* 最小化闲置GPU资源：对于在策略RL，最近的工作如SGLang【索引69，L. Zheng et al.，SGLang: Efficient Execution of Structured Language Model Programs，2024】和vLLM可能会在训练过程中保留一些GPU，这反过来可能导致训练GPU闲置。在训练和推理之间共享相同的设备会更高效。
* 动态扩展能力：在某些情况下，通过增加推理节点的数量同时保持训练过程不变，可以实现显著的加速。我们的系统能够在需要时高效利用闲置的GPU节点。

混合部署框架实现。如图4所示，我们在Megatron和vLLM之上实现了这个混合部署框架（2.6.3节），实现了从训练到推理阶段的切换时间小于一分钟，反向切换约十秒。

混合部署策略。我们提出了一种用于训练和推理任务的混合部署策略，该策略利用Kubernetes Sidecar容器共享所有可用的GPU，将两种工作负载并置于一个pod中。该策略的主要优点是：
* 促进高效的资源共享和管理，当训练和推理节点部署在不同节点上时，可以防止训练节点在等待推理节点时闲置。
* 利用不同的部署镜像，训练和推理可以各自独立迭代以获得更好的性能。
* 该架构不限于vLLM，其他框架可以方便地集成。

Checkpoint Engine。Checkpoint Engine负责管理vLLM进程的生命周期，暴露HTTP API以触发vLLM上的各种操作。为保证整体一致性和可靠性，我们利用由etcd服务管理的全局元数据系统来广播操作和状态。由于CUDA图、NCCL缓冲区和NVIDIA驱动程序的原因，完全由vLLM卸载来释放GPU内存可能具有挑战性。为了最小化对vLLM的修改，我们在需要时终止并重启它，以实现更好的GPU利用率和容错能力。Megatron中的工作节点将其拥有的检查点转换为Hugging Face格式并存放在共享内存中。此转换还考虑了流水线并行和专家并行，因此这些检查点中只剩下张量并行。共享内存中的检查点随后被分成多个分片并在全局元数据系统中注册。我们使用Mooncake通过RDMA在对等节点之间传输检查点。需要对vLLM进行一些修改以加载权重文件并执行张量并行转换。

2.6.4 代码沙箱

沙箱设计与功能。我们开发了沙箱作为一个安全环境，用于执行用户提交的代码，并针对代码执行和代码基准评估进行了优化。通过动态切换容器镜像，沙箱支持不同的用例，包括MultiPL-E【索引6，Cassano, Gouwar, D. Nguyen, S. Nguyen, et al.，“MultiPL-E: A Scalable and Polyglot Approach to Benchmarking Neural Code Generation”，2023】、DMOJ Judge Server 2、Lean、Jupyter Notebook等镜像。对于编码任务中的RL，沙箱通过提供一致且可重复的评估机制，确保了训练数据判断的可靠性。其反馈系统支持多阶段评估，如代码执行反馈和仓库级编辑，同时保持统一的上下文，以确保跨编程语言的公平、公正的基准比较。

部署与优化。我们在Kubernetes上部署该服务以实现可扩展性和弹性，并通过HTTP端点暴露以供外部集成。Kubernetes的自动重启和滚动更新等功能确保了可用性和容错能力。为了优化性能并支持RL环境，我们在代码执行服务中融入了多种技术以提高效率、速度和可靠性。这些技术包括：
* 使用Crun：我们使用crun作为容器运行时而不是Docker，显著减少了容器启动时间。
* Cgroup复用：我们为容器预创建cgroup，这在高并发场景中至关重要，因为为每个容器创建和销毁cgroup可能成为瓶颈。
* 磁盘使用优化：使用一个上层挂载为tmpfs的覆盖文件系统来控制磁盘写入，提供了一个固定大小、高速的存储空间。这种方法对于临时性工作负载很有利。

容器启动时间
容器启动时间

16核机器上每秒启动的最大容器数
16核机器上每秒启动的最大容器数

这些优化提高了代码执行中RL的效率，为评估RL生成的代码提供了一个一致可靠的环境，这对于迭代训练和模型改进至关重要。

A4 实验环境

A4 实验结果

A5 结论

本文展示了Kimi k1.5的训练配方和系统设计。核心洞见之一是上下文长度的扩展对LLM的持续改进至关重要，通过优化的学习算法和部分展开等基础设施优化实现了高效的长上下文RL训练。未来，如何进一步提升长上下文RL训练的效率和可扩展性仍是重要问题。

另一项贡献是实现改进策略优化的一系列技术组合，包括为长CoT RL建立公式并推导出一种在线镜像下降的变体，同时结合采样策略、长度惩罚和数据配方优化。研究表明，仅通过长上下文扩展和改进的策略优化，即使不使用蒙特卡洛树搜索、价值函数等复杂技术，也能实现强大性能。未来研究方向包括改进信用分配和在不损害模型探索能力的情况下减少“过度思考”。

此外，本文还观察到long2short方法的潜力，这些方法能大幅提升短CoT模型的性能。通过将long2short方法与长CoT RL迭代结合，有望进一步提高token效率，并在给定的上下文长度预算内发挥出最佳性能。

A6 附录

B 预训练

预训练数据的重要性。强化学习（RL）的效率与基础模型的性能密切相关。前沿模型如Gemini和Llama凸显了预训练数据质量在实现高性能方面的重要性。本节主要关注多模态预训练数据配方，并简要讨论模型架构和训练阶段。

B.1 语言数据

语料库构成及质量控制。我们的预训练语料库旨在为训练大型语言模型（LLM）提供全面和高质量的数据。它涵盖五个领域：英语、中文、代码、数学与推理以及知识。我们对每个领域都采用了复杂的过滤和质量控制机制，以确保最高质量的训练数据。对于所有预训练数据，我们对每个数据源都进行了严格的独立验证，以评估其对整体训练配方的具体贡献。

中英文文本数据处理。我们开发了一个多维度的质量过滤框架，该框架结合了多种评分方法以减少个体偏见并确保全面的质量评估。我们的框架包括：
1. 基于规则的过滤：我们实施特定领域的启发式规则来移除有问题的内容，包括重复内容、机器翻译文本和低质量的网络抓取内容。我们还过滤掉含有过多特殊字符、不寻常格式或垃圾邮件模式的文档。
2. 基于FastText的分类：我们训练了专门的FastText模型【索引18，Joulin et al.，“Bag of tricks for efficient text classification”，2016】；【索引24，J. Li et al.，“Datacomp-lm: In search of the next generation of training sets for language models”，2024】，根据语言特征和语义连贯性来识别内容质量。这有助于识别具有自然语言流畅性和正确语法结构的文档。
3. 基于嵌入的相似性分析：使用文档嵌入【索引7，Jianlv Chen et al.，“Bge m3-embedding: Multi-lingual, multi-functionality, multi-granularity text embeddings through self-knowledge distillation”，2024】，我们计算文档级别的相似度得分，以识别和移除近似重复的内容，同时保留有语义价值的变体。这种方法有助于维持我们训练语料库的多样性。
4. 基于LLM的质量评估：遵循【索引39，Penedo et al.，“The fineweb datasets: Decanting the web for the finest text data at scale”，2024】的方法，我们利用LLM根据连贯性、信息量和潜在教育价值对文档进行评分。这种方法在识别简单方法可能错过的细微质量指标方面特别有效。
每个文档的最终质量得分是这些独立得分的组合。基于广泛的实