Kimi k1.5技术报告速读
1. 强大的多模态推理能力:
Kimi k1.5 模型在多个基准测试和模态上实现了最先进的推理性能,包括数学、代码、文本和视觉推理等任务。
它不仅能处理纯文本,还能理解图像和文本的组合,实现真正的多模态推理。
long-CoT 版本和 short-CoT 版本都展现了强大的性能。
在AIME、MATH 500、Codeforces等数据集上均达到了领先水平。
在视觉基准测试如MathVista上也有较强的表现。
2. 长上下文 RL 训练的突破:
通过扩展上下文窗口到 128k,模型性能持续提高,验证了上下文长度是强化学习扩展的关键维度。
- 采用了部分 rollout 技术,显著提高了长上下文 RL 训练的效率,使得训练更长、更复杂的推理过程成为可能。
- 证明了长上下文对复杂问题的解决能力至关重要。
3. 有效的 long2short 方法:
- 提出了有效的 long2short 方法,利用长上下文模型的知识来改进短上下文模型的性能,并提高 Token 效率。
- 通过模型合并、最短拒绝采样、DPO 等技术,实现了短模型的性能提升,并且比直接训练短模型更高效。
证明了可以通过迁移 long-CoT 模型的思维先验知识来提高短模型的性能。
4. 优化的 RL 训练框架:
提出了改进的策略优化方法,无需依赖复杂技术,即可实现强大的性能。
探索了多种采样策略、长度惩罚和数据配方优化,提升了 RL 训练的效率和效果。
通过消融研究,验证了所提出的方法和训练技巧的有效性。
5. 详细的训练过程和系统设计:
- 公开了 Kimi k1.5 的详细训练配方,包括预训练、微调和强化学习阶段。
- 介绍了混合部署策略,优化训练和推理过程中的资源利用率。
- 详细介绍了数据管道和质量控制机制,确保训练数据的质量。
- 开发了安全的代码沙箱,用于执行和评估生成的代码。
6. 模型规模与上下文长度的探索:
研究了模型规模和上下文长度对性能的影响,发现规模较大的模型性能更好,但使用更长上下文的小型模型也能接近甚至媲美大模型的性能,证明了 RL 训练对模型性能的增强。
摘要
使用下一个 Token 预测进行语言模型预训练已被证明可以有效地扩展计算量,但受限于可用的训练数据量。扩展强化学习 (RL) 为人工智能的持续改进开辟了一个新的维度,有望通过学习探索奖励来扩展大型语言模型 (LLM) 的训练数据。然而,先前发表的工作并未产生具有竞争力的结果。鉴于此,我们报告了 Kimi k1.5 的训练实践,Kimi k1.5 是我们最新的使用 RL 训练的多模态 LLM,包括其 RL 训练技术、多模态数据配方和基础设施优化。长上下文扩展和改进的策略优化方法是我们方法的关键要素,它建立了一个简单有效的 RL 框架,而无需依赖更复杂的技术,例如蒙特卡罗树搜索、价值函数和过程奖励模型。值得注意的是,我们的系统在多个基准测试和模态上实现了最先进的推理性能——例如,在 AIME 上达到 77.5,在 MATH 500 上达到 96.2,在 Codeforces 上达到第 94 个百分位数,在 MathVista 上达到 74.9,与 OpenAI 的 01 相匹配。此外,我们提出了有效的 long2short 方法,该方法使用长 CoT 技术来改进 short-CoT 模型,产生最先进的 short-CoT 推理结果——例如,在 AIME 上达到 60.8,在 MATH500 上达到 94.6,在 LiveCodeBench 上达到 47.3,其性能大幅优于现有的 short-CoT 模型,例如 GPT-40 和 Claude Sonnet 3.5(高达 +550%)。
1 简介
在扩展定律的背景下,语言模型使用下一个 Token 预测进行预训练已被研究,其中按比例缩放模型参数和数据大小会导致智能的持续改进。(Kaplan et al. 2020; Hoffmann et al. 2022) 然而,这种方法受限于可用的高质量训练数据量。(Villalobos et al. 2024; Muennighoff et al. 2023) 在本报告中,我们介绍了 Kimi k1.5 的训练配方,Kimi k1.5 是我们最新的使用强化学习 (RL) 训练的多模态 LLM。目标是探索一种可能的新的持续扩展维度。使用带有 LLM 的 RL,模型学习探索奖励,因此不限于预先存在的静态数据集。
关于 k1.5 的设计和训练,有几个关键要素。
- 长上下文扩展: 我们将 RL 的上下文窗口扩展到 128k,并观察到随着上下文长度的增加,性能持续提高。我们方法背后的一个关键思想是使用部分 rollout 来提高训练效率——即,通过重用先前轨迹的大部分来采样新轨迹,从而避免从头开始重新生成新轨迹的成本。我们的观察结果表明,上下文长度是 LLM 中 RL 持续扩展的关键维度。
- 改进的策略优化: 我们推导了具有 long-CoT 的 RL 公式,并采用在线镜像下降的变体来实现稳健的策略优化。通过我们有效的采样策略、长度惩罚和数据配方的优化,进一步改进了该算法。
- 简化的框架: 长上下文扩展与改进的策略优化方法相结合,建立了一个使用 LLM 进行学习的简化 RL 框架。由于我们能够扩展上下文长度,因此学习到的 CoT 表现出规划、反思和纠正的特性。增加的上下文长度会增加搜索步骤的数量。因此,我们表明,无需依赖更复杂的技术(例如蒙特卡罗树搜索、价值函数和过程奖励模型)即可实现强大的性能。
- 多模态: 我们的模型是使用文本和视觉数据联合训练的,这使其具有对这两种模态进行联合推理的能力。
此外,我们提出了有效的 long2short 方法,该方法使用 long-CoT 技术来改进 short-CoT 模型。具体来说,我们的方法包括应用带有 long-CoT 激活的长度惩罚和模型合并。
我们的 long-CoT 版本在多个基准测试和模态上实现了最先进的推理性能——例如,在 AIME 上达到 77.5,在 MATH 500 上达到 96.2,在 Codeforces 上达到第 94 个百分位数,在 MathVista 上达到 74.9,与 OpenAI 的 01 相匹配。我们的模型还实现了最先进的 short-CoT 推理结果——例如,在 AIME 上达到 60.8,在 MATH500 上达到 94.6,在 LiveCodeBench 上达到 47.3,其性能大幅优于现有的 short-CoT 模型,例如 GPT-40 和 Claude Sonnet 3.5(高达 +550%)。结果如图 1 和图 2 所示。
2 方法:使用 LLM 进行强化学习
Kimi k1.5 的开发包括几个阶段:预训练、vanilla 监督微调 (SFT)、long-CoT 监督微调和强化学习 (RL)。本报告侧重于 RL,首先概述 RL 提示集管理(第 2.1 节)和 long-CoT 监督微调(第 2.2 节),然后深入讨论第 2.3 节中的 RL 训练策略。有关预训练和 vanilla 监督微调的更多详细信息可以在第 2.5 节中找到。
2.1 RL 提示集管理
通过我们的初步实验,我们发现 RL 提示集的质量和多样性在确保强化学习的有效性方面起着至关重要的作用。精心构建的提示集不仅引导模型进行稳健的推理,而且还减轻了奖励攻击和过度拟合表面模式的风险。具体而言,三个关键属性定义了高质量的 RL 提示集:
- 多样化的覆盖范围: 提示应涵盖各种学科,例如 STEM、编码和一般推理,以增强模型的适应性并确保其在不同领域中的广泛适用性。
- 均衡的难度: 提示集应包括分布均匀的简单、中等和困难问题,以促进渐进式学习并防止过度拟合特定的复杂性级别。
- 准确的可评估性: 提示应允许验证者进行客观和可靠的评估,以确保模型性能是基于正确的推理而不是表面模式或随机猜测来衡量的。
为了在提示集中实现多样化的覆盖范围,我们采用自动过滤器来选择需要丰富推理且易于评估的问题。我们的数据集包括来自各个领域的问题,例如 STEM 领域、竞赛和一般推理任务,包含纯文本和图像文本问答数据。此外,我们开发了一个标记系统来按领域和学科对提示进行分类,以确保不同主题领域的均衡表示 (M. Li et al. 2023; W. Liu et al. 2023)。
我们采用了一种基于模型的方法,该方法利用模型自身的能力来自适应地评估每个提示的难度。具体来说,对于每个提示,SFT 模型使用相对较高的采样温度生成十次答案。然后计算通过率,并将其用作提示难度的代理——通过率越低,难度越高。这种方法允许难度评估与模型的内在能力保持一致,使其在 RL 训练中非常有效。通过利用这种方法,我们可以预先筛选掉大多数琐碎的案例,并在 RL 训练期间轻松探索不同的采样策略。
为了避免潜在的奖励攻击 (Everitt et al. 2021; Pan et al. 2022),我们需要确保可以准确验证每个提示的推理过程和最终答案。经验观察表明,一些复杂的推理问题可能具有相对简单且易于猜测的答案,从而导致虚假的正向验证——模型通过不正确的推理过程达到正确的答案。为了解决这个问题,我们排除了容易出现此类错误的问题,例如多项选择题、真/假题和基于证明的问题。此外,对于一般的问答任务,我们提出了一种简单而有效的方法来识别和删除易于攻击的提示。具体来说,我们提示模型在没有任何 CoT 推理步骤的情况下猜测可能的答案。如果模型在 N 次尝试中预测出正确的答案,则该提示被认为太容易攻击而被删除。我们发现,设置 N = 8 可以删除大多数容易攻击的提示。开发更高级的验证模型仍然是未来研究的一个开放方向。
2.2 Long-CoT 监督微调
通过改进的 RL 提示集,我们采用提示工程来构建一个小而高质量的 long-CoT 热身数据集,其中包含针对文本和图像输入的准确验证的推理路径。这种方法类似于拒绝采样 (RS),但侧重于通过提示工程生成 long-CoT 推理路径。由此产生的热身数据集旨在封装对类人推理至关重要的关键认知过程,例如规划(模型系统地概述执行前的步骤);评估(包括对中间步骤的批判性评估);反思(使模型能够重新考虑和改进其方法);以及探索(鼓励考虑替代解决方案)。通过在这个热身数据集上执行轻量级的 SFT,我们可以有效地引导模型内化这些推理策略。因此,经过微调的 long-CoT 模型表现出生成更详细且逻辑上连贯的响应的改进能力,从而增强了其在各种推理任务中的性能。
2.3 强化学习
2.3.1 问题设定
给定一个训练数据集 D = {(xi, yi)}^n_i=1,其中包含问题 xi 和对应的真实答案 yi,我们的目标是训练策略模型 πθ 以准确地解决测试问题。在复杂的推理背景下,问题 x 到解决方案 y 的映射并非易事。为了应对这一挑战,思维链 (CoT) 方法建议使用一系列中间步骤 z = (z1, z2, ..., zm) 来桥接 x 和 y,其中每个 zi 都是一个连贯的 Token 序列,可作为解决问题的重要的中间步骤 (J. Wei et al. 2022)。在解决问题 x 时,以自回归方式采样思维 zt ~ πθ(·|x, z1, ..., zt−1),然后采样最终答案 y ~ πθ(·|x, z1,..., zm)。我们使用 y, z ~ πθ 来表示此采样过程。请注意,思维和最终答案都是作为语言序列采样的。
为了进一步增强模型的推理能力,我们采用了规划算法来探索各种思维过程,从而在推理时生成改进的 CoT (Yao et al. 2024; Y. Wu et al. 2024; Snell et al. 2024)。这些方法的核心见解是显式构建由值估计引导的思维搜索树。这允许模型探索思维过程的各种延续,或者在遇到死胡同时回溯以调查新的方向。更详细地说,令 T 为搜索树,其中每个节点表示部分解决方案 s = (x, z1:|s|)。这里 s 由问题 x 和一系列思维 z1:|s| = (z1, ..., z|s|) 组成,这些思维通向该节点,其中 |s| 表示序列中思维的数量。规划算法使用评论模型 v 来提供反馈 v(x, z1:|s|),这有助于评估当前解决问题的进展并识别现有部分解决方案中的任何错误。我们注意到,反馈可以由判别分数或语言序列提供 (L. Zhang et al. 2024)。在所有 s ∈ T 的反馈的指导下,规划算法选择最有可能扩展的节点,从而扩大搜索树。上述过程会重复迭代,直到得出完整的解决方案。
我们还可以从算法的角度来处理规划算法。给定第 t 次迭代时可用的过去搜索历史 (s1, v(s1), ..., st−1, v(st−1)),规划算法 A 迭代地确定下一个搜索方向 A(st|s1, v(s1),..., st−1, v(st−1)) 并为当前搜索进度 A(v(st)|s1, v(s1),..., st) 提供反馈。由于思维和反馈都可以被视为中间推理步骤,并且这些组件都可以表示为语言 Token 序列,因此我们使用 z 来替换 s 和 v 以简化表示法。因此,我们将规划算法视为直接作用于一系列推理步骤 A(·|z1, z2, ... ) 的映射。在这个框架中,规划算法使用的搜索树中存储的所有信息都被展平到提供给算法的完整上下文中。这为生成高质量 CoT 提供了一个有趣的视角:我们可以训练一个模型来近似此过程,而不是显式构建搜索树并实现规划算法。在这里,思维的数量(即,语言 Token)类似于传统上分配给规划算法的计算预算。最近在长上下文窗口方面的进步促进了训练和测试阶段的无缝可扩展性。如果可行,此方法使模型能够通过自回归预测直接在推理空间上运行隐式搜索。因此,模型不仅学习解决一组训练问题,而且还开发了有效解决单个问题的能力,从而提高了对未见测试问题的泛化能力。
因此,我们考虑训练模型以使用强化学习 (RL) 生成 CoT (OpenAI 2024)。令 r 为奖励模型,该模型根据真实答案 y* 为给定问题 x 提出的答案 y 的正确性提供依据,方法是分配一个值 r(x, y, y*) ∈ {0, 1}。对于可验证的问题,奖励直接由预定义的标准或规则确定。例如,在编码问题中,我们会评估答案是否通过了测试用例。对于具有自由形式真实答案的问题,我们训练奖励模型 r(x, y, y*) 来预测答案是否与真实答案匹配。给定问题 x,模型 πθ 通过采样过程 z ~ πθ(·|x),y ~ πθ(·|x, z) 生成 CoT 和最终答案。生成的 CoT 的质量取决于它是否可以产生正确的最终答案。总之,我们考虑以下目标来优化策略
通过扩展 RL 训练,我们旨在训练一个模型,该模型可以利用简单基于提示的 CoT 和规划增强的 CoT 的优势。该模型仍然在推理期间自回归地采样语言序列,从而避免了部署期间高级规划算法所需的复杂并行化。然而,与简单的基于提示的方法的一个关键区别在于,模型不应仅仅遵循一系列推理步骤。相反,它还应该学习关键的规划技能,包括错误识别、回溯和解决方案改进,方法是利用整个探索思维集作为上下文信息。
2.3.2 策略优化
我们应用在线策略镜像下降的变体作为我们的训练算法 (Abbasi-Yadkori et al. 2019; Mei et al. 2019; Tomar et al. 2020)。该算法迭代执行。在第 i 次迭代中,我们使用当前模型 πθi 作为参考模型,并优化以下相对熵正则化策略优化问题,
其中 τ > 0 是控制正则化程度的参数。此目标具有封闭形式的解
其中 Z = Σy',z' πθi (y', z'|x) exp(r(x, y', y*)/τ) 是归一化因子。取两边的对数,我们得到对于任何 (y, z),以下约束都满足,这允许我们在优化期间利用离线策略数据
这激发了以下替代损失
为了近似 τ log Z,我们使用样本 (y1, z1), ..., (yk, zk) ~ πθ:τ log Z ≈ log 1/k Σ^k_j=1 exp(r(x, yj, y*)/τ)。我们还发现使用采样奖励的经验平均值 τ = mean(r(x, y1, y*), ..., r(x, yk, y*)) 会产生有效的实际结果。这是合理的,因为当 τ → ∞ 时,τ log Z 接近 πθi 下的预期奖励。最后,我们通过取替代损失的梯度来得出我们的学习算法。对于每个问题 x,使用参考策略 πθi 采样 k 个响应,并且梯度由下式给出
对于那些熟悉策略梯度方法的人来说,此梯度类似于使用采样奖励的平均值作为基线的 (2) 的策略梯度 (Kool et al. 2019; Ahmadian et al. 2024)。主要区别在于响应是从 πθi 而不是在线策略中采样的,并且应用了 l2 正则化。因此,我们可以将其视为通常的在线正则化策略梯度算法到离线策略情况的自然扩展 (Nachum et al. 2017)。我们从 D 中采样一批问题,并将参数更新为 θi+1,随后该参数将作为下一次迭代的参考策略。由于每次迭代都会由于参考策略的变化而考虑不同的优化问题,因此我们还在每次迭代开始时重置优化器。
我们排除了训练系统中的价值网络,该网络也在之前的研究中被利用 (Ahmadian et al. 2024)。虽然这种设计选择显着提高了训练效率,但我们也假设经典 RL 中用于信用分配的价值函数的传统使用可能不适合我们的背景。考虑这样一种情况:模型生成了部分 CoT (z1, z2, ..., zt) 并且有两个潜在的下一个推理步骤:zt+1 和 z't+1。假设 zt+1 直接导致了正确的答案,而 z't+1 包含一些错误。如果可以访问 oracle 价值函数,它将表明 zt+1 比 z't+1 保留更高的价值。根据标准信用分配原则,选择 z't+1 将受到惩罚,因为它相对于当前策略具有负优势。但是,探索 z't+1 对于训练模型生成 long CoT 非常有价值。通过使用从 long CoT 得出的最终答案的理由作为奖励信号,模型可以从采用 z't+1 的试错模式中学习,只要它成功恢复并得出正确的答案。从这个例子中得出的关键结论是,我们应该鼓励模型探索各种推理路径,以增强其解决复杂问题的能力。这种探索性方法会产生大量的经验,这些经验支持关键规划技能的开发。我们的主要目标不是将训练问题的准确性提高到最高水平,而是侧重于使模型掌握有效的问题解决策略,从而最终提高其在测试问题上的性能。
2.3.3 长度惩罚
我们观察到一种过度思考现象,即模型在 RL 训练期间的响应长度显着增加。虽然这会导致更好的性能,但过度冗长的推理过程在训练和推理期间成本高昂,并且人类通常不喜欢过度思考。为了解决这个问题,我们引入了长度奖励来限制 Token 长度的快速增长,从而提高模型的 Token 效率。给定问题 x 的 k 个采样响应 (y1, z1), ..., (yk, zk) 和真实答案 y*,令 len(i) 为 (yi, zi) 的长度,min_len = min_i len(i) 且 max_len = max_i len(i)。如果 max_len = min_len,我们为所有响应设置长度奖励为零,因为它们具有相同的长度。否则,长度奖励由下式给出
本质上,我们鼓励较短的响应,并惩罚正确答案中较长的响应,同时明确惩罚错误答案中较长的响应。然后将此基于长度的奖励与权重参数添加到原始奖励中。
在我们的初步实验中,长度惩罚可能会减慢初始阶段的训练速度。为了缓解这个问题,我们建议在训练的初始阶段逐步热身长度惩罚。具体来说,我们在没有长度惩罚的情况下采用标准策略优化,然后在训练的其余部分采用恒定的长度惩罚。
2.3.4 采样策略
尽管 RL 算法本身具有相对较好的采样属性(难度更大的问题提供更大的梯度),但其训练效率受到限制。因此,一些定义明确的先验采样方法可以产生更大的性能提升。我们利用多种信号来进一步改进采样策略。首先,我们收集的 RL 训练数据自然带有不同的难度标签。例如,数学竞赛问题比小学数学题更难。其次,由于 RL 训练过程多次采样相同的问题,我们还可以跟踪每个问题的成功率作为难度的指标。我们提出了两种采样方法来利用这些先验知识来提高训练效率。
课程采样 我们首先在较简单的任务上进行训练,然后逐步过渡到更具挑战性的任务。由于初始 RL 模型的性能有限,因此将有限的计算预算花费在非常困难的问题上通常会产生很少的正确样本,从而导致较低的训练效率。同时,我们收集的数据自然包含等级和难度标签,这使得基于难度的采样成为提高训练效率的一种直观有效的方法。
优先采样 除了课程采样外,我们还使用优先采样策略来关注模型性能不佳的问题。我们跟踪每个问题 i 的成功率 si,并以与 1/si 成比例的方式采样问题,以便成功率较低的问题获得更高的采样概率。这会将模型的努力导向其最薄弱的领域,从而加快学习速度并提高整体性能。
2.3.5 有关训练配方的更多详细信息
用于编码的测试用例生成 由于网上许多编码问题没有可用的测试用例,我们设计了一种自动生成测试用例的方法,这些测试用例可以用作使用 RL 训练模型的奖励。我们的重点主要放在不需要特殊裁判的问题上。我们还假设这些问题的真实答案是可用的,以便我们可以利用这些答案来生成更高质量的测试用例。
我们利用广泛认可的测试用例生成库 CYaRon¹ 来增强我们的方法。我们使用我们的基础 Kimi k1.5 基于问题陈述生成测试用例。CYaRon 的使用说明和问题描述作为生成器的输入。对于每个问题,我们首先使用生成器生成 50 个测试用例,并随机采样每个测试用例的 10 个真实答案提交。我们针对提交运行测试用例。如果至少 7 个提交中有 10 个产生匹配的结果,则认为测试用例有效。经过此轮筛选后,我们获得了一组选定的测试用例。如果至少 9 个提交中有 10 个通过了整套选定的测试用例,则将问题及其关联的选定测试用例添加到我们的训练集中。
在统计方面,从 1,000 个在线竞赛问题的样本中,大约有 614 个问题不需要特殊裁判。我们开发了 463 个测试用例生成器,这些生成器产生了至少 40 个有效的测试用例,从而导致我们的训练集中包含 323 个问题。
数学的奖励建模 评估数学解决方案的一个挑战是,不同的书写形式可以表示相同的基本答案。例如,a² - 4 和 (a + 2)(a - 2) 可能都是有效解决方案。我们采用了两种方法来提高奖励模型的评分准确性:
- 经典 RM:从 InstructGPT (Ouyang et al. 2022) 方法中汲取灵感,我们实现了一个基于值头的奖励模型,并收集了大约 80 万个用于微调的数据点。该模型最终将“问题”、“参考答案”和“响应”作为输入,并输出一个表示响应是否正确的标量。
- 思维链 RM:最近的研究 (Ankner et al. 2024; McAleese et al. 2024) 表明,使用思维链 (CoT) 推理增强的奖励模型可以显着优于经典方法,尤其是在需要细微的正确性标准(例如数学)的任务中。因此,我们收集了大约 80 万个 CoT 标记的示例的同等大小的数据集来微调 Kimi 模型。在与经典 RM 相同的输入基础上,思维链方法在 JSON 格式中明确生成逐步推理过程,然后提供最终的正确性判断,从而实现更强大且可解释的奖励信号。
在我们的手动抽查中,经典 RM 的准确率约为 84.4,而思维链 RM 的准确率达到 98.5。在 RL 训练过程中,我们采用了思维链 RM 来确保更正确的反馈。
视觉数据 为了提高模型的真实世界图像推理能力并实现视觉输入与大型语言模型 (LLM) 之间更有效的对齐,我们的视觉强化学习 (Vision RL) 数据主要来自三个不同的类别:真实世界数据、合成视觉推理数据和文本渲染数据。
- 真实世界数据涵盖了各个年级的各种科学问题,这些问题需要图形理解和推理;需要视觉感知和推理的定位猜测任务;以及涉及理解复杂图表的分析任务等。这些数据集提高了模型在真实世界场景中执行视觉推理的能力。
- 合成视觉推理数据是人为生成的,包括旨在提高特定视觉推理技能的程序创建的图像和场景,例如理解空间关系、几何图案和对象交互。这些合成数据集为测试模型的视觉推理能力提供了受控环境,并提供了无穷无尽的训练示例。
- 文本渲染数据是通过将文本内容转换为视觉格式创建的,使模型在处理跨不同模态的基于文本的查询时保持一致性。通过将文本文档、代码片段和结构化数据转换为图像,我们确保模型无论输入是纯文本还是以图像形式渲染的文本(如屏幕截图或照片),都能提供一致的响应。这也有助于增强模型在处理文本繁重图像时的能力。
每种类型的数据对于构建一个全面的视觉语言模型至关重要,该模型可以有效地管理各种真实世界的应用,同时确保在各种输入模式下保持一致的性能。
2.4 Long2short:用于 Short-CoT 模型的上下文压缩
尽管 long-CoT 模型实现了强大的性能,但与标准的 short-CoT LLM 相比,它消耗了更多的测试时 Token。但是,可以将在 long-CoT 模型中的思维先验知识转移到 short-CoT 模型中,以便即使在有限的测试时 Token 预算下也可以提高性能。我们针对此 long2short 问题提出了几种方法,包括模型合并 (Yang et al. 2024)、最短拒绝采样、DPO (Rafailov et al. 2024) 和 long2short RL。以下提供了对这些方法的详细描述:
模型合并 模型合并已被发现有助于保持泛化能力。我们还发现它在合并 long-CoT 模型和 short-CoT 模型时提高了 Token 效率。这种方法将 long-CoT 模型与较短的模型相结合,无需训练即可获得新模型。具体来说,我们通过简单地平均它们的权重来合并这两个模型。
最短拒绝采样 我们观察到我们的模型为同一个问题生成长度差异很大的响应。基于此,我们设计了最短拒绝采样方法。这种方法对同一个问题采样 n 次(在我们的实验中,n = 8),并选择最短的正确响应进行监督微调。
DPO 与最短拒绝采样类似,我们利用 Long CoT 模型来生成多个响应样本。选择最短的正确解决方案作为正样本,而较长的响应则被视为负样本,包括错误较长的响应和正确较长的响应(比选择的正样本长 1.5 倍)。这些正负对构成了用于 DPO 训练的成对偏好数据。
Long2short RL 在标准的 RL 训练阶段之后,我们选择一个在性能和 Token 效率之间提供最佳平衡的模型作为基础模型,并进行单独的 long2short RL 训练阶段。在第二阶段,我们应用第 2.3.3 节中介绍的长度惩罚,并显着减少最大 rollout 长度,以进一步惩罚超出所需长度但可能正确的响应。
2.5 其他训练详情
2.5.1 预训练
Kimi k1.5 基础模型是在一个多样化、高质量的多模态语料库上训练的。语言数据涵盖五个领域:英语、中文、代码、数学推理和知识。多模态数据,包括字幕、图像文本交错、OCR、知识和 QA 数据集,使我们的模型能够获得视觉语言能力。严格的质量控制确保了整体预训练数据集的相关性、多样性和平衡性。我们的预训练分三个阶段进行:(1) 视觉语言预训练,其中建立了强大的语言基础,然后逐步进行多模态集成;(2) 冷静期,利用精选和合成数据巩固能力,特别是对于推理和基于知识的任务;以及 (3) 长上下文激活,将序列处理扩展到 131,072 个 Token。有关我们预训练工作的更多详细信息,请参见附录 B。
2.5.2 Vanilla 监督微调
我们创建了涵盖多个领域的 vanilla SFT 语料库。对于非推理任务,包括问答、写作和文本处理,我们最初通过人工注释构建种子数据集。此种子数据集用于训练种子模型。随后,我们收集各种提示,并使用种子模型为每个提示生成多个响应。然后,注释者对这些响应进行排名,并改进排名最高的响应以生成最终版本。对于诸如数学和编码问题之类的推理任务,其中基于规则和奖励建模的验证比人工判断更准确、更高效,我们利用拒绝采样来扩展 SFT 数据集。
我们的 vanilla SFT 数据集包含大约 100 万个文本示例。具体来说,50 万个示例用于一般问答,20 万个用于编码,20 万个用于数学和科学,5 千个用于创意写作,以及 2 万个用于长上下文任务,例如摘要、文档问答、翻译和写作。此外,我们构建了 100 万个文本视觉示例,涵盖了各种类别,包括图表解释、OCR、基于图像的对话、视觉编码、视觉推理以及具有视觉辅助的数学/科学问题。
我们首先在 32k 个 Token 的序列长度下训练模型 1 个 epoch,然后在 128k 个 Token 的序列长度下再训练 1 个 epoch。在第一阶段 (32k),学习率从 2 × 10-5 衰减到 2 × 10-6,然后在第二阶段 (128k) 重新升温到 1 × 10-5,最后衰减到 1 × 10-6。为了提高训练效率,我们将多个训练示例打包到每个训练序列中。
2.6 RL 基础设施
(a) 系统概述
(b) 部分 rollout
2.6.1 用于 LLM 的大规模强化学习训练系统
在人工智能领域,强化学习 (RL) 已经成为大型语言模型 (LLM) 的一种关键训练方法 (Ouyang et al. 2022) (Jaech et al. 2024),它从其在掌握复杂游戏(如围棋、星际争霸 II 和 Dota 2)方面的成功中汲取灵感,例如 AlphaGo (Silver et al. 2017)、AlphaStar (Vinyals et al. 2019) 和 OpenAI Dota Five (Berner et al. 2019) 等系统。遵循这一传统,Kimi k1.5 系统采用了迭代同步 RL 框架,该框架经过精心设计,可通过持久学习和适应来增强模型的推理能力。此系统中的一项关键创新是引入了部分 rollout 技术,该技术旨在优化复杂推理轨迹的处理。
如图 3a 所示,RL 训练系统通过迭代同步方法运行,每次迭代都包含一个 rollout 阶段和一个训练阶段。在 rollout 阶段,由中央主服务器协调的 rollout 工作程序通过与模型交互来生成 rollout 轨迹,从而产生针对各种输入的响应序列。然后,这些轨迹存储在重放缓冲区中,这确保了一个多样化且无偏差的训练数据集,方法是中断时间相关性。在随后的训练阶段,训练工作程序会访问这些经验以更新模型的权重。此循环过程允许模型从其操作中不断学习,并随着时间的推移调整其策略以提高性能。
中央主服务器充当中央指挥,管理 rollout 工作程序、训练工作程序、具有奖励模型的评估和重放缓冲区之间的数据流和通信。它确保系统协调运行,平衡负载并促进高效的数据处理。
训练工作程序访问这些 rollout 轨迹,无论是在单个迭代中完成还是在多个迭代中划分,以计算梯度更新,从而改进模型的参数并增强其性能。此过程由奖励模型监督,该模型评估模型输出的质量并提供必要的反馈来指导训练过程。奖励模型的评估对于确定模型策略的有效性和引导模型走向最佳性能至关重要。
此外,该系统还包含代码执行服务,该服务专门设计用于处理代码相关的问题,并且是奖励模型的组成部分。此服务在实际编码场景中评估模型的输出,确保模型的学习与真实世界的编程挑战紧密结合。通过针对实际代码执行验证模型的解决方案,此反馈循环对于改进模型的策略和增强其在代码相关任务中的性能至关重要。
2.6.2 用于 Long CoT RL 的部分 Rollout
我们工作的主要思想之一是扩展长上下文 RL 训练。部分 rollout 是一种关键技术,它通过管理长轨迹和短轨迹的 rollout,有效地解决了处理 long-CoT 特性的挑战。此技术建立了固定的输出 Token 预算,限制每个 rollout 轨迹的长度。如果轨迹在 rollout 阶段超出 Token 限制,则将未完成的部分保存到重放缓冲区中,并在下一次迭代中继续。它确保没有单个冗长的轨迹垄断系统的资源。此外,由于 rollout 工作程序是异步运行的,当某些工作程序正在处理长轨迹时,其他工作程序可以独立处理新的、较短的 rollout 任务。异步操作通过确保所有 rollout 工作程序都积极为训练过程做出贡献来最大化计算效率,从而优化系统的整体性能。
如图 3b 所示,部分 rollout 系统通过将长响应分解为跨迭代的片段(从迭代 n-m 到迭代 n)来工作。重放缓冲区充当中央存储机制,用于维护这些响应片段,其中只有当前迭代(迭代 n)需要在线计算。以前的片段(迭代 n-m 到 n-1)可以从缓冲区中有效地重用,从而无需重复 rollout。这种分段方法显着减少了计算开销:系统不是一次性 rollout 整个响应,而是以增量方式处理和存储片段,从而可以在保持快速迭代时间的同时生成更长的响应。在训练期间,某些片段可以从损失计算中排除,以进一步优化学习过程,从而使整个系统既高效又可扩展。
部分 rollout 的实现还提供了重复检测。该系统会识别生成内容中重复的序列,并提前终止它们,从而在保持输出质量的同时减少不必要的计算。检测到的重复项可能会被分配额外的惩罚,从而有效地阻止提示集中出现冗余的内容生成。
2.6.3 训练和推理的混合部署
RL 训练过程包括以下阶段:
- 训练阶段:首先,Megatron (Shoeybi et al. 2020) 和 vLLM (Kwon et al. 2023) 在单独的容器中执行,这些容器由称为检查点引擎(第 2.6.3 节)的 shim 进程封装。Megatron 开始训练过程。训练完成后,Megatron 会卸载 GPU 内存,并准备将当前权重传输到 vLLM。
- 推理阶段:在 Megatron 卸载后,vLLM 从虚拟模型权重开始,并使用通过 Mooncake (Qin et al. 2024) 从 Megatron 传输的最新权重更新它们。在 rollout 完成后,检查点引擎将停止所有 vLLM 进程。
- 后续训练阶段:一旦分配给 vLLM 的内存被释放,Megatron 就会卸载内存并启动另一轮训练。
我们发现现有工作难以同时支持以下所有特性。
- 复杂的并行策略:Megatron 可能具有与 vLLM 不同的并行策略。在 Megatron 的多个节点中分配的训练权重可能难以与 vLLM 共享。
- 最大限度地减少空闲 GPU 资源:对于在线策略 RL,最近的工作(如 SGLang (L. Zheng et al. 2024) 和 VLLM)可能会在训练过程中保留一些 GPU,这反过来可能会导致空闲训练 GPU。在训练和推理之间共享相同的设备会更有效率。
- 动态扩展能力:在某些情况下,可以通过增加推理节点的数量,同时保持训练过程不变来实现显着加速。我们的系统可以有效地利用所需的空闲 GPU 节点。
如图 4 所示,我们在 Megatron 和 vLLM 之上实现了此混合部署框架(第 2.6.3 节),从而实现了从训练到推理阶段不到一分钟的时间,反之则约为 10 秒。
混合部署策略 我们提出了一种用于训练和推理任务的混合部署策略,该策略利用 Kubernetes Sidecar 容器共享所有可用的 GPU,以将两种工作负载并置在一个 pod 中。此策略的主要优点是:
- 它有助于高效的资源共享和管理,防止训练节点在等待推理节点时闲置(当两者部署在单独的节点上时)。
- 利用不同的部署映像,训练和推理可以独立迭代以获得更好的性能。
- 该架构不限于 vLLM,其他框架也可以方便地集成。
检查点引擎 检查点引擎负责管理 vLLM 进程的生命周期,公开 HTTP API,这些 API 能够触发 vLLM 的各种操作。为了保证整体一致性和可靠性,我们利用由 etcd 服务管理的全局元数据系统来广播操作和状态。
当 vLLM 卸载时,它可能难以完全释放 GPU 内存,这主要是由于 CUDA 图、NCCL 缓冲区和 NVIDIA 驱动程序。为了最大程度地减少对 vLLM 的修改,我们会在需要时终止并重新启动它,以获得更好的 GPU 利用率和容错能力。
Megatron 中的工作程序将拥有的检查点转换为共享内存中的 Hugging Face 格式。此转换还考虑了管道并行性和专家并行性,以便在这些检查点中仅保留张量并行性。共享内存中的检查点随后被分成碎片并注册到全局元数据系统中。我们使用 Mooncake 通过 RDMA 在对等节点之间传输检查点。需要对 vLLM 进行一些修改才能加载权重文件并执行张量并行转换。
我们开发了沙箱作为执行用户提交代码的安全环境,该环境针对代码执行和代码基准评估进行了优化。通过动态切换容器映像,沙箱通过 MultiPL-E (Cassano, Gouwar, D. Nguyen, S. Nguyen, et al. 2023)、DMOJ Judge Server²、Lean、Jupyter Notebook 和其他映像支持不同的用例。
对于编码任务中的 RL,沙箱通过提供一致且可重复的评估机制来确保训练数据判断的可靠性。它的反馈系统支持多阶段评估,例如代码执行反馈和 repo 级编辑,同时保持统一的上下文,以确保跨编程语言进行公平和公正的基准比较。
我们将该服务部署在 Kubernetes 上以实现可扩展性和弹性,通过 HTTP 端点将其公开以进行外部集成。Kubernetes 的功能(如自动重启和滚动更新)确保了可用性和容错能力。
为了优化性能并支持 RL 环境,我们将几种技术集成到代码执行服务中,以提高效率、速度和可靠性。其中包括:
- 使用 Crun:我们使用 crun 作为容器运行时而不是 Docker,从而显着减少了容器启动时间。
- Cgroup 重用:我们预先创建用于容器使用的 cgroup,这在高并发场景中至关重要,在这些场景中,为每个容器创建和销毁 cgroup 可能会成为瓶颈。
- 磁盘使用优化:使用安装为 tmpfs 的上层的覆盖文件系统来控制磁盘写入,从而提供固定大小的高速存储空间。这种方法有利于临时工作负载。
这些优化提高了代码执行中的 RL 效率,从而为评估 RL 生成的代码提供了持续可靠的环境,这对于迭代训练和模型改进至关重要。
3 实验
3.1 评估
由于 k1.5 是多模态模型,我们针对不同模态在各种基准测试上进行了全面评估。详细的评估设置可在附录 C 中找到。我们的基准测试主要包括以下三个类别:
- 文本基准测试:MMLU (Hendrycks et al. 2020)、IF-Eval (J. Zhou et al. 2023)、CLUEWSC (L. Xu et al. 2020)、C-EVAL (Y. Huang et al. 2023)
- 推理基准测试:HumanEval-Mul、LiveCodeBench (Jain et al. 2024)、Codeforces、AIME 2024、MATH-500 (Lightman et al. 2023)
- 视觉基准测试:MMMU (Yue, Ni, et al. 2024)、MATH-Vision (K. Wang et al. 2024)、MathVista (Lu et al. 2023)
3.2 主要结果
K1.5 long-CoT 模型 表 2 中显示了 Kimi k1.5 long-CoT 模型的性能。通过 long-CoT 监督微调(在第 2.2 节中描述)和视觉文本联合强化学习(在第 2.3 节中讨论),模型的长期推理能力得到显着增强。测试时计算扩展进一步增强了其性能,使模型能够在各种模式中实现最先进的结果。我们的评估揭示了模型在扩展上下文中推理、理解和综合信息的能力方面的显着改进,代表了多模态 AI 能力的进步。
K1.5 short-CoT 模型 表 3 中显示了 Kimi k1.5 short-CoT 模型的性能。此模型集成了多种技术,包括传统的监督微调(在第 2.5.2 节中讨论)、强化学习(在第 2.3 节中探讨)和长到短的提炼(在第 2.4 节中概述)。结果表明,与领先的开源和专有模型相比,k1.5 short-CoT 模型在多项任务中提供了有竞争力的或卓越的性能。其中包括文本、视觉和推理挑战,在自然语言理解、数学、编码和逻辑推理方面具有显著优势。
3.3 长上下文扩展
我们采用中型模型来研究具有 LLM 的 RL 的扩展属性。图 5 说明了在数学提示集上训练的小模型变体的训练迭代中训练准确率和响应长度的演变。随着训练的进行,我们观察到响应长度和性能准确率同步增加。值得注意的是,更具挑战性的基准测试的响应长度增加得更快,这表明该模型学习为复杂问题生成更详细的解决方案。图 6 表明模型的输出上下文长度与其问题解决能力之间存在很强的相关性。我们最终运行的 k1.5 将上下文长度扩展到 128k,并观察到在硬推理基准测试上持续改进。
3.4 Long2short
我们将提议的 long2short RL 算法与第 2.4 节中介绍的 DPO、最短拒绝采样和模型合并方法进行了比较,重点关注 long2short 问题的 Token 效率 (X. Chen et al. 2024),特别是如何使获得的 long-cot 模型能够使短模型受益。在图 7 中,k1.5-long 表示我们选择用于 long2short 训练的 long-cot 模型。k1.5-short w/ rl 指的是使用 long2short RL 训练获得的短模型。k1.5-short w/ dpo 表示通过 DPO 训练提高 Token 效率的短模型。k1.5-short w/ merge 表示模型合并后的模型,而 k1.5-short w/ merge + rs 表示通过对合并的模型应用最短拒绝采样获得的短模型。k1.5-shortest 表示我们在 long2short 训练期间获得的最短模型。如图 7 所示,与 DPO 和模型合并等其他方法相比,提议的 long2short RL 算法表现出最高的 Token 效率。值得注意的是,k1.5 系列中的所有模型(以橙色标记)都表现出比其他模型(以蓝色标记)更好的 Token 效率。例如,k1.5-short w/ rl 在 AIME2024 上实现了 60.8 的 Pass@1 分数(平均 8 次运行),同时平均仅使用 3,272 个 Token。同样,k1.5-shortest 在 MATH500 上获得了 88.2 的 Pass@1 分数,同时消耗的 Token 数量与其他短模型大致相同。
3.5 消融研究
模型大小和上下文长度的扩展 我们的主要贡献是将 RL 应用于增强模型生成扩展 CoT 的能力,从而提高其推理能力。一个自然而然的问题是:这与简单地增加模型大小相比如何?为了证明我们方法的有效性,我们使用相同的数据集训练了两个不同大小的模型,并记录了在 RL 训练期间所有检查点的评估结果和平均推理长度。这些结果如图 8 所示。值得注意的是,尽管较大的模型最初优于较小的模型,但较小的模型可以通过利用通过 RL 优化的较长 CoT 来实现可比的性能。但是,较大的模型通常比较小的模型表现出更好的 Token 效率。这也表明,如果目标是获得尽可能好的性能,则扩展较大模型的上下文长度具有更高的上限,并且 Token 效率更高。但是,如果测试时计算具有预算,则训练具有较长上下文长度的较小模型可能是可行的解决方案。
使用负梯度的影响 我们研究了在我们的设置中使用 ReST (Gulcehre et al. 2023) 作为策略优化算法的有效性。ReST 与其他基于 RL 的方法(包括我们的方法)之间的主要区别在于,ReST 通过拟合从当前模型采样的最佳响应来迭代地改进模型,而不会应用负梯度来惩罚不正确的响应。如图 10 所示,与 ReST 相比,我们的方法表现出更高的样本复杂性,表明纳入负梯度显着提高了模型生成 long-CoT 的效率。我们的方法不仅提高了推理质量,而且还优化了训练过程,从而使用更少的训练样本实现了稳健的性能。这一发现表明,策略优化算法的选择在我们的设置中至关重要,因为 ReST 和其他基于 RL 的方法之间的性能差距在其他领域并不像 (Gulcehre et al. 2023) 中那么明显。因此,我们的结果突出了选择适当的优化策略以最大程度地提高生成 long-CoT 的效率的重要性。
采样策略 我们进一步证明了第 2.3.4 节中介绍的课程采样策略的有效性。我们的训练数据集 D 包含各种难度级别的问题。通过我们的课程采样方法,我们最初使用 D 进行热身阶段,然后仅关注硬问题来训练模型。将此方法与采用不进行任何课程调整的统一采样策略的基线方法进行比较。如图 9 所示,我们的结果清楚地表明,提议的课程采样方法显着提高了性能。这种改进可以归因于该方法逐步挑战模型的能力,使其能够在处理复杂问题时形成更强大的理解和能力。通过在初步的概括性介绍之后将训练精力集中在更困难的问题上,模型可以更好地加强其推理和问题解决能力。
4 结论
我们介绍了 k1.5(我们最新的使用 RL 训练的多模态 LLM)的训练配方和系统设计。我们从实践中得出的一个关键见解是,上下文长度的扩展对于 LLM 的持续改进至关重要。我们采用了优化的学习算法和基础设施优化,例如部分 rollout,以实现高效的长上下文 RL 训练。如何进一步提高长上下文 RL 训练的效率和可扩展性仍然是一个重要的前进方向。
我们做出的另一个贡献是结合了多种技术,从而可以改进策略优化。具体来说,我们使用 LLM 公式化了 long-CoT RL,并推导了用于鲁棒优化的在线镜像下降的变体。我们还尝试了采样策略、长度惩罚和优化数据配方,以实现强大的 RL 性能。
我们表明,即使不使用更复杂的技术(例如蒙特卡罗树搜索、价值函数和过程奖励模型),也可以通过长上下文扩展和改进的策略优化来实现强大的性能。在未来,研究如何在不损害模型探索能力的情况下改进信用分配和减少过度思考也将很有趣。
我们还观察到了 long2short 方法的潜力。这些方法在很大程度上提高了 short-CoT 模型的性能。此外,可以通过迭代方式将 long2short 方法与 long-CoT RL 相结合,以进一步提高 Token 效率,并在给定的上下文长度预算中提取最佳性能。