跳转至

第12章 线性注意力

保留 softmax 的稀疏注意力,本质上只是减少每次查询要读取的 key 数量,但并没有解决 softmax 本身带来的底层约束。线性注意力另辟蹊径:直接替换 softmax,改写相似度计算形式,调换矩阵运算顺序,把原本 \(O(L^{2})\) 的二次复杂度,降低为随序列长度线性增长的 \(O(L)\)。

12.1 核函数

标准自注意力公式:

\[ \operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d}}\right)V \]

前面基础知识章节已经说明,softmax 依靠指数函数,属于强非线性运算,它作用在 \(QK^{\top}\) 算出的相似度矩阵上,卡在两次矩阵乘法中间,破坏了矩阵乘法的结合律。我们只能先计算昂贵的 \(L\times L\) 相似度矩阵,带来 \(O(L^{2}d)\) 的平方级算力开销,序列越长开销爆炸,同时还会产生注意力稀释问题。

线性注意力用核函数 \(\phi(\cdot)\) 替换 softmax。 原版注意力直接对原始向量做内积 \(q^{\top}k\),再套 softmax;线性注意力先通过核函数把向量 \(q,k\) 各自映射到新特征空间,再在新空间计算内积,以此衡量相似度:

\[ \operatorname{sim}(q,k)=\phi(q)^{\top}\phi(k) \]

常用核选择:早期使用 ReLU,也可使用保证输出非负的指数核;硬性要求:\(\phi\) 输出所有元素 \(\ge 0\)。 为什么要求非负?内积结果将作为加权权重,全部非负才能保证状态累加过程数值稳定。

当然这里如果直接拿原始 \(q\)、\(k\) 做内积不也能算相似度吗?可以,但原始向量内积是纯线性运算。如果模型所有运算都是线性,无论堆叠多少层,整体等价于单层矩阵乘法,不足以拟合人类语言这种高度复杂的非线性语义关系。

线性注意力的设计巧思:非线性由核函数 \(\phi\) 单独施加在每一个独立的 \(q\)、\(k\) 向量上,而不是作用在 \(QK^{\top}\) 得到的两两相似度矩阵中间。

做完核映射之后,表达式为:

\[ \operatorname{LinearAttention}(Q,K,V)=\frac{\bigl(\phi(Q)\phi(K)^{\top}\bigr)V}{\bigl(\phi(Q)\sum\phi(K)\bigr)} \]

\(\phi(Q),\phi(K)\) 完成映射后,剩下的全部都是矩阵乘法,满足矩阵结合律,可以调换计算顺序,优先计算 \(\phi(K)^{\top}V\)。 \(K^{\top}\) 是 \([d,L]\),\(V\) 是 \([L,d]\),\(\phi(K)^{\top}V\) 得到 \(d\times d\) 固定大小的矩阵,这个矩阵就是汇总全部历史信息的状态 \(S\):

\[ S_{t}=\sum_{i\le t}\phi(k_{i})v_{i}^{\top} \]

这里最关键的特点:\(S\) 永远是 \(d\times d\) 固定大小矩阵,不会随输入序列长度 \(L\) 变大而膨胀。

  • 原版注意力:必须算出 \(L\times L\) 相似度矩阵,复杂度 \(O(L^{2}d)\),文本变长,算力平方增长;
  • 线性注意力:计算量 \(O(L d^{2})\),计算开销随 token 数量线性增长。

在自回归推理的时候,不需要保存每一个历史 token 对应的 K、V。每新来一个 token,只需要更新这个固定大小的状态 \(S\),极大降低 KV 缓存的显存压力。

注:这里依然保留因果约束,\(t\) 时刻看不到未来位置。 朴素累加式线性注意力自带两个短板:旧信息无限累积、只能追加信息,无法修改历史记录。

12.2 门控遗忘

解决旧信息无限堆积的问题。 在更新状态前,用门控系数对旧状态做缩放衰减:久远内容的权重慢慢变小,近期信息保留更高权重。 衰减系数可以固定,也可以做成动态可学习:模型看输入文本,自主决定哪些信息要长期记住、哪些可以快速遗忘。

举例:闲聊类内容快速衰减遗忘,文档里的关键事实保留更久。

12.3 Delta 更新

12.1 的朴素写法是每来一个 token 就做一次外积累加:

\[ S_{t}=S_{t-1}+\phi(k_{t})v_{t}^{\top}. \]

(下面为写清更新规则,把 \(\phi(k)\) 仍记作 \(k\),与常见 DeltaNet 记号一致;核映射若在,就作用在进状态之前。)

这只解决「怎么把历史塞进固定矩阵」,不解决「塞进去之后怎么改」。同一条 key 方向上再写一次,新旧 \(v\) 在状态里叠在一起;读出来是糊掉的混合,没有「覆盖」语义。门控遗忘(12.2)能把整块旧状态按比例变淡,但变淡是全局衰减,不是针对某条关联做定点修正。

行程助手里更具体:前面写过「博物馆周一闭馆」,后面用户说「这周周一照常开」。标准注意力可以回头看两处原文再裁决;固定状态若只会累加,两句矛盾的 \(v\) 会叠在同一片 \(S\) 里,读出来两边都糊着一点。需要的是:用新事实去改写旧关联,而不是再贴一层。

Delta 更新(delta rule)的做法是:先问状态「按现在的 \(k_t\),你以为 \(v\) 是什么」,再只把误差写回去。

\[ \hat{v}_{t}=S_{t-1}^{\top}k_{t}, \qquad e_{t}=v_{t}-\hat{v}_{t}, \qquad S_{t}=S_{t-1}+\beta_{t}\,k_{t}\,e_{t}^{\top}. \]

逐步看:

  1. 读:\(\hat{v}_{t}=S_{t-1}^{\top}k_{t}\) 是用当前 key 从旧状态里取出的预测,形状与 \(v_t\) 相同。状态此时扮演一张可寻址的联想表:\(k\) 当地址,读出的是「这个地址上现在记着什么」。
  2. 比:\(e_{t}=v_{t}-\hat{v}_{t}\)。若状态里已经记对了,\(e_{t}\approx 0\),后面几乎不写,避免同一事实被反复糊加。若记错了或还是空的,\(e_{t}\) 就是要补上的差。
  3. 写:只把 \(\beta_{t}\,k_{t}\,e_{t}^{\top}\) 这一秩一更新加进 \(S\)。方向仍由 \(k_t\) 决定(改的是这条关联),幅度由误差 \(e_t\) 和写入系数 \(\beta_t\in(0,1]\) 决定。\(\beta_t=1\) 时,在理想情形下可以把该 key 上的读出拉向 \(v_t\);\(\beta_t\) 较小则改得更保守。

和朴素累加对照:累加每次写入完整的 \(k_t v_t^{\top}\);delta 写入的是 \(k_t(v_t-\hat{v}_{t})^{\top}\)。已经对齐的部分被减掉,重复出现的相同关联不再指数级堆厚;纠正时,\(e_t\) 指向「旧预测该往哪边扳」,等价于在状态里抵消一部分旧写入。

这就是 DeltaNet 的核心递推:把线性注意力的状态更新从「只追加」换成「按误差修正」。它本身还不带 12.2 那种对整块 \(S\) 的遗忘缩放;带上遗忘门 \(\alpha_t\) 之后,就是下一节的 GDN(Gated DeltaNet)。名字上:delta rule 是规则,DeltaNet 是用这条规则做层的模块,GDN 是再加门控的常用工程形态。

仍要记住:改写发生在固定大小的 \(S\) 里,不是回翻原始 token。容量上限还在,只是同地址上的冲突比瞎叠加好处理一些。

12.4 GDN

GDN 把门控遗忘和 Delta 增量更新融合到一套状态更新公式,Qwen3 系列大量层使用该结构。

\[ \widetilde{S}_{t-1}=\alpha_{t}\odot S_{t-1},\qquad S_{t}=\widetilde{S}_{t-1}+\beta_{t} k_{t}\bigl(v_{t}-\widetilde{S}_{t-1}^{\top}k_{t}\bigr)^{\top} \]
  • \(\alpha_{t}\):遗忘门,缩放旧状态,控制历史保留程度;
  • \(\beta_{t}\):写入门,控制本次新信息写入状态的强弱。

它可以动态做到:不重要的历史慢慢忘掉,遇到冲突事实时,修正状态里已经存储的旧内容。

12.5 天花板:固定状态带来的信息损失

不管怎么优化门控和 delta,核心局限无法消除:所有历史信息都压缩存入一个固定尺寸 \(d\times d\) 矩阵。 固定容量相当于有损压缩。多条不同 token 的信息融合在一起,原始独立的 token 细粒度信息丢失。 长文本做精确检索、查找变量名、定位原文细节这类「大海捞针」任务,纯线性注意力会弱于可以读取原始 KV 的稀疏注意力。


核心问答

Q1:线性注意力为什么不需要 KV Cache? 历史已累进固定形状的状态;推理带着状态走即可。

Q2:delta rule 解决了什么? 按 \(v-\hat v\) 写回,减少同地址糊状叠加,并允许改写旧关联。不取消 \(d\times d\) 容量上限;遗忘要靠门控(GDN)。

Q3:又快又省,为什么不全用它? 精确检索弱。主流答案是与精确注意力混合,见第 14 章。

Q4:和下一章的 SSM 是一回事吗? 推理时都是固定状态,出身不同。本章从注意力换序推过来;第 13 章从状态方程推过来。旗舰主栈目前多见本章的 GDN 一路。