Blog Tags GitHub
Back to blog
efficient-attention mixture-of-experts long-context kv-cache

Kimi Linear: 表达能力强的高效注意力架构

标准 Transformer 注意力机制的二次复杂度在长上下文场景下成为瓶颈。虽然线性注意力方法（如 Mamba2、Gated DeltaNet）可以降低计算复杂度，但它们在表达能力和实际性能上通常不如全注意力机制。

Kimi Linear: 表达能力强的高效注意力架构

一、论文概述

项目内容
标题Kimi Linear: An Expressive, Efficient Attention Architecture
作者Kimi Team (Moonshot AI) - 59 位作者，由 Yu Zhang 领衔，核心贡献者包括 Songlin Yang (MIT)、Jiaxi Hu (HKUST-GZ)、Jianlin Su 等
论文https://arxiv.org/abs/2510.26692
发布2025-10-30 (v1), 2025-11-01 (v2)
代码https://github.com/MoonshotAI/Kimi-Linear
模型https://huggingface.co/moonshotai/Kimi-Linear-48B-A3B-Instruct

二、核心思想

问题定义

标准 Transformer 注意力机制的二次复杂度在长上下文场景下成为瓶颈。虽然线性注意力方法（如 Mamba2、Gated DeltaNet）可以降低计算复杂度，但它们在表达能力和实际性能上通常不如全注意力机制。

核心问题：能否设计一种混合线性注意力架构，在公平比较下首次超越全注意力？

解决方案概述

Kimi Linear 是一种混合线性注意力架构，在短上下文、长上下文和强化学习（RL）缩放场景下均超越全注意力。

核心组件：

  1. Kimi Delta Attention (KDA)：扩展 Gated DeltaNet，使用更细粒度的通道级门控机制
  2. 3:1 KDA:MLA 混合架构：交替使用 KDA 层和 Multi-Head Latent Attention (MLA) 层
  3. 专用 DPLR 变体：对角加低秩过渡矩阵的高效实现

关键优势：

  • 首次在公平比较下超越全注意力
  • KV cache 减少高达 75%
  • 1M 上下文解码吞吐量提升 6×
  • 在短上下文、长上下文和 RL 场景均表现优异

三、技术架构

Kimi Delta Attention (KDA)

核心公式（递归形式）：

St=(I−βtktktT)⋅Diag(αt)⋅St−1+βtktvtTS_t = (I - \beta_t k_t k_t^T) \cdot \text{Diag}(\alpha_t) \cdot S_{t-1} + \beta_t k_t v_t^T

ot=StTqto_t = S_t^T q_t

与 Gated DeltaNet (GDN) 的关键区别：

特性GDNKDA
门控粒度标量（head-wise）通道级对角矩阵
遗忘率每个 head 一个值每个特征维度独立
类比Mamba2Gated Linear Attention (GLA)

为什么重要：

  • 细粒度衰减使有限状态 RNN 内存的调控更精确
  • 每个通道可以独立决定遗忘程度
  • 可解释为可学习的、数据依赖的位置编码

神经网络参数化：

  • q, k：ShortConv + Swish 激活后 L2 归一化
  • v：ShortConv + Swish（无 L2 归一化）
  • αt\alpha_t：低秩投影 (W_up_alpha * W_down_alpha * x) + 衰减函数
  • βt\beta_t：Sigmoid(W_beta * x) 标量门控
  • 输出：head-wise RMSNorm + Sigmoid 数据依赖输出门控

合成任务结果：

  • KDA 在 Palindrome、MQAR、Stack 任务上始终达到最高精度
  • 序列长度从 256 增加到 2048 token 时表现优异
  • 在召回密集型任务上比 GDN 收敛更快
  • Mamba2（无 delta rule）在所有任务上失败

混合架构

架构设计：

  • 3:1 比例：每 3 个 KDA 层后跟 1 个 MLA 层
  • 层间混合（交替整层）而非层内混合（混合 head）
  • MLA 层使用 NoPE（无位置编码）

消融实验（混合比例）：

KDA:MLA 比例训练 PPL验证 PPL
0:1 (纯 MLA)9.455.77
1:19.295.66
3:19.235.65
7:19.235.70
15:19.345.82

关键设计选择：

  • NoPE for MLA 层：

    • MLA 层可转换为高效纯 MQA 推理
    • 简化长上下文训练（无需调整 RoPE 参数）
    • 位置信息完全由 KDA 层编码

  • MoE 配置：

    • 基于 Moonlight 架构
    • MoE 稀疏度增加到 32（256 个专家中激活 8 个，含 1 个共享专家）

模型规模：

  • 总参数：48B
  • 每次前向传播激活参数：3B

DPLR 过渡矩阵

通用 DPLR 公式：

St=(D−atbtT)⋅St−1+ktvtTS_t = (D - a_t b_t^T) \cdot S_{t-1} + k_t v_t^T

KDA 的专用变体：

D=Diag(αt),at=βtkt,bt=kt⋅αtD = \text{Diag}(\alpha_t), \quad a_t = \beta_t k_t, \quad b_t = k_t \cdot \alpha_t

关键约束：a 和 b 都绑定到 key 向量 k

效率提升：

  1. 减少二次分块步骤：从 4 步减少到 2 步
  2. 消除 3 个矩阵乘法：在块间和输出计算中
  3. 整体 ~100% 内核加速：相比通用 DPLR

块级算法：

  • 使用 WY 表示的并行公式
  • 块间递归 + 块内并行策略
  • UT 变换减少非矩阵乘法 FLOP

四、核心公式

KDA 状态更新

St=Diag(αt)⏟细粒度衰减⋅(I−βtktktT)⏟Householder 变换⋅St−1+βtktvtTS_t = \underbrace{\text{Diag}(\alpha_t)}_{\text{细粒度衰减}} \cdot \underbrace{(I - \beta_t k_t k_t^T)}_{\text{Householder 变换}} \cdot S_{t-1} + \beta_t k_t v_t^T

缩放定律

MLA: L=2.3092⋅C−0.0536L = 2.3092 \cdot C^{-0.0536}

Kimi Linear: L=2.2879⋅C−0.0527L = 2.2879 \cdot C^{-0.0527}

计算效率：Kimi Linear 相比 MLA 实现约 1.16× 计算效率

KV Cache 大小

KV Cache=dk×dv=128×128 per head\text{KV Cache} = d_k \times d_v = 128 \times 128 \text{ per head}

固定大小，独立于序列长度

五、实验结果

预训练结果（1.4T token）

基准MLAGDN-HKimi Linear
HellaSwag81.782.282.9
ARC-challenge64.666.567.3
MMLU71.672.273.8
MMLU-Pro47.247.951.0
BBH71.670.672.9
GSM8K83.781.783.9
CEval79.379.179.5

SFT 结果（1.4T token）

基准MLAGDN-HKimi Linear
MMLU75.775.677.0
MMLU-Pro65.764.867.4
GPQA-Diamond (Avg@8)57.158.662.1
AIME 2025 (Avg@64)20.621.121.3
HMMT 2025 (Avg@32)11.311.312.5
PolyMath-en (Avg@4)41.341.543.6
LiveCodeBench v625.125.426.0

长上下文结果（128K 上下文）

基准MLAGDN-HKimi Linear (RoPE)Kimi Linear
RULER81.380.578.884.3
MRCR22.623.922.029.6
HELMET-ICL88.085.588.090.0
RepoQA63.063.066.568.5
Long Code Arena (Lib)32.834.731.337.1
平均52.251.251.854.5

Kimi Linear @5.7T 结果（vs Moonlight）

基准Moonlight-InstructKimi Linear @5.7T
RULER@1M-94.8
MMLU-Pro43.872.7
AIME 2025-58.6
MATH50058.094.6
LiveCodeBench v611.945.7

RL 训练结果

  • Kimi Linear 在 RL 收敛上优于 MLA
  • 在 MATH500 和 AIME2025 上实现更快更好的改进
  • 差距随训练步骤逐渐扩大

效率结果

解码效率：

  • 1M token 时 TPOT：1.84ms vs MLA 的 11.48ms（6.3× 更快）
  • KV cache 减少：高达 75%
  • 解码吞吐量：高达 6×

预填充效率：

  • 1M 上下文：2.9× 加速 over MLA
  • 解码（batch=1）：2.3× 加速 over MLA

KDA 内核速度：

  • 相比通用 DPLR：近 2× 速度（序列长度达 64K）

六、核心创新总结

创新点说明优势
Kimi Delta Attention (KDA)通道级细粒度门控扩展 Gated DeltaNet更精确的 RNN 内存调控
专用 DPLR 变体a=b=k 约束，消除冗余计算~100% 内核加速
3:1 KDA:MLA 混合层间混合，MLA 使用 NoPE首次超越全注意力
可学习位置编码KDA 过渡矩阵类比 RoPE数据依赖的位置信息
高效块级算法WY 表示 + UT 变换高硬件利用率

七、技术影响

对注意力架构的改进

  • 首次超越全注意力：在公平比较下，短上下文、长上下文和 RL 场景均表现优异
  • 表达能力强：通道级门控提供更细粒度的内存控制
  • 效率大幅提升：75% KV cache 减少，6× 解码吞吐量

与现有方法对比

方法优势局限
全注意力 (MLA)表达能力强二次复杂度，KV cache 大
Mamba2线性复杂度无 delta rule，表达能力有限
Gated DeltaNet有 delta rule标量门控，粒度粗
Kimi Linear通道级门控，混合架构需要训练专用内核

实际应用价值

  • 长上下文处理：1M token 上下文高效处理
  • 推理加速：6× 解码吞吐量，降低延迟
  • RL 训练：更好的收敛性
  • 代码理解：RepoQA 和 Long Code Arena 表现优异

八、局限性

  1. 模型规模：48B 总参数，3B 激活参数，较小规模可能效果不同
  2. 训练数据量：5.7T token 大规模训练，数据需求高
  3. 内核实现：需要定制 CUDA 内核（已开源）
  4. 混合比例：3:1 比例可能因任务和规模而异
  5. NoPE 设计：MLA 层无位置编码，可能影响某些任务

九、相关工作

线性注意力

  • Gated Linear Attention (GLA)：通道级门控线性注意力
  • Gated DeltaNet (GDN)：delta rule + 标量门控
  • Mamba2：选择性状态空间模型

混合架构

  • Jamba：Mamba + Transformer 混合
  • Zamba：高效混合架构
  • Griffin：Google 的混合线性注意力

高效注意力

  • Multi-Head Latent Attention (MLA)：DeepSeek 的低秩注意力
  • Flash Attention：高效注意力实现
  • Ring Attention：分布式长上下文

十、参考资源

论文与代码

相关工作

  • Gated DeltaNet: Yang et al., 2024
  • Mamba2: Dao & Gu, 2024
  • MLA: DeepSeek-V2
  • Flash Linear Attention: fla-org

基准

  • RULER: Hsieh et al., 2024
  • MRCR: Microsoft
  • HELMET: Li et al., 2024
  • RepoQA: Liu et al., 2024
  • Long Code Arena: Bayer et al., 2024