刚刚,Kimi开源新架构,开始押注线性注意力

Kimi 押注线性注意力,MiniMax 青睐全注意力,究竟哪条路能走得更远?

在智能体时代,推理的计算需求正成为一个核心瓶颈,尤其是在长时程和强化学习场景中。此时,标准注意力机制中存在的低效问题变得更加突出。

线性注意力为降低计算复杂度提供了一种有前景的方法,但由于表达能力有限,它在语言建模方面的表现历来不如 softmax 注意力,即使对于短序列也是如此。

最近的进展显著缩小了这一差距,主要得益于两项创新:门控或衰减机制以及 delta 规则。这些进展共同推动线性注意力在中等长度序列上的性能接近 softmax 水平。尽管如此,纯粹的线性结构从根本上仍受限于有限状态容量,这使得长序列建模和上下文内检索在理论上仍具有挑战性。

因此,结合 softmax 注意力和线性注意力的混合架构成为在质量和效率之间的一种折衷方案。但之前的混合模型往往规模较小,缺乏多样化基准评估。关键挑战是开发出一种新的注意力架构,能够在速度和内存上显著提高效率,同时保证或超过全注意力的质量,推动下一代解码密集型 LLM 的发展。

最近,月之暗面在这一方向有所突破。在一篇新的技术报告中,他们提出了一种新的混合线性注意力架构 ------Kimi Linear。该架构在各种场景中都优于传统的全注意力方法,包括短文本、长文本以及强化学习的 scaling 机制。

Kimi Linear 的核心是 Kimi Delta 注意力(KDA)------ 这是 Gated DeltaNet(GDN)的改进版本,引入了更高效的门控机制,以优化有限状态 RNN 内存的使用。作者表示,虽然 GDN 与 Mamba2 类似,采用了粗糙的 head-wise 遗忘门,但 KDA 引入了一种 channel-wise 的变体,其中每个特征维度都保持独立的遗忘率,类似于门控线性注意力(GLA)。

这种细粒度的设计能够更精确地调控有限状态 RNN 的记忆,从而在混合架构中释放 RNN 风格模型的潜力。关键在于,KDA 通过 Diagonal-Plus-Low-Rank(DPLR)矩阵的一种专门变体来参数化其转换动态,这使得一种定制的分块并行算法成为可能,该算法相较于通用的 DPLR 公式能显著减少计算量,同时仍与经典的 delta 规则保持一致。Kimi Linear 以 3:1 的固定比例将 KDA 与周期性的全注意力层交错排列。

作者基于 KDA 与多头潜在注意力(MLA)的分层混合架构,预训练了一个 Kimi Linear 模型。该模型激活参数为 3B,总参数达 48B。

Kimi Linear 实现了卓越的性能和硬件效率,尤其在长上下文任务中表现突出。它最多可将对大型 KV 缓存的需求减少 75%,并且在处理长达 100 万个 token 的上下文时,能将解码吞吐量提升到完整 MLA 模型的 6 倍。

月之暗面团队在 FLA 中开源了 KDA 内核,并发布了两个版本的模型检查点,这些检查点是用 5.7 万亿个 token 训练的。

整个项目的核心看点包括:

  • Kimi Delta 注意力(KDA):一种线性注意力机制,通过细粒度门控优化门控 Delta 规则。

  • 混合架构:3:1 的 KDA 与全局 MLA 比例在降低内存使用的同时,保持甚至超越全注意力的质量。

  • 卓越性能:在多种任务中表现优于全注意力,包括在 1.4 万亿 token 的训练运行中,通过公平对比,在长上下文和强化学习风格的基准测试中均有出色表现。

  • 高吞吐量:实现了高达 6 倍的更快解码速度,并显著减少了每个输出 token 的时间(TPOT)。

目前,vLLM 已经官宣支持 Kimi Linear。

在技术报告发布后,Kimi 的研究员「熊狸」在 X 上表示:「我很荣幸在过去一年中见证了这项伟大的工作。线性注意力在表达能力上具有巨大潜力,但在处理长上下文时存在较高的风险。KDA 不是一个玩具式的想法,它经受住了各种不可作弊的内部评估,这些评估的目的就是为了验证其有效性。」

该项目的重要贡献者 Zongyu Lin(目前在 UCLA)表示,「坦白说,这只是一个中间阶段,最终我们仍然在朝着实现无限上下文模型迈进。只要我们使用全局注意力,长时间解码依然受到其限制,而线性注意力背后仍然有一些基础设施挑战。但我相信这些问题都会被解决,而且来自不同实验室 / 公司的更多令人激动的成果即将到来。」

所以,K3 也会延续线性注意力架构吗?

而另外的大模型玩家中,Qwen 也曾表示要大胆押注线性注意力:

但 MiniMax 似乎更青睐全注意力。

到底哪个方向能走得更远?我们拭目以待。

以下是 Kimi Linear 的技术细节。

Kimi Delta Attention:通过细粒度门控改进 Delta 规则

Kimi Delta Attention (KDA) 是一种新型的门控线性注意力变体。它通过引入一个细粒度的对角化门控 来改进 GDN 的标量衰减,从而实现了对记忆衰减和位置感知的细粒度控制。

硬件高效的分块算法

通过将公式 1 的递归部分展开为分块公式,可得到:

  • WY Representation

通常用于将一系列秩 - 1 更新打包成单个紧凑的表示。该方法遵循了 Comba 中 P 的公式,以减少后续计算中额外矩阵求逆的需求。

其中,辅助向量 通过以下递归关系计算得出:

  • UT transform

该算法应用了 UT transform 来减少非矩阵乘法的 FLOPs,这对于在训练期间提升硬件利用率至关重要。

下三角矩阵的逆可以通过高斯消元法中的前向替换,采用逐行迭代的方法高效计算。等效地,以矩阵形式,可以按如下方式分块更新状态:

在输出阶段,该算法采用块间递归和块内并行的策略来最大化矩阵乘法吞吐量,从而充分利用 Tensor Cores 的计算潜力。

效率分析

在表示能力方面,KDA 与广义的 DPLR 公式一致,即

两者都表现出细粒度的衰减行为。然而,这种细粒度的衰减会在除法运算(例如,公式 9 中的块内计算)过程中引入数值精度问题。

为了解决这个问题,以往的工作(如 GLA)在对数域中执行计算,并在全精度下引入二级分块。然而,这种方法妨碍了半精度矩阵乘法的充分利用,并显著降低了算子速度。

通过将变量 a 和 b 同时绑定到 k,KDA 有效地缓解了这一瓶颈 ------ 将二级分块矩阵计算的数量从四次减少到两次,并进一步消除了三次额外的矩阵乘法。因此,与 DPLR 公式相比,KDA 的算子效率提升了大约 100%。

在批量大小统一为 1 且头数为 16 的条件下,算子随输入长度变化的执行时间。

Kimi 线性模型架构

该模型架构的主干遵循 Moonlight 的设计。除了细粒度门控之外,该模型还利用了几个组件来进一步提高 Kimi Linear 的表达能力。Kimi Linear 的整体架构如图 3 所示。

神经参数化

是第 t 个 token 的输入表示,KDA 每个头 h 的输入计算如下:

其中 dk, dv 代表 key (键) 和 value (值) 的头部维度,在所有实验中均设置为 128。对于 q, k, v,该模型应用了一个短卷积,随后是一个 Swish 激活函数。q 和 k 的表示被进一步使用 L2 范数归一化进行归一化,以确保特征值的稳定性。

每通道衰减 是通过一个低秩投影,其秩等于头部维度)和一个类似于 GDN 和 Mamba 中使用的衰减函数 来参数化的。

在通过进行输出投影之前,模型使用了一个逐头 RMSNorm 和一个数据依赖的门控机制,其参数化如下:

在这里,输出门采用了类似于遗忘门的低秩参数化,以确保公平的参数比较,同时保持与全秩门控相当的性能,并缓解注意力下沉问题。

混合模型架构

长上下文检索仍然是纯线性注意力的主要瓶颈,因此,本文将 KDA 与少数几个全全局注意力 (Full MLA) 层进行了混合。

对于 Kimi Linear,研究人员选择了一种逐层方法(交替使用整个层),而不是逐头方法(在层内混合不同的头),因为前者在基础设施简单性和训练稳定性方面更具优势。

经验表明,采用统一的 3:1 比例,即每 3 个 KDA 层重复 1 个全 MLA 层,可提供最佳的质量 - 吞吐量权衡。

MLA 层的无位置编码 (NoPE)

在 Kimi Linear 中,NoPE 被应用于所有全注意力层。这种设计将编码位置信息和近期偏好的全部责任委托给了 KDA 层。

因此,KDA 被确立为主要的位置感知算子,其扮演的角色类似于(甚至可以说强于)短卷积或 SWA 等辅助组件。这一发现与先前的研究结果一致,后者同样证明了:使用一个专用的位置感知机制来补充全局 NoPE 注意力,可以产生具有竞争力的长上下文性能。

值得注意的是,NoPE 提供了显著的实用优势,特别是对于 MLA 层。首先,NoPE 使得这些层在推理过程中可以转换为高效的纯多查询注意力。其次,它简化了长上下文训练,因为它避免了 RoPE(旋转位置编码)的参数调整需求,例如频率基调优或使用像 YaRN 这样的方法。

实验结果

对 Kimi Linear 关键组件的消融实验

作者通过将不同模型与 first-scale scaling law 模型(即 16 个注意力头、16 层)进行直接比较,开展了一系列消融研究。下表 1 记录了训练和验证的困惑度(PPL)。

从表中还可以得出以下信息:

  • 输出门控:移除门控会降低性能,swish 门控性能明显不如 Sigmoid。

  • 卷积层:卷积层在混合模型中仍然发挥着不可忽视的作用。

  • 混合比例:在测试的配置中,3:1 的比例(每 1 个 MLA 层对应 3 个 KDA 层)产生了最佳结果,实现了最低的训练损失和验证损失。

NoPE vs. RoPE 的对比结果见表 5。从中可以看出,采用 NoPE 的 Kimi Linear 在不同的长上下文基准测试中取得了最佳的平均分数。

Kimi Linear 的 scaling law

作者基于 Moonlight 架构,对一系列 MoE 模型进行了 scaling law 实验。在所有实验中,他们从 64 个专家中激活了 8 个,并使用了 Muon 优化器。详细信息和超参数列于表 2 中。

对于 MLA,他们遵循 Chinchilla scaling law,训练了五个不同规模的语言模型,并通过网格搜索仔细调整它们的超参数,以确保每个模型都能达到最佳性能。对于 KDA,他们保持了表 1 中验证的 3:1 这一最佳混合比例。除此之外,他们严格遵循 MLA 的训练配置,未做任何修改。如图 5 所示,与经过计算优化训练的 MLA 基线相比,Kimi Linear 的计算效率约为其 1.16 倍。作者预计,仔细的超参数调整将为 KDA 带来更优的缩放曲线。

与基线对比的主要实验结果

Kimi Linear @1.4T 结果:

  • 预训练结果

在表 3 中,团队使用一个 1.4T 的预训练语料库,比较了 Kimi Linear 模型与两个基线模型(MLA 和 hybrid GDN-H)。评估重点关注三个领域:通用知识、推理(数学和代码)以及中文任务。Kimi Linear 在几乎所有类别中都持续优于两个基线模型。

  • 通用知识: Kimi Linear 在所有关键基准(如 BBH, MMLU 和 HellaSwag)上均获得最高分。

  • 推理: 它在数学(GSM8K)和大多数代码任务(CRUXEval)上处于领先地位。然而,在 EvalPlus 上的得分略低于 GDN-H。

  • 中文任务: Kimi Linear 在 CEval 和 CMMLU 上取得了最高分。

总之,Kimi Linear 展现了最强的性能,使其成为短上下文预训练中全注意力架构的有力替代方案。

  • SFT 结果

在经历了相同的监督微调(SFT)流程后,Kimi Linear 在通用任务以及数学和代码任务上均表现出强劲性能,持续优于 MLA 和 GDN-H。

在通用任务中,Kimi Linear 全面领先,在各种 MMLU 基准、BBH 和 GPQA-Diamond 上均取得了最高分。在数学和代码任务中,它在 AIME 2025、HMMT 2025、PolyMath-en 和 LiveCodeBench 等高难度基准上超越了两个基线模型。

尽管在 MATH500 和 EvalPlus 等个别项目上存在微小例外,但 Kimi Linear 在各项任务中均显示出稳健的优势,证实了其相较于其他测试模型(GDN-H 和 MLA)的明显优越性。

  • 长上下文性能评估

团队在 128k 上下文长度下,评估了 Kimi Linear 相对于三个基线模型 ------MLA、GDN-H 和 Kimi Linear (RoPE)------ 在几个基准上的长上下文性能(见表 5)。

结果凸显了 Kimi Linear 在这些长上下文任务中的明显优势。它持续优于 MLA 和 GDN-H,在 RULER (84.3) 和 RepoQA (68.5) 上以显著优势取得了最高分。这种超越模式在除 LongBench V2 和 Frames 之外的大多数其他任务中也得以保持。

总体而言,Kimi Linear 取得了最高的平均分 (54.5),进一步巩固了其作为长上下文场景中领先注意力架构的有效性。

  • RL 结果

为了比较 Kimi Linear 和 MLA 的强化学习(RL)收敛特性,团队使用了 内部数学训练集进行了 RLVR,并在数学测试集(例如 AIME 2025, MATH500)上进行评估,同时保持算法和所有超参数一致,以确保公平的性能比较。

如图 6 所示,与 MLA 相比,Kimi Linear 展示了更高的效率。在训练集上,尽管两个模型起点相似,但 Kimi Linear 的训练准确率增长速度明显高于 MLA,且差距逐渐拉大。在测试集上也观察到了类似现象。例如,在 MATH500 和 AIME2025 上,Kimi Linear 相比 MLA 取得了更快、更好的提升。

总体而言,团队根据经验观察到,在强化学习下的推理密集型长文本生成中,Kimi Linear 的表现明显优于 MLA。

  • 总体研究结果总结

在预训练和 SFT 阶段,一个清晰的性能层级得以确立:Kimi Linear 优于 GDN-H,而 GDN-H 又优于 MLA。然而,这个层级在长上下文评估中发生了变化。虽然 Kimi Linear 保持了其领先地位,但 GDN-H 的性能有所下降,使其排名落后于 MLA。

此外,在 RL 阶段,Kimi Linear 也表现出优于 MLA 的性能。总体而言,Kimi Linear 在所有阶段始终名列前茅,确立了其作为全注意力架构的卓越替代方案的地位。

效率对比结果

作者在图 7a 和图 7b 中比较了全注意力 MLA、GDN-H 和 Kimi Linear 的训练及解码时间。

作者观察到,尽管融入了更精细的衰减机制,但 Kimi Linear 在预填充期间相较于 GDN-H 仅引入了可忽略不计的延迟开销。如图 7a 所示,它们的性能曲线几乎难以区分,这证实了 Kimi Linear 的方法保持了较高的效率。随着序列长度的增加,混合的 Kimi Linear 模型相较于 MLA 基线展现出明显的效率优势。虽然在较短长度(4k--16k)时其性能与 MLA 相当,但从 128k 开始,它的速度显著提升。这种效率差距在规模扩大时急剧拉大。如图 1b 所示,Kimi Linear 在解码阶段充分展现了其优势。在 1M 上下文长度下进行解码时,Kimi Linear 的速度是全注意力的 6 倍。

更多技术细节请参见原论文。

相关推荐
IT_陈寒6 小时前
JavaScript 性能优化:3个V8引擎隐藏技巧让你的代码提速50%
前端·人工智能·后端
故事挺秃然7 小时前
NLP模型优化
人工智能·自然语言处理·nlp
文火冰糖的硅基工坊7 小时前
[人工智能-大模型-78]:模型层技术 - 深度神经网络的网络架构的演进,这不仅是一条技术路线图,更是一部 “机器如何逐步逼近人类认知方式” 的进化史诗。
人工智能·架构·dnn
IT古董7 小时前
【第六章:项目实战之推荐/广告系统】3.精排算法-(2)精排算法模型精讲: DNN、deepFM、ESMM、PLE、MMOE算法精讲与实现- DNN 精排模型
人工智能·神经网络·dnn
IT古董7 小时前
【第六章:项目实战之推荐/广告系统】3.精排算法-(2)精排算法模型精讲: DNN、deepFM、ESMM、PLE、MMOE算法精讲与实现- PLE 模型
人工智能·神经网络·dnn
Serverless 社区7 小时前
吉利汽车携手阿里云函数计算,打造新一代 AI 座舱推理引擎
人工智能·阿里云·汽车
没有梦想的咸鱼185-1037-16637 小时前
AI大模型支持下的:ArcGIS数据处理、空间分析、可视化及多案例综合应用
人工智能·arcgis·chatgpt·数据分析
青春不败 177-3266-05207 小时前
AI+ArcGIS:数据处理、空间分析、可视化前沿技术应
人工智能·arcgis·gis·生态学·可视化·数据处理
新智元8 小时前
老黄亲自站台,英伟达编程神器!Cursor 2.0 自研模型狂飙 4 倍
人工智能·openai