Kimi Linear：让长上下文大模型真正实用的线性注意力架构

拖拖765
132 阅读4分钟

一、为什么我们需要新的注意力机制？

Transformer 已成为大语言模型（LLM）的核心，但它的 softmax 全注意力（Full Attention） 存在两个关键瓶颈：

  1. 计算复杂度高（O(n²)） —— 上下文越长，推理越慢。

  2. KV Cache 线性增长 —— 显存和内存消耗随上下文线性增加。

这些问题在长上下文（如百万 token）、强化学习（RL）推理、Agent 长时间对话等场景下尤其突出。

线性注意力（Linear Attention）理论上能将复杂度降到 O(n)，但过去它的 表达能力不足，无法与 full attention 媲美。

于是，Moonshot AI 提出了——Kimi Linear。

二、Kimi Linear 是什么？

Kimi Linear 是一种 混合线性注意力架构（Hybrid Linear Attention Architecture） ，

它首次在公平对比下，在 短上下文、长上下文、强化学习（RL） 等多种场景中全面超越 full attention。

它的核心模块是：

👉 Kimi Delta Attention (KDA) —— 一种结合了高表达力与高效率的线性注意力机制。

三、关键创新与技术突破

Kimi Linear 的技术创新主要集中在三个层面：

1️⃣ Kimi Delta Attention（KDA）—— 通道级可学习遗忘

在传统 Gated DeltaNet（GDN）中，每个注意力头只有一个遗忘率 αₜ。

而 KDA 做到了 Fine-grained Gating（细粒度门控） ：

每个通道维度都拥有独立遗忘率 αₜ[i]。

🧩 这意味着模型可以精确控制“哪些信息保留，哪些丢弃”。

它像是一个可学习的、神经化的“记忆遗忘机制”，

在长序列任务中表现更稳定、更精准。

2️⃣ DPLR 优化：高效且稳定的矩阵变体

KDA 使用一种特制的 DPLR（Diagonal-Plus-Low-Rank）矩阵：

  • 保留了表达力；

  • 避免了传统 DPLR 中昂贵的矩阵逆与二次分块计算；

  • 实现了 chunk 并行和 GPU 高效利用。

结果：速度提升约 2×，精度更高。****

3️⃣ 混合架构（Hybrid Design）：3:1 层间交替

“三层线性注意力 + 一层全注意力（Full MLA）”

这种结构：

  • 既保留全局信息传播（来自 MLA）；
  • 又保持线性层的高效率；
  • 内存占用降低高达 75% ；
  • 在 1M token 上下文时推理速度提升 6×。

4️⃣ 无显式位置编码（NoPE）

传统 Transformer 用 RoPE（旋转位置编码）表示位置信息。

Kimi Linear 直接移除 RoPE，让 KDA 的遗忘因子 αₜ 自主学习位置信息。

✅ 实验显示：Kimi Linear 的长文本外推能力比 RoPE 更强。

5️⃣ Attention 的新解释：在线梯度下降（Delta Rule）

论文从学习理论角度重新解释 Attention：

Attention 的记忆更新过程，其实等价于对一个记忆矩阵做在线梯度下降。

这将 Transformer、RNN、DeltaNet 统一在一个数学框架中。

KDA 的更新规则：

S_t = (I - β_t k_t k_t^T) \cdot \text{Diag}(α_t) S_{t-1} + β_t k_t v_t^T

其中：

  • βₜ：学习率

  • αₜ：遗忘率

  • Sₜ：记忆状态矩阵

这就是「Attention as Learning」的核心理念。

四、实验结果一览

场景指标MLA（Full）GDN-HKimi Linear
短上下文MMLU-Pro47.247.951.0
长上下文 (128k)RULER81.380.584.3
KV Cache 占用-100%75%25%
解码速度 (1M tokens)-1×2×6×

Kimi Linear 在几乎所有维度都领先。

五、实际应用场景

应用领域场景示例Kimi Linear 优势
🔹 长上下文 LLM文档/代码检索、论文问答百万级上下文、低显存
🔹 AI Agent 推理工具调用链、游戏代理、机器人连续记忆 + 实时响应
🔹 高并发推理服务AI 浏览器、API 接口、搜索摘要6× 推理加速、75% 缓存减少
🔹 多模态任务视频字幕生成、时间序列预测时序记忆优化
🔹 数学与科学推理MATH、AIME、STEM RL更稳定的收敛特性

Moonshot AI 已在其产品 Kimi Chat 中使用此架构，实现百万 token 上下文。

六、最小可运行 Demo

Kimi Linear 的完整模型开源在 Hugging Face 上，下面是一个最小可运行示例：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "moonshotai/Kimi-Linear-48B-A3B-Instruct"

# 1️⃣ 加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# 2️⃣ 编写输入
prompt = "Explain the key idea behind Kimi Linear in simple terms."

# 3️⃣ 生成输出
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(output[0], skip_special_tokens=True))

💡 环境要求：

  • Python ≥ 3.10

  • torch ≥ 2.2

  • transformers ≥ 4.46

  • fla-core ≥ 0.4.0

  • 一张 ≥ 24GB 显存的 GPU（或使用推理引擎如 vLLM）

如需快速体验，可使用：

👉 Hugging Face SpaceAttachment.tiff

七、总结

Kimi Linear 的意义在于，它不仅让线性注意力“能跑”，而是真正超越了全注意力。

它的关键特性：

✅ 表达力媲美 full attention

✅ 复杂度降至 O(n)

✅ 6× 推理加速

✅ 75% 内存节省

✅ 支持百万级上下文

它代表了一种趋势：未来 LLM 架构将不再是 “Transformer = Full Attention”，

而是 “Hybrid Attention = 高效 + 智能 + 可扩展”。

avatar
文章
阅读
粉丝