Kimi k1.5：使用LLM扩展强化学习
Kimi k1.5 技术报告
摘要
通过下一个词预测进行语言模型预训练已被证明在扩展计算能力方面是有效的，但它受限于可用训练数据的数量。扩展强化学习（RL）为人工智能的持续改进开辟了一个新的方向，其承诺在于大型语言模型（LLM）可以通过奖励学习来探索，从而扩展其训练数据。然而，之前发表的工作并未产生具有竞争力的结果。鉴于此，我们报告了Kimi k1.5的训练实践，这是我们最新的多模态LLM，使用RL进行训练，包括其RL训练技术、多模态数据配方和基础设施优化。长上下文扩展和改进的策略优化方法是我们方法的关键组成部分，它建立了一个简单而有效的RL框架，而不依赖于更复杂的技术，如蒙特卡洛树搜索、价值函数和过程奖励模型。值得注意的是，我们的系统在多个基准测试和模态上实现了最先进的推理性能——例如，AIME得分为77.5，MATH 500得分为96.2，Codeforces百分位数为94，MathVista得分为74.9——与OpenAI的o1模型相当。此外，我们提出了有效的长到短方法，使用长链思维（long-CoT）技术来改进短链思维（short-CoT）模型，从而实现了最先进的短链思维推理结果——例如，AIME得分为60.8，MATH500得分为94.6，LiveCodeBench得分为47.3——大幅超越了现有的短链思维模型，如GPT-4o和Claude Sonnet 3.5（提升高达+550%）。Kimi k1.5的服务将在kimi.ai上很快推出。
图1：Kimi k1.5的长链思维结果。
1 引言
在扩展定律的背景下，通过下一个词预测进行语言模型预训练已被广泛研究，其中按比例扩展模型参数和数据规模可以持续提升智能水平（Kaplan等，2020；Hoffmann等，2022）。然而，这种方法受限于可用高质量训练数据的数量（Villalobos等，2024；Muennighoff等，2023）。在本报告中，我们介绍了Kimi k1.5的训练配方，这是我们最新的多模态LLM，使用强化学习（RL）进行训练。目标是探索一种可能的新的扩展方向。通过使用RL与LLM，模型可以通过奖励学习进行探索，从而不受限于预先存在的静态数据集。
关于k1.5的设计和训练，有几个关键要素：
长上下文扩展：我们将RL的上下文窗口扩展到128k，并观察到随着上下文长度的增加，性能持续提升。我们方法的一个关键思想是使用部分rollout来提高训练效率——即通过重用大部分先前的轨迹来采样新的轨迹，避免从头重新生成新轨迹的成本。我们的观察表明，上下文长度是RL与LLM持续扩展的一个关键维度。
改进的策略优化：我们推导了长链思维RL的公式，并采用了一种在线镜像下降的变体来进行鲁棒策略优化。该算法通过我们有效的采样策略、长度惩罚和数据配方的优化进一步改进。
简化框架：长上下文扩展与改进的策略优化方法相结合，建立了一个简化的RL框架，用于LLM的学习。由于我们能够扩展上下文长度，学习到的链式思维表现出规划、反思和修正的特性。增加的上下文长度具有增加搜索步数的效果。因此，我们展示了即使不依赖更复杂的技术，如蒙特卡洛树搜索、价值函数和过程奖励模型，也能实现强大的性能。
多模态：我们的模型在文本和视觉数据上联合训练，具备跨两种模态的联合推理能力。
此外，我们提出了有效的长到短方法，使用长链思维技术来改进短链思维模型。具体来说，我们的方法包括应用长链思维激活的长度惩罚和模型合并。
我们的长链思维版本在多个基准测试和模态上实现了最先进的推理性能——例如，AIME得分为77.5，MATH 500得分为96.2，Codeforces百分位数为94，MathVista得分为74.9——与OpenAI的o1模型相当。我们的模型还实现了最先进的短链思维推理结果——例如，AIME得分为60.8，MATH500得分为94.6，LiveCodeBench得分为47.3——大幅超越了现有的短链思维模型，如GPT-4o和Claude Sonnet 3.5（提升高达+550%）。结果如图1和图2所示。
图2：Kimi k1.5的短链思维结果。
2 方法：使用LLM进行强化学习
Kimi k1.5的开发包括几个阶段：预训练、普通的监督微调（SFT）、长链思维监督微调和强化学习（RL）。本报告重点介绍RL，首先概述RL提示集的构建（第2.1节）和长链思维监督微调（第2.2节），然后在第2.3节深入讨论RL训练策略。关于预训练和普通监督微调的更多细节可以在第2.5节中找到。
RL提示集构建
通过初步实验，我们发现RL提示集的质量和多样性在确保强化学习的有效性方面起着关键作用。一个构建良好的提示集不仅能够引导模型进行鲁棒的推理，还能降低奖励黑客和过度拟合表面模式的风险。具体来说，高质量的RL提示集具有以下三个关键属性：
多样化的覆盖范围：提示应涵盖广泛的学科，如STEM、编码和一般推理，以增强模型的适应性并确保其在不同领域的广泛应用。
平衡的难度：提示集应包括分布均匀的简单、中等和困难问题，以促进渐进式学习，并防止过度拟合特定复杂度的任务。
准确的可评估性：提示应允许验证者进行客观和可靠的评估，确保模型的表现基于正确的推理，而不是表面模式或随机猜测。
为了实现提示集的多样化覆盖，我们使用自动过滤器来选择需要丰富推理且易于评估的问题。我们的数据集包括来自各个领域的问题，如STEM领域、竞赛和一般推理任务，涵盖了纯文本和图像-文本问答数据。此外，我们开发了一个标签系统，按领域和学科对提示进行分类，确保不同学科领域的平衡表示（M. Li等，2023；W. Liu等，2023）。
我们采用了一种基于模型的方法，利用模型自身的能力自适应地评估每个提示的难度。具体来说，对于每个提示，SFT模型使用相对较高的采样温度生成十次答案。然后计算通过率，并将其作为提示难度的代理——通过率越低，难度越高。这种方法使难度评估与模型的内在能力保持一致，从而使其在RL训练中非常有效。通过这种方法，我们可以预先过滤大多数简单案例，并在RL训练期间轻松探索不同的采样策略。
为了避免潜在的奖励黑客问题（Everitt等，2021；Pan等，2022），我们需要确保每个提示的推理过程和最终答案都能被准确验证。经验观察表明，一些复杂的推理问题可能具有相对简单且易于猜测的答案，从而导致错误的验证——即模型通过错误的推理过程得出正确答案。为了解决这个问题，我们排除了容易出错的题目，如选择题、判断题和证明题。此外，对于一般的问答任务，我们提出了一种简单但有效的方法来识别和删除容易被黑客攻击的提示。具体来说，我们提示模型在没有链式思维推理步骤的情况下猜测潜在答案。如果模型在$N$次尝试内预测出正确答案，则该提示被认为太容易被黑客攻击并被删除。我们发现，设置$N=8$可以删除大多数容易被黑客攻击的提示。开发更高级的验证模型仍然是未来研究的一个开放方向。
长链思维监督微调
通过精炼的RL提示集，我们使用提示工程构建了一个小而高质量的长链思维预热数据集，其中包含经过准确验证的文本和图像输入的推理路径。这种方法类似于拒绝采样（RS），但侧重于通过提示工程生成长链思维推理路径。生成的预热数据集旨在封装人类推理的关键认知过程，如规划，即模型在执行前系统地概述步骤；评估，涉及对中间步骤的批判性评估；反思，使模型能够重新考虑并改进其方法；以及探索，鼓励考虑替代解决方案。通过对该预热数据集进行轻量级的SFT，我们有效地引导模型内化这些推理策略。结果，经过微调的长链思维模型在生成更详细和逻辑一致的响应方面表现出改进的能力，从而提高了其在多样化推理任务中的表现。
强化学习
2.3.1 问题设置
给定一个训练数据集$\mathcal{D}={(x_{i},y^{}{i})}^{n}$，其中$x_{i}$是问题，$y^{}{i}$是对应的真实答案，我们的目标是训练一个策略模型$\pi{\theta}$，以准确解决测试问题。在复杂推理的背景下，从问题$x$到解决方案$y$的映射是非平凡的。为了应对这一挑战，链式思维（CoT）方法提出使用一系列中间步骤$z=(z_{1},z_{2},\ldots,z_{m})$来桥接$x$和$y$，其中每个$z_{i}$是一个连贯的标记序列，作为解决问题的重要中间步骤（J. Wei等，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}$表示此采样过程。请注意，思维和最终答案都是作为语言序列采样的。
为了进一步增强模型的推理能力，规划算法被用于探索各种思维过程，在推理时生成改进的链式思维（Yao等，2024；Y. Wu等，2024；Snell等，2024）。这些方法的核心见解是显式构建一个由价值估计引导的思维搜索树。这使得模型能够探索思维过程的不同延续，或在遇到死胡同时回溯以调查新的方向。更详细地说，设$\mathcal{T}$为一个搜索树，其中每个节点表示一个部分解决方案$s=(x,z_{1:|s|})$。这里$s$由问题$x$和一系列思维$z_{1:|s|}=(z_{1},\ldots,z_{|s|})$组成，$|s|$表示序列中的思维数量。规划算法使用一个批评模型$v$来提供反馈$v(x,z_{1:|s|})$，这有助于评估当前解决问题的进展并识别现有部分解决方案中的任何错误。我们注意到，反馈可以由判别分数或语言序列提供（L. Zhang等，2024）。在所有$s\in\mathcal{T}$的反馈指导下，规划算法选择最有希望的节点进行扩展，从而增长搜索树。上述过程迭代重复，直到得出完整的解决方案。
我们也可以从算法角度来看待规划算法。给定在第$t$次迭代时可用的过去搜索历史$(s_{1},v(s_{1}),\ldots,s_{t-1},v(s_{t-1}))$，规划算法$\mathcal{A}$迭代地确定下一个搜索方向$\mathcal{A}(s_{t}|s_{1},v(s_{1}),\ldots,s_{t-1},v(s_{t-1}))$，并为当前搜索进度提供反馈$\mathcal{A}(v(s_{t})|s_{1},v(s_{1}),\ldots,s_{t})$。由于思维和反馈都可以被视为中间推理步骤，并且这些组件都可以表示为语言标记序列，我们使用$z$替换$s$和$v$以简化符号。因此，我们将规划算法视为直接作用于推理步骤序列$\mathcal{A}(\cdot|z_{1},z_{2},\ldots)$的映射。在这个框架中，规划算法使用的搜索树中存储的所有信息都被扁平化为提供给算法的完整上下文。这为生成高质量的链式思维提供了一个有趣的视角：与其显式构建搜索树并实现规划算法，我们可能可以训练一个模型来近似这个过程。在这里，思维的数量（即语言标记的数量）与传统上分配给规划算法的计算预算类似。最近的长上下文窗口的进展在训练和测试阶段都促进了无缝扩展。如果可行，这种方法使模型能够通过自回归预测直接在推理空间中进行隐式搜索。因此，模型不仅学会解决一组训练问题，还发展出有效解决单个问题的能力，从而提高了对未见测试问题的泛化能力。
因此，我们考虑使用强化学习（RL）训练模型生成链式思维（OpenAI，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)$生成链式思维和最终答案。生成的链式思维的质量通过其是否能够引导出正确的最终答案来评估。总之，我们考虑以下目标来优化策略：
$$ \max_{\theta}\mathbb{E}{(x,y^{*})\sim\mathcal{D},(y,z)\sim\pi{\theta} }\left[r(x,y,y^{*})\right],. $$
通过扩展RL训练，我们的目标是训练一个模型，该模型结合了简单提示链式思维和规划增强链式思维的优势。模型在推理时仍然自回归采样语言序列，从而避免了部署时高级规划算法所需的复杂并行化需求。然而，与简单的提示方法的一个关键区别在于，模型不应仅仅遵循一系列推理步骤。相反，它还应通过学习关键规划技能，包括错误识别、回溯和解决方案改进，利用所有探索过的思维作为上下文信息。
2.3.2 策略优化
我们应用了一种在线策略镜像下降的变体作为我们的训练算法（Abbasi-Yadkori等，2019；Mei等，2019；Tomar等，2020）。该算法迭代执行。在第$i$次迭代时，我们使用当前模型$\pi_{\theta_{i} }$作为参考模型，并优化以下相对熵正则化策略优化问题：
$$ \max_{\theta}\mathbb{E}{(x,y^{*})\sim\mathcal{D} }\left[\mathbb{E}{(y,z)\sim\pi_{\theta} }\left[r(x,y,y^{*})\right]-\tau\mathrm{KL}(\pi_{\theta}(x)||\pi_{\theta_{i} }(x))\right],, $$
其中$\tau>0$是控制正则化程度的参数。该目标有一个闭式解：
$$ \pi^{}(y,z|x)=\pi_{\theta_{i} }(y,z|x)\exp(r(x,y,y^{})/\tau)/Z,. $$
这里$Z=\sum_{y^{\prime},z^{\prime} }\pi_{\theta_{i} }(y^{\prime},z^{\prime}|x)\exp(r(x,y^{\prime},y^{*})/\tau)$是归一化因子。对两边取对数，我们得到对于任何$(y,z)$，以下约束成立，这允许我们在优化过程中利用离策略数据：
$$ r(x,y,y^{})-\tau\log Z=\tau\log\frac{\pi^{}(y,z|x)}{\pi_{\theta_{i} }(y,z|x)},. $$
这激发了以下替代损失：
$$ L(\theta)=\mathbb{E}{(x,y^{*})\sim\mathcal{D} }\left[\mathbb{E}{(y,z)\sim\pi_{\theta_{i} } }\left[\left(r(x,y,y^{*})-\tau\log Z-\tau\log\frac{\pi_{\theta}(y,z|x)}{\pi_{\theta_{i} }(y,z|x)}\right)^{2}\right]\right],. $$
为了近似$\tau\log Z$，我们使用样本$(y_{1},z_{1}),\ldots,(y_{k},z_{k})\sim\pi_{\theta_{i} }$：$\tau\log Z\approx\tau\log\frac{1}{k}\sum_{j=1}^{k}\exp(r(x,y_{j},y^{})/\tau)$。我们还发现，使用采样奖励的均值$\overline{r}=\mathrm{mean}(r(x,y_{1},y^{}),\ldots,r(x,y_{k},y^{*}))$在实践中产生了有效的结果。这是合理的，因为随着$\tau\rightarrow\infty$，$\tau\log Z$接近$\pi_{\theta_{i} }$下的期望奖励。最后，我们通过对替代损失求梯度来总结我们的学习算法。对于每个问题$x$，使用参考策略$\pi_{\theta_{i} }$采样$k$个响应，梯度由下式给出：
$$ \frac{1}{k}\sum_{j=1}^{k}\left(\nabla_{\theta}\log\pi_{\theta}(y_{j},z_{j}|x)(r(x,y_{j},y^{*})-\overline{r})-\frac{\tau}{2}\nabla_{\theta}\left(\log\frac{\pi_{\theta}(y_{j},z_{j}|x)}{\pi_{\theta_{i} }(y_{j},z_{j}|x)}\right)^{2}\right),. $$
对于熟悉策略梯度方法的人来说，这个梯度类似于使用采样奖励均值作为基线的策略梯度（Kool等，2019；Ahmadian等，2024）。主要区别在于响应是从$\pi_{\theta_{i} }$采样的，而不是在策略上采样的，并且应用了$l_{2}$正则化。因此，我们可以将其视为通常的在策略正则化策略梯度算法在离策略情况下的自然扩展（Nachum等，2017）。我们从$\mathcal{D}$中采样一批问题，并将参数更新为$\theta_{i+1}$，随后将其作为下一次迭代的参考策略。由于每次迭代都考虑不同的优化问题，由于参考策略的变化，我们还在每次迭代开始时重置优化器。
我们在训练系统中排除了价值网络，这在之前的研究中也被利用过（Ahmadian等，2024）。虽然这一设计选择显著提高了训练效率，但我们还假设，在经典RL中用于信用分配的价值函数的传统使用可能不适合我们的上下文。考虑一个场景，模型生成了一个部分链式思维$(z_{1},z_{2},\ldots,z_{t})$，并且有两个潜在的下一个推理步骤：$z_{t+1}$和$z^{\prime}{t+1}$。假设$z$直接导致正确答案，而$z^{\prime}{t+1}$包含一些错误。如果可以访问一个预言价值函数，它将表明$z$相对于$z^{\prime}{t+1}$保留了更高的价值。根据标准的信用分配原则，选择$z^{\prime}$将受到惩罚，因为它相对于当前策略具有负优势。然而，探索$z^{\prime}{t+1}$对于训练模型生成长链式思维非常有价值。通过使用从长链式思维得出的最终答案的合理性作为奖励信号，模型可以从采取$z^{\prime}$中学习试错的模式，只要它成功恢复并达到正确答案。这个例子的关键启示是，我们应该鼓励模型探索多样化的推理路径，以增强其解决复杂问题的能力。这种探索性方法生成了丰富的经验，支持关键规划技能的发展。我们的主要目标不仅限于在训练问题上获得高准确性，而是专注于为模型配备有效的问题解决策略，最终提高其在测试问题上的表现。
2.3.3 长度惩罚
我们观察到在RL训练过程中，模型的响应长度显著增加。虽然这导致了更好的性能，但过长的推理过程在训练和推理期间成本高昂，而且过度思考通常不被人类所偏好。为了解决这个问题，我们引入了一个长度奖励，以抑制标记长度的快速增长，从而提高模型的标记效率。给定问题$x$的$k$个采样响应$(y_{1},z_{1}),\ldots,(y_{k},z_{k})$，其中真实答案为$y^{*}$，设$\mathsf{len}(i)$为$(y_{i},z_{i})$的长度，$\min_\mathsf{len}=\min_{i}\mathsf{len}(i)$，$\max_\mathsf{len}=\max_{i}\mathsf{len}(i)$。如果$\max_\mathsf{len}=\min_\mathsf{len}$，我们将所有响应的长度奖励设为零，因为它们具有相同的长度。否则，长度奖励由下式给出：
$$ \mathsf{len_reward}(i)=\begin{cases}&\lambda\quad\text{如果 }r(x,y_{i},y^{})=1\ \min(0,\lambda)&\text{如果 }r(x,y_{i},y^{})=0\end{cases},\quad\text{其中 }\lambda=0.5-\frac{\mathsf{len}(i)-\min_\mathsf{len} }{\max_\mathsf{len}-\min_\mathsf{len} },. $$
本质上，我们促进较短的响应，并在正确的响应中惩罚较长的响应，同时明确惩罚具有错误答案的长响应。然后，这个基于长度的奖励被添加到原始奖励中，并带有加权参数。
在我们的初步实验中，长度惩罚可能会在初始阶段减缓训练速度。为了缓解这个问题，我们建议在训练过程中逐渐预热长度惩罚。具体来说，我们首先使用标准策略优化而不使用长度惩罚，然后在训练的其余部分应用恒定的长度惩罚。
2.3.4 采样策略
尽管RL算法本身具有良好的采样特性（较难的问题提供较大的梯度），但其训练效率有限。因此，一些定义良好的先验采样方法可能会产生更大的性能提升。我们利用多个信号来进一步改进采样策略。首先，我们收集的RL训练数据自然带有不同的难度标签。例如，数学竞赛问题比小学数学问题更难。其次，由于RL训练过程对同一问题进行多次采样，我们还可以跟踪每个问题的成功率作为难度的度量。我们提出了两种采样方法，以利用这些先验知识来提高训练效率。
课程采样：我们首先在较简单的任务上训练，然后逐渐过渡到更具挑战性的任务。由于初始RL模型的性能有限，将有限的计算预算花费在非常困难的问题上通常会产生很少的正确样本，导致训练效率较低。同时，我们收集的数据自然包括年级和难度标签，使得基于难度的采样成为一种直观且有效的提高训练效率的方法。
优先采样：除了课程采样外，我们还使用优先采样策略，专注于模型表现不佳的问题。我们跟踪每个问题$i$的成功率$s_{i}$，并按$1-s_{i}$的比例采样问题，使得成功率较低的问题获得更高的采样概率。这使模型的努力集中在最薄弱的领域，从而加快学习速度并提高整体性能。
2.3.5 训练配方的更多细节
编码问题的测试用例生成：由于许多来自网络的编码问题没有可用的测试用例，我们设计了一种自动生成测试用例的方法，作为训练我们模型的奖励。我们主要关注不需要特殊评判的问题。我们还假设这些问题的真实解决方案是可用的，以便我们可以利用这些解决方案生成更高质量的测试用例。
我们利用广泛认可的测试用例生成库CYaRon1来增强我们的方法。我们使用基础的Kimi k1.5根据问题描述生成测试用例。CYaRon的使用说明和问题描述作为生成器的输入。对于每个问题，我们首先使用生成器生成50个测试用例，并随机采样10个真实提交的解决方案。我们运行测试用例以验证提交的解决方案。如果至少7个提交的解决方案产生匹配的结果，则测试用例被视为有效。经过这轮过滤后，我们获得了一组选定的测试用例。如果至少9个提交的解决方案通过了整个选定的测试用例集，则将问题及其相关的选定测试用例添加到我们的训练集中。
数学问题的奖励建模：评估数学解决方案的一个挑战是，不同的书写形式可能代表相同的底层答案。例如，$a^{2}-4$和$(a+2)(a-2)$可能是同一问题的有效解决方案。我们采用了两种方法来提高奖励模型的评分准确性：
经典奖励模型：借鉴InstructGPT（Ouyang等，2022）的方法，我们实现了一个基于价值头的奖励模型，并收集了大约800k个数据点进行微调。该模型最终将“问题”、“参考答案”和“响应”作为输入，并输出一个标量，指示响应是否正确。
链式思维奖励模型：最近的研究（Ankner等，2024；McAleese等，2024）表明，增强链式思维（CoT）推理的奖励模型可以显著优于经典方法，特别是在涉及细微正确性标准的任务上——如数学。因此，我们收集了大约800k个带有CoT标签的示例来微调Kimi模型。基于与经典奖励模型相同的输入，链式思维方法显式生成逐步推理过程，然后以JSON格式提供最终的正确性判断，从而实现更鲁棒和可解释的奖励信号。
在我们的手动抽查中，经典奖励模型的准确率约为84.4，而链式思维奖励模型的准确率达到98.5。在RL训练过程中，我们采用了链式思维奖励模型，以确保更正确的反馈。
视觉数据：为了提高模型在现实世界图像推理能力，并实现视觉输入与大型语言模型（LLM）之间更有效的对齐，我们的视觉强化学习（Vision RL）数据主要来自三个不同的类别：现实世界数据、合成视觉推理数据和文本渲染数据。
现实世界数据：包括需要图形理解和推理的各个年级的科学问题、需要视觉感知和推理的位置猜测任务，以及涉及理解复杂图表的数据分析等。这些数据集提高了模型在现实世界场景中进行视觉推理的能力。
合成视觉推理数据：是人工生成的，包括程序生成的图像和场景，旨在提高特定的视觉推理技能，如理解空间关系、几何图案和对象交互。这些合成数据集提供了一个受控的环境来测试模型的视觉推理能力，并提供了无尽的训练示例。
文本渲染数据：通过将文本内容转换为视觉格式创建，使模型在处理跨不同模态的基于文本的查询时保持一致。通过将文本文档、代码片段和结构化数据转换为图像，我们确保模型无论输入是纯文本还是文本渲染为图像（如截图或照片）时都能提供一致的响应。这也有助于增强模型在处理文本密集型图像时的能力。
每种类型的数据对于构建一个全面的视觉语言模型都是必不可少的，该模型可以有效地管理广泛的现实世界应用，同时确保跨各种输入模态的一致性能。
长到短：短链思维模型的上下文压缩
尽管长链思维模型表现出强大的性能，但与标准的短链思维LLM相比，它在测试时消耗了更多的标记。然而，可以将长链思维模型的思维先验转移到短链思维模型中，从而即使在有限的测试标记预算下也能提高性能。我们提出了几种解决长到短问题的方法，包括模型合并（Yang等，2024）、最短拒绝采样、DPO（Rafailov等，2024）和长到短RL。这些方法的详细描述如下：
模型合并：模型合并已被证明在保持泛化能力方面非常有用。我们还发现，在合并长链思维模型和短链思维模型时，它在提高标记效率方面也非常有效。这种方法将长链思维模型与较短的模型结合起来，无需训练即可获得一个新模型。具体来说，我们通过简单地平均它们的权重来合并这两个模型。
最短拒绝采样：我们观察到，我们的模型对同一问题生成的响应长度变化很大。基于此，我们设计了最短拒绝采样方法。该方法对同一问题采样$n$次（在我们的实验中，$n=8$），并选择最短的正确响应进行监督微调。
DPO：与最短拒绝采样类似，我们利用长链思维模型生成多个响应样本。最短的正确解决方案被选为正样本，而较长的响应被视为负样本，包括错误的较长响应和正确的较长响应（比选定的正样本长1.5倍）。这些正负对形成了用于DPO训练的成对偏好数据。
长到短RL：在标准的RL训练阶段之后，我们选择一个在性能和标记效率之间提供最佳平衡的模型作为基础模型，并进行单独的长到短RL训练阶段。在这个第二阶段，我们应用第2.3.3节中引入的长度惩罚，并显著减少最大rollout长度，以进一步惩罚超过所需长度的响应，即使它们可能是正确的。
其他训练细节
2.5.1 预训练
Kimi k1.5基础模型是在多样化、高质量的多模态语料库上进行训练的。语言数据涵盖五个领域：英语、中文、代码、数学推理和知识。多模态数据，包括字幕、图像-文本交错、OCR、知识和问答数据集，使我们的模型能够获得视觉-语言能力。严格的质量控制确保了整个预训练数据集的相关性、多样性和平衡性。我们的预训练分为三个阶段：（1）视觉-语言预训练，建立强大的语言基础，然后逐步引入多模态集成；（2）冷却阶段，使用精选和合成数据巩固能力，特别是推理和基于知识的任务；（3）长上下文激活，将序列处理扩展到131,072个标记。有关我们预训练工作的更多细节可以在附录B中找到。
2.5.2 普通监督微调
我们创建了涵盖多个领域的普通SFT语料库。对于非推理任务，包括问答、写作和文本处理，我们首先通过人工注释构建一个种子数据集。该种子数据集用于训练种子模型。随后，我们收集了多样化的提示，并使用种子模型为每个提示生成多个响应。注释者随后对这些响应进行排名，并完善排名最高的响应以生成最终版本。对于推理任务，如数学和编码问题，基于规则和奖励建模的验证比人工判断更准确和高效，我们使用拒绝采样来扩展SFT数据集。
我们的普通SFT数据集包含大约100万个文本示例。具体来说，50万个示例用于一般问答，20万个用于编码，20万个用于数学和科学，5千个用于创意写作，2万个用于长上下文任务，如摘要、文档问答、翻译和写作。此外，我们构建了100万个文本-视觉示例，涵盖各种类别，包括图表解释、OCR、基于图像的对话、视觉编码、视觉推理以及带有视觉辅助的数学/科学问题。
我们首先在32k标记的序列长度上训练模型1个epoch，然后在128k标记的序列长度上训练另一个epoch。在第一阶段（32k），学习率从$2\times 10^{-5}$衰减到$2\times 10^{-6}$，然后在第二阶段（128k）重新预热到$1\times 10^{-5}$，最后衰减到$1\times 10^{-6}$。为了提高训练效率，我们将多个训练示例打包到每个训练序列中。
RL基础设施
图3：用于LLM的大规模强化学习训练系统
2.6.1 用于LLM的大规模强化学习训练系统
在人工智能领域，强化学习（RL）已成为大型语言模型（LLM）的关键训练方法（Ouyang等，2022）（Jacch等，2024），其灵感来自于通过AlphaGo（Silver等，2017）、AlphaStar（Vinyals等，2019）和OpenAI Dota Five（Berner等，2019）等系统在掌握复杂游戏（如围棋、星际争霸II和Dota 2）方面的成功。遵循这一传统，Kimi k1.5系统采用了一种迭代同步RL框架，精心设计以通过持续学习和适应来增强模型的推理能力。该系统的一个关键创新是引入了部分rollout技术，旨在优化复杂推理轨迹的处理。
如图3a所示，RL训练系统通过迭代同步方法运行，每次迭代包括一个rollout阶段和一个训练阶段。在rollout阶段，rollout工作器由中央主控协调，通过与模型交互生成rollout轨迹，产生对各种输入的响应序列。这些轨迹随后存储在回放缓冲区中，通过破坏时间相关性确保训练数据的多样性和无偏性。在随后的训练阶段，训练工作器访问这些经验以更新模型的权重。这个循环过程使模型能够从其行为中不断学习，随着时间的推移调整其策略以提高性能。
中央主控作为中央指挥，管理rollout工作器、训练工作器、奖励模型评估和回放缓冲区之间的数据流和通信。它确保系统和谐运行，平衡负载并促进高效的数据处理。
训练工作器访问这些rollout轨迹，无论是在单次迭代中完成还是跨多次迭代分割，以计算梯度更新，从而优化模型的参数并提高其性能。这个过程由奖励模型监督，奖励模型评估模型输出的质量并提供必要的反馈以指导训练过程。奖励模型的评估在确定模型策略的有效性和引导模型实现最佳性能方面尤为重要。
此外，系统还集成了一个代码执行服务，专门用于处理与代码相关的问题，并且是奖励模型的重要组成部分。该服务在实际编码场景中评估模型的输出，确保模型的学习与真实世界的编程挑战紧密对齐。通过将模型的解决方案与实际代码执行进行验证，这个反馈循环对于优化模型的策略和提高其在代码相关任务中的表现至关重要。
2.6.2 长链思维RL的部分rollout
我们工作的一个主要思想是扩展长上下文RL训练。部分rollout是一项关键技术，通过管理长和短轨迹的rollout，有效解决了处理长链思维特征的挑战。该技术建立了一个固定的输出标记预算，限制了每个rollout轨迹的长度。如果轨迹在rollout阶段超过了标记限制，未完成的部分将保存到回放缓冲区，并在下一次迭代中继续。这确保了没有单个长轨迹独占系统资源。此外，由于rollout工作器是异步运行的，当一些工作器处理长轨迹时，其他工作器可以独立处理新的、较短的rollout任务。异步操作通过确保所有rollout工作器都积极参与训练过程，从而优化了系统的整体性能。
如图3b所示，部分rollout系统通过将长响应分解为跨迭代的片段（从迭代n-m到迭代n）来工作。回放缓冲区充当中央存储机制，维护这些响应片段，其中只有当前迭代（迭代n）需要策略计算。先前的片段（迭代n-m到n-1）可以有效地从缓冲区中重用，无需重复rollout。这种分段方法显著减少了计算开销：与其一次性rollout整个响应，系统增量处理和存储片段，从而允许生成更长的响应，同时保持快速的迭代时间。在训练期间，某些片段可以从损失计算中排除，以进一步优化学习过程，使整个系统既高效又可扩展。
部分rollout的实现还提供了重复检测。系统识别生成内容中的重复序列并提前终止它们，减少了不必要的计算，同时保持了输出质量。检测到的重复可以被分配额外的惩罚，从而有效抑制提示集中冗余内容的生成。
2.6.3 训练和推理的混合部署
RL训练过程包括以下阶段：
训练阶段：最初，Megatron（Shoeybi等，2020）和vLLM（Kwon等，2023）在单独的容器中执行，由称为checkpoint-engine的shim进程封装（第2.6.3节）。Megatron开始训练过程。训练完成后，Megatron卸载GPU内存，并准备将当前权重传输到vLLM。
推理阶段：在Megatron卸载后，vLLM以虚拟模型权重启动，并使用通过Mooncake（Qin等，2024）从Megatron传输的最新权重进行更新。rollout完成后，checkpoint-engine停止所有vLLM进程。
后续训练阶段：一旦分配给vLLM的内存被释放，Megatron重新加载内存并启动另一轮训练。
我们发现现有的工作难以同时支持以下所有特性。
复杂的并行策略：Megatron可能与vLLM有不同的并行策略。分布在多个节点上的训练权重在Megatron中可能难以与vLLM共享。
最小化空闲GPU资源：对于策略RL，最近的工作如SGLang（L. Zheng等，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（Cassano等，2023）、DMOJ Judge Server2、Lean等（2023）、Jupyter Notebook和其他镜像的不同用例。
对于编码任务中的RL，沙箱通过提供一致和可重复的评估机制来确保训练数据判断的可靠性。其反馈系统支持多阶段评估，如代码执行反馈和仓库级编辑，同时保持统一的上下文以确保跨编程语言的公平和公正的基准比较。
我们将服务部署在Kubernetes上以实现可扩展性和弹性，并通过HTTP端点暴露它以进行外部集成。Kubernetes的自动重启和滚动更新等功能确保了可用性和容错性。
为了优化性能并支持RL环境，我们在代码执行服务中集成了几种技术以提高效率、速度和可靠性。这些技术包括：
使用Crun：我们使用crun作为容器运行时而不是Docker，显著减少了容器启动时间。
Cgroup重用：我们预先为容器使用创建cgroup，这在具有高并发性的场景中至关重要，其中为每个容器创建和销毁cgroup可能成为瓶颈。
磁盘使用优化：使用带有上层挂载为tmpfs的覆盖文件系统来控制磁盘写入，提供固定大小的高速存储空间。这种方法对于短暂的工作负载非常有益。
这些优化提高了代码执行中的RL效率，为评估RL生成的代码提供了一个一致且可靠的环境，这对于迭代训练和模型改进至关重要。
3 实验
评估
由于k1.5是一个多模态模型，我们在不同模态的各种基准上进行了全面评估。详细的评估设置可以在附录C中找到。我们的基准主要包括以下三类：
文本基准：MMLU（Hendrycks等，2020）、IF-Eval（J. Zhou等，2023）、CLUEWSC（L. Xu等，2020）、C-EVAL（Y. Huang等，2023）
推理基准：HumanEval-Mul等（2024）、Codeforces等（2024）、MATH-500（Lightman等，2023）
视觉基准：MMMU（Yue等，2024）、MATH-Vision（K. Wang等，2024）、MathVista（Lu等，2023）
主要结果
K1.5长链思维模型：Kimi k1.5长链思维模型的性能如表2所示。通过长链思维监督微调（第2.2节描述）和视觉-文本联合强化学习（第2.3节讨论），模型的长期推理能力显著增强。测试时计算扩展进一步增强了其性能，使模型能够在一系列模态上实现最先进的结果。我们的评估揭示了模型在扩展上下文中的推理、理解和综合信息能力的显著提升，代表了多模态AI能力的进步。
K1.5短链思维模型：Kimi k1.5短链思维模型的性能如表3所示。该模型集成了多种技术，包括传统的监督微调（第2.5.2节讨论）、强化学习（第2.3节探讨）和长到短蒸馏（第2.4节概述）。结果表明，k1.5短链思维模型在多个任务上表现出与领先的开源和专有模型相当或更优的性能。这些任务包括文本、视觉和推理挑战，在自然语言理解、数学、编码和逻辑推理方面表现出显著的优势。
长上下文扩展
我们使用一个中等大小的模型来研究RL与LLM的扩展特性。图5展示了在数学提示集上训练的小模型变体的训练准确率和响应长度随训练迭代的变化。随着训练的进行，我们观察到响应长度和性能准确率同时增加。值得注意的是，更具挑战性的基准测试显示出响应长度的更陡峭增加，表明模型学会了为复杂问题生成更详细的解决方案。图6表明，模型的输出上下文长度与其解决问题的能力之间存在强相关性。我们最终的k1.5运行扩展到128k上下文长度，并在困难的推理基准上观察到持续改进。
长到短
我们将提出的长到短RL算法与第2.4节中介绍的DPO、最短拒绝采样和模型合并方法进行了比较，重点关注长到短问题的标记效率（X. Chen等，2024），特别是获得的长链思维模型如何使短模型受益。在图7中，k1.5-long代表我们选择用于长到短训练的长链思维模型。k1.5-short w/ rl指的是使用长到短RL训练获得的短模型。k1.5-short w/ dpo表示通过DPO训练提高标记效率的短模型。k1.5-short w/ merge表示合并后的模型，而k1.5-short w/ merge + rs表示通过对合并模型应用最短拒绝采样获得的短模型。k1.5-shortest表示我们在长到短训练中获得的最短模型。如图7所示，与DPO和模型合并等其他方法相比，提出的长到短RL算法表现出最高的标记效率。值得注意的是，k1.5系列中的所有模型（标记为橙色）都表现出比其他模型（标记为蓝色）更优的标记效率。例如，k1.5-short w/ rl在AIME2024上实现了60.8的Pass@1得分（8次运行的平均值），同时平均仅使用3,272个标记。同样，k1.5-shortest在MATH500上实现了88.2的Pass@1得分，同时消耗的标记数量与其他短模型大致相同。
消融研究
模型大小和上下文长度的扩展：我们的主要贡献是应用RL来增强模型生成扩展链式思维的能力，从而提高其推理能力。一个自然的问题是：这与简单地增加模型大小相比如何？为了证明我们方法的有效性，我们使用相同的数据集训练了两个不同大小的模型，并记录了RL训练期间所有检查点的评估结果和平均推理长度。这些结果如图8所示。值得注意的是，尽管较大的模型最初表现优于较小的模型，但较小的模型可以通过利用通过RL优化的较长链式思维实现相当的性能。然而，较大的模型通常表现出比小模型更好的标记效率。这也表明，如果目标是实现最佳性能，扩展较大模型的上下文长度具有更高的上限，并且更节省标记。然而，如果测试时计算有预算，训练具有较大上下文长度的较小模型可能是可行的解决方案。
使用负梯度的效果：我们研究了在我们的设置中使用ReST（Gulcehre等，2023）作为策略优化算法的有效性。ReST与其他基于RL的方法（包括我们的方法）的主要区别在于，ReST通过拟合从当前模型采样的最佳响应来迭代优化模型，而不应用负梯度来惩罚错误的响应。如图10所示，与ReST相比，我们的方法表现出更优的样本复杂性，表明负梯度的加入显著提高了模型生成长链式思维的效率。我们的方法不仅提高了推理的质量，还优化了训练过程，以更少的训练样本实现了鲁棒的性能。这一发现表明，在我们的设置中，选择适当的优化策略对于最大化生成长链式思维的有效性至关重要。
采样策略：我们进一步证明了我们课程采样策略的有效性，如第2.3.4节所述。我们的训练数据集$\mathcal{D}$包含具有不同难度的问题。通过我们的课程采样方法，我们首先使用$\mathcal{D}$进行预热阶段，然后专注于难题来训练模型。这种方法与不进行任何课程调整的均匀采样基线方法进行了比较。如图9所示，我们的结果清楚地表明，提出的课程采样方法显著提高了性能。这种改进可以归因于该方法逐步挑战模型的能力，使其在处理复杂问题时发展出更强大的理解和能力。通过在初步介绍后专注于更困难的问题，模型可以更好地加强其推理和问题解决能力。
4 结论
我们介绍了k1.5的训练配方和系统设计，这是我们最新的多模态LLM，使用RL进行训练。我们从实践中提取的一个关键见解是，上下文长度的扩展对于LLM的持续改进至关重要。我们采用了优化的学习算法和基础设施优化（如部分rollout）来实现高效的长上下文RL训练。如何进一步提高长上下文RL训练的效率和可扩展性仍然是一个重要的问题。
我们做出的另一个贡献是结合了多种技术，以实现改进的策略优化。具体来说，我们制定了长链思维RL与LLM的公式，并推导了一种在线镜像下降的变体以进行鲁棒优化。我们还通过采样策略、长度惩罚和优化数据配方进行了实验，以实现强大的RL性能。
我们展示了通过长上下文扩展和改进的策略优化，即使不使用更复杂的技术（如蒙特卡洛树搜索、价值函数和过程奖励模型），也能实现强大的性能。未来，研究如何在不损害模型探索能力的情况下改进信用分配和减少过度思考也将是一个有趣的方向。
我们还观察到了长到短方法的潜力。这些方法在很大程度上提高了短链思维模型的性能。此外，可以将长到短方法与长链思维RL以迭代方式结合，以进一步提高标记效率，并在给定的上下文长度预算内提取最佳性能。