Kimi Linear 高效注意力架构
基本信息
- 标题:Kimi Linear: An Expressive, Efficient Attention Architecture
- 作者:Kimi Team(Moonshot AI)
- 机构:Moonshot AI(月之暗面)
- 年份:2025
- arXiv:2510.26692
- 开源:KDA kernel + vLLM 实现、预训练和指令微调模型权重(GitHub: MoonshotAI/Kimi-Linear)
核心论点
- 首次在公平比较下超越全注意力:Kimi Linear 是首个在短上下文、长上下文和强化学习(RL)scaling 场景中,均在公平对比条件下超越全注意力(full attention)的混合线性注意力架构。
- 细粒度门控是关键:相比 Gated DeltaNet(GDN)的 head-wise 标量遗忘门,KDA 引入 channel-wise 对角门控(每个特征维度独立的衰减率),更精细地控制有限状态 RNN 的记忆,显著提升线性注意力的表达能力。
- 混合架构是实用方案:3:1 的 KDA-to-MLA 混合比例在性能和效率之间取得最优平衡,纯线性注意力在精确检索上仍有瓶颈,混合设计是当前最佳实践。
- 位置编码可以内化:KDA 层承担了位置信息编码的全部职责,全局注意力层(MLA)采用 NoPE(No Position Encoding),简化了长上下文训练并支持高效 MQA 推理。
- 线性注意力在 RL scaling 中表现优异:在数学 RL 训练中,Kimi Linear 的收敛速度和最终性能均显著优于全注意力 MLA 基线,表明线性注意力在推理密集型长序列生成中具有独特优势。
关键技术方法
Kimi Delta Attention(KDA)
KDA 的核心状态更新公式:
与标准注意力的关键区别:
- 线性复杂度:注意力计算从 降至 ,KV cache 固定为 (每头 128x128),不随序列长度增长
- Delta Rule:采用重建目标的在线梯度下降,通过 rank-1 修正更新(Householder 变换)精确纠正关联记忆,而非仅做加法积累
- Channel-wise 衰减:GDN 使用标量 (head-wise),KDA 使用向量 (channel-wise),类似 GLA 的细粒度门控,但通过约束 DPLR 结构避免了 GLA 的数值精度问题
- 硬件高效分块算法:通过将 DPLR 中的 约束,消除了二次分块需求,矩阵乘法运算减少约 100%,kernel 速度接近 DPLR 的 2 倍
与标准注意力的对比
| 维度 | 标准 Softmax 注意力 | KDA(线性注意力) |
|---|---|---|
| 时间复杂度 | ||
| KV cache | 随序列长度线性增长 | 固定 每头 |
| 位置编码 | 需要 RoPE 等显式编码 | 内化于衰减机制(数据依赖的乘性位置编码) |
| 表达能力 | 理论上界更高 | 有限状态容量,但 Delta Rule 提供修正能力 |
| 检索能力 | 强 | 较弱(通过混合架构弥补) |
模型架构
- 基于 Moonlight 架构(MoE),3B 激活参数 / 48B 总参数,8/256 专家激活
- 混合比例:3 层 KDA + 1 层 Full MLA(3:1),均匀交替
- MLA 层使用 NoPE,位置信息完全由 KDA 层捕获
- KDA 层包含:ShortConv + Swish 激活 + L2Norm(q/k)、低秩参数化的 channel-wise 衰减门、Sigmoid 输出门 + RMSNorm
主要结果
短上下文(1.4T tokens 预训练)
- MMLU-Pro:Kimi Linear 51.0 vs MLA 47.2 vs GDN-H 47.9
- MMLU:73.8 vs 71.6 vs 72.2
- BBH:72.9 vs 71.6 vs 70.6
- 几乎所有通用知识和推理基准均优于 MLA 和 GDN-H
长上下文(128k)
- RULER:84.3(最优)vs MLA 81.3 vs GDN-H 80.5
- MRCR:29.6 vs 22.6 vs 23.9
- 平均分:54.5 vs MLA 52.2 vs GDN-H 51.2
效率
- KV cache 减少最高 75%
- 1M 上下文解码吞吐量提升最高 6.3x(1.84ms vs 11.48ms per token)
- 预填充速度:与 GDN-H 持平,512k 以上显著快于 MLA(2.3x@512k, 2.9x@1M)
- Scaling law:Kimi Linear 相比 MLA 基线有约 1.16x 的计算效率优势
RL 训练
- 数学 RL 训练全程,Kimi Linear 的训练和测试准确率增长速度均显著高于 MLA
局限性
- 精确检索仍是瓶颈:纯线性注意力在精确记忆检索和复制任务上仍有不足,必须通过混合架构(引入全注意力层)弥补
- 有限状态容量:RNN 隐状态 的固定大小意味着理论上存在信息压缩上限,极长上下文下的信息损失是固有约束
- 混合比例敏感:7:1 或更高比例会导致验证性能下降,3:1 是经验最优但可能因任务而异
- EvalPlus 等个别基准表现略逊:在部分代码生成基准上,GDN-H 略优于 Kimi Linear,说明细粒度门控并非在所有场景都有正向收益
- 硬件生态依赖:需要专用 kernel 实现和 vLLM 集成才能发挥效率优势,当前硬件和推理框架仍以全注意力优化为主
与相关工作的关系
线性注意力演进
- Linear Attention(Katharopoulos 2020):基础形式,无衰减,状态无界增长
- RetNet(Sun 2023):数据无关标量衰减
- Mamba2 / SSD(Dao & Gu 2024):数据相关标量衰减,S6 结构
- GLA(Yang 2024):channel-wise 对角门控,但使用通用 DPLR 公式,计算效率受限
- DeltaNet(Yang 2024):引入 Delta Rule(重建目标的在线梯度下降),rank-1 修正更新
- Gated DeltaNet / GDN(Yang 2024):在 DeltaNet 基础上加标量遗忘门
- RWKV7(Peng 2025):类似 DPLR 结构,使用 形式的低秩更新
- Comba(Zhang 2025):DPLR 变体,使用数据相关衰减
KDA 的定位:在 GDN 基础上将标量门控升级为 channel-wise 对角门控,同时通过约束 DPLR 结构()避免了 GLA 级别 DPLR 的计算开销,实现了表达能力和硬件效率的兼顾。
混合架构
- Jamba(AI21, 2024)、Nemotron-H(NVIDIA, 2025):Mamba + Transformer 混合
- Zamba(Zyphra, 2024):SSM + Attention 混合
- Falcon-H(TII, 2025):近 NoPE 的高 base frequency RoPE
- SwanGPT(2025):RoPE 层与 NoPE 层交替
Kimi Linear 采用层间交替(inter-layer)混合而非层内混合(intra-layer),3:1 比例在系统简洁性和训练稳定性上更优。
稀疏注意力(对比方向)
- NSA(DeepSeek, 2025):chunk-level 动态稀疏选择,依赖 GQA 大头数配置
- MoBA(2025):top-k chunk 选择,利用 LSE 分数
- DSA / DeepSeek-V3.2(2025):token-level 稀疏,低精度 fp8 indexer
线性注意力与稀疏注意力是两条不同的高效长上下文路径:线性注意力通过压缩实现固定状态大小,稀疏注意力通过选择性检索保留完整 KV cache。两者不互斥,未来可能融合。
后续问题
- KDA 在更大规模(如 MoE 200B+)上的 scaling 行为
- channel-wise 衰减 vs head-wise 衰减在不同任务上的系统性对比
- 线性注意力与稀疏注意力的混合架构可能性
- KDA 在多模态场景(视觉 token 长序列)中的表现
- NoPE 策略在不同混合比例下的影响