17

线性注意力——用递推状态替代 KV cache

阅读约 7 分钟第 17/60 课

上一章节看到 KV cache 随序列长度线性增长，是长上下文推理的瓶颈。线性注意力用一种完全不同的方式来处理历史信息——不存储每一个历史 token 的 K 和 V，而是把它们压缩成一个固定大小的状态。

标准注意力的瓶颈在哪

让我们重新看标准（softmax）注意力在推理时做的事：

新 token 的 Q  ×  [所有历史 K]  →  注意力分数（一个分数对应一个历史 token）
softmax(分数)  →  注意力权重
注意力权重  ×  [所有历史 V]  →  输出

核心操作是：新 token 要跟每一个历史 token 单独算一个分数。100 个历史 token 就算 100 次点积，10,000 个就算 10,000 次。计算量和内存都跟序列长度成正比。

softmax 是这个瓶颈的关键原因。因为 softmax 需要看到所有分数之后才能归一化（需要知道总和才能算出每个分数的占比），所以必须先算出所有点积，然后才能做 softmax。这迫使模型存储每一个历史 token 的 K 和 V。

线性注意力的思路

线性注意力的核心思想是：去掉 softmax，用一种可以增量更新的计算方式替代它。

标准注意力中，输出可以写成：

输出 = Σᵢ (softmax权重ᵢ × Vᵢ)

如果把 softmax 去掉（或者替换成一个可以分解的函数），公式就变成了：

输出 ∝ Σᵢ (φ(Q) · φ(Kᵢ)) × Vᵢ

这里 φ 是某种变换函数。关键在于，这个求和可以通过改变计算顺序来避免显式存储所有的 K 和 V。

具体来说，定义一个状态矩阵 S：

S = Σᵢ φ(Kᵢ) × Vᵢ的转置

这是一个固定大小的矩阵——不管有多少历史 token，S 的大小不变（维度 × 维度）。每来一个新的历史 token，只需要把它的贡献累加到 S 上：

S_new = S_old + φ(K_new) × V_new的转置

查询时，新 token 直接用 S 来得到输出：

输出 = φ(Q) × S

S 就是一个固定大小的「压缩记忆」。 它把所有历史信息浓缩成一个矩阵，查询时只需要一次矩阵乘法——不需要回看每一个历史 token。

代价是什么

压缩必然有信息损失。

标准注意力能精确地对每一个历史 token 分配不同的注意力权重——哪怕是 1000 步之前的某个特定词，只要 Q 和 K 匹配，它就能得到很高的注意力。

线性注意力把 1000 个 token 的信息压缩进了一个固定大小的矩阵 S。如果矩阵不够大，早期的信息会被后来的信息「覆盖」——类似于一个有限容量的笔记本，写满之后新的内容会冲掉旧的。

这意味着线性注意力在需要精确回忆远处特定细节时表现不如标准注意力。

KDA：Sophia 的线性注意力

Sophia 使用的线性注意力变体叫 KDA（Kernel Dynamic Attention）。它在基本的线性注意力框架上做了一些改进，但核心思路一样：维护一个递推状态，每个新 token 更新状态，查询时从状态中读取。

Sophia 的 28 层按 [KDA, KDA, KDA, MLA] × 7 的模式排列：每四层中前三层使用 KDA，第四层使用 MLA（标准的多头注意力变体，需要 KV cache）。

为什么不全用 KDA？因为线性注意力的压缩记忆有容量限制，完全靠它会损失太多精确检索的能力。每四层插入一层标准注意力，让模型每隔几层就有机会精确地回看整个历史序列——这是一个在效率和能力之间的折中。

这个设计的效果是：

KDA 层（21 层）：不需要 KV cache，只维护固定大小的递推状态
MLA 层（7 层）：需要 KV cache，但只有 7 层而不是 28 层

KV cache 的大小减少了 75%——因为只有四分之一的层需要存储它。这让 Sophia 在同等显存下能支持更长的上下文，或者在同等上下文长度下需要更少的显存。

递推状态有多大

KDA 的递推状态是一个矩阵，大小取决于头的维度。对 Sophia 来说，每个 KDA 头的状态矩阵大约是 96 × 96 = 9,216 个数字。16 个头 × 21 层 = 336 个状态矩阵，共约 310 万个数字。

用 BF16 存储 ≈ 6 MB。

对比上一章节算的 KV cache 对 4096 个 token：688 MB（28 层全用标准注意力时）。即使 Sophia 的 7 层 MLA 仍然需要一些 KV cache，总量也远小于全标准注意力的方案。

这个单元学了什么

六个章节走完了注意力机制的全部核心概念：

  1. 词嵌入（第 12 章节）：一张查找表，把 token ID 变成向量。
  2. 自注意力（第 13 章节）：Q、K、V 三个角色，点积算相关性，加权汇总。
  3. 缩放与 softmax（第 14 章节）：除以 √d 控制分数范围，softmax 转成概率分布。
  4. 多头注意力（第 15 章节）：16 个头各看不同的关系，计算量不变。
  5. KV cache（第 16 章节）：推理时存储历史 K、V，随序列长度线性增长。
  6. 线性注意力（本章节）：用固定大小的递推状态替代 KV cache，牺牲精确检索换取效率。

注意力机制让 token 之间能够交换信息——每个 token 的向量从「只包含自身信息」变成了「包含上下文信息」。但注意力不是 Transformer 的全部。每一层除了注意力之外，还有残差连接、归一化和前馈网络——这些是下一个单元的内容。