线性注意力从头到尾讲清楚：从注意力矩阵到可递推状态
关键不是少算几次乘法，而是把 token-to-token 交互改写成“先写入状态，再由 query 读取”
Linear Attention 的目标是避免显式构造 \(n\times n\) 的注意力矩阵，而是把“每个 query 与所有 key 的交互”改写成“query 读取一个累积状态”。
这句话比“Linear Attention 把复杂度从平方降到线性”更接近本质。复杂度下降只是结果，真正发生的结构变化是：标准 Attention 显式维护 token 与 token 之间的关系，而 Linear Attention 先把所有 key–value 信息写进一个状态，再让 query 从状态中读出结果。
这篇文章从标准 Attention 出发，把这次改写完整推一遍。读完后，你应该能回答四个问题：为什么不能直接先算 \(K^\top V\)；归一化需要保存什么额外状态；因果 Linear Attention 为什么像 RNN；以及它用固定状态换走 \(n\times n\) 矩阵时究竟牺牲了什么。
如果只看三句话
- 标准 Attention 是“每个 query 逐个匹配所有 key”，因此会产生 \(n\times n\) 的注意力矩阵。
- 如果相似度可以分解为 \(\operatorname{sim}(q,k)=\phi(q)^\top\phi(k)\)，就能利用结合律先计算 \(\sum_j\phi(k_j)v_j^\top\)，再由 query 读取。
- 在因果场景中，这个汇总量可以逐 token 更新，于是 Attention 变成了一个具有固定大小状态的递推模型。
最值得记住的图景是：
\[\boxed{\text{显式 token-to-token 注意力图}\quad\longrightarrow\quad\text{可累积的记忆状态}}\]1. 先把标准 attention 写清楚
标准 scaled dot-product attention 是：
\[ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V \]其中：
- \(Q \in \mathbb{R}^{n \times d_k}\)
- \(K \in \mathbb{R}^{n \times d_k}\)
- \(V \in \mathbb{R}^{n \times d_v}\)
- \(n\) 是序列长度
对第 \(i\) 个 token 来说，它的输出可以写成：
\[ o_i = \sum_{j=1}^{n}\alpha_{ij} v_j \]其中权重是：
\[ \alpha_{ij} = \frac{\exp(q_i^\top k_j / \sqrt{d_k})}{\sum_{l=1}^{n}\exp(q_i^\top k_l / \sqrt{d_k})} \]这就是 attention 的核心：\(q_i\) 和所有 \(k_j\) 做匹配，然后按匹配结果对 \(v_j\) 加权求和。
问题也在这里。
因为每个 \(q_i\) 都要和所有 \(k_j\) 比一遍，所以你会得到一个 \(n \times n\) 的注意力矩阵：
\[ QK^\top \in \mathbb{R}^{n \times n} \]所以长度维度上的代价是二次的。序列一长，时间和显存就都顶不住。
2. 为什么不能直接先算 \(K^\top V\)
如果没有 softmax，attention 其实就是普通矩阵乘法：
\[ O = QK^\top V \]这时由于矩阵乘法满足结合律，可以改写成：
\[ O = Q(K^\top V) \]这样一来，先算 \(K^\top V\)，就不需要显式构造 \(n \times n\) 的中间矩阵了。
但标准 attention 不是这个式子，而是：
\[ O = \mathrm{softmax}(QK^\top)V \]这里的 softmax 是按行做归一化的。它把每一行的打分变成一个概率分布，所以每个 \(o_i\) 都依赖于“这一行所有 token 的相对关系”。
这意味着：
\[ \mathrm{softmax}(QK^\top)V \neq Q(K^\top V) \]原因很简单：softmax 不是线性的，也不能从中间那一步直接拆出去。
所以“先算 \(K^\top V\)”这个想法，只有在没有 softmax 时才成立。
3. linear attention 的真正思路：让相似度可以分解
如果你想把 attention 变成线性复杂度，最自然的思路不是去硬拆矩阵乘法，而是想办法把注意力写成某种可分解形式。
更一般地，可以把归一化 attention 写成：
\[ o_i=\frac{\sum_{j=1}^{n}\operatorname{sim}(q_i,k_j)v_j}{\sum_{j=1}^{n}\operatorname{sim}(q_i,k_j)} \]如果相似度可以写成某个特征映射的内积：
\[ \operatorname{sim}(q,k)=\phi(q)^\top\phi(k) \]这样一来，注意力权重就不再是必须显式计算的两两匹配，而可以变成“先聚合 \(K\) 和 \(V\)，再让 \(Q\) 去读这个聚合结果”。
为了看清楚，先看分子：
\[ \sum_{j=1}^{n}\operatorname{sim}(q_i,k_j)v_j \]代入可分解相似度：
\[ \sum_{j=1}^{n}\phi(q_i)^\top \phi(k_j) v_j \]由于 \(\phi(q_i)\) 和求和变量 \(j\) 无关，可以把它提到外面：
\[ \phi(q_i)^\top \sum_{j=1}^{n}\phi(k_j)v_j \]这一步就是 linear attention 的核心。分母也可以用同样的方法拆开，并用另一个累积量完成归一化。
注意，这里不是把 \(K^\top V\) 原封不动提前算掉，而是把“核化后的 key/value 聚合”提前算掉。
这里有一条很重要的概念边界：有些 Linear Attention 直接选择新的可分解核，例如 \(\operatorname{ELU}(x)+1\)；另一些方法才是去近似 softmax 的指数点积核。前者重新定义了注意力相似度，后者试图逼近原来的 softmax attention，不能笼统地说成同一件事。
4. 详细推导：从逐 token 公式到可实现形式
我们从单个位置 \(i\) 开始。
标准 softmax attention 的输出是：
\[ o_i = \sum_{j=1}^{n}\frac{\exp(q_i^\top k_j)}{\sum_{l=1}^{n}\exp(q_i^\top k_l)}v_j \]把分子分母分开写：
\[ o_i = \frac{\sum_{j=1}^{n}\exp(q_i^\top k_j)v_j}{\sum_{l=1}^{n}\exp(q_i^\top k_l)} \]现在引入特征映射 \(\phi(\cdot)\)，近似：
\[ \exp(q_i^\top k_j)\approx \phi(q_i)^\top\phi(k_j) \]于是分子变成：
\[ \sum_{j=1}^{n}\phi(q_i)^\top\phi(k_j)v_j \]把 \(\phi(q_i)\) 提到外面：
\[ \phi(q_i)^\top\left(\sum_{j=1}^{n}\phi(k_j)v_j^\top\right) \]这里为了维度一致，把 \(v_j\) 写成列向量时，里面是一个外积累积。 定义：
\[ S=\sum_{j=1}^{n}\phi(k_j)v_j^\top \]那么分子就是：
\[ \phi(q_i)^\top S \]分母同理：
\[ \sum_{j=1}^{n}\phi(q_i)^\top\phi(k_j) = \phi(q_i)^\top\left(\sum_{j=1}^{n}\phi(k_j)\right) \]定义：
\[ z=\sum_{j=1}^{n}\phi(k_j) \]于是最终可以写成：
\[ o_i \approx \frac{\phi(q_i)^\top S}{\phi(q_i)^\top z + \varepsilon} \]这里加一个很小的 \(\varepsilon\)，是为了避免分母太小导致数值不稳定。
这就是 linear attention 最核心的形状。
它的意义非常明确：
- \(S\) 负责存“历史 key 和 value 的压缩摘要”
- \(z\) 负责存“历史 key 的归一化摘要”
- 当前 query \(q_i\) 只需要和这个摘要交互一次
于是原本每个 token 要和所有历史 token 两两交互，变成了：
- 先把所有历史信息压成一个固定大小状态；
- 再用当前 query 去读这个状态。
5. 矩阵形式更直观
把所有 token 一次写完，会更好看。
令：
\[ \Phi(Q)\in\mathbb{R}^{n\times r},\quad \Phi(K)\in\mathbb{R}^{n\times r} \]其中 \(r\) 是特征映射后的维度。 那么可以写成：
\[ O = \frac{\Phi(Q)\left(\Phi(K)^\top V\right)}{\Phi(Q)\left(\Phi(K)^\top \mathbf{1}\right)} \]这里：
- \(\Phi(K)^\top V \in \mathbb{R}^{r \times d_v}\)
- \(\Phi(K)^\top \mathbf{1} \in \mathbb{R}^{r}\)
\(\mathbf{1}\) 是全 1 向量，表示对所有 key 做求和。
这个式子非常关键，因为它直接告诉你复杂度怎么降下来：
- \(\Phi(K)^\top V\) 只需要扫一遍序列
- \(\Phi(K)^\top \mathbf{1}\) 也只需要扫一遍序列
- 最后 \(\Phi(Q)\) 再扫一遍序列
所以长度 \(n\) 上的复杂度从二次变成线性。
6. 复杂度到底降在哪里
标准 attention 的主要开销是构造 \(n \times n\) 的注意力矩阵。
如果忽略常数，复杂度可以粗略看成：
\[ O(n^2 d_k + n^2 d_v) \]而 linear attention 先算：
\[ \Phi(K)^\top V \]这是 \(r \times d_v\) 的累积，代价约为：
\[ O(n r d_v) \]再算：
\[ \Phi(Q)(\cdot) \]代价约为：
\[ O(n r d_v) \]所以只要 \(r\) 是固定的，长度维度就是线性的：
\[ O(n) \]严格说，上面只计算 attention mixing 的主要项；生成 \(\Phi(Q)\) 和 \(\Phi(K)\) 还会带来特征映射本身的成本。若 \(r\approx d_k\approx d_v\approx d\)，常见写法是 \(O(nd^2)\)。所谓“线性”是指它关于序列长度 \(n\) 线性，并不是其他维度都消失了。
7. causal linear attention 为什么还能递推
如果是自回归生成场景，我们通常要求不能看未来 token。
这时 linear attention 还有一个额外好处：它可以递推更新。
定义到时刻 \(t\) 为止的状态：
\[ S_t = \sum_{j=1}^{t}\phi(k_j)v_j^\top \]\[ z_t = \sum_{j=1}^{t}\phi(k_j) \]那么它们的递推式就是：
\[ S_t = S_{t-1} + \phi(k_t)v_t^\top \]\[ z_t = z_{t-1} + \phi(k_t) \]当前输出则是：
\[ o_t = \frac{\phi(q_t)^\top S_t}{\phi(q_t)^\top z_t + \varepsilon} \]这意味着在推理时，你不需要保存完整的 KV cache，只要保存固定大小的 \(S_t\) 和 \(z_t\)。
这就是 causal linear attention 在长上下文生成里特别有吸引力的原因：记忆占用和长度无关，而是和特征维度有关。
更严格地说，这个“固定大小状态”直接对应流式执行和逐 token 解码。如果训练时为了并行而保存每个位置的前缀状态，中间激活仍可能达到 \(O(nrd_v)\)；因此不能把推理阶段的常数状态，直接等同于所有训练实现都只占常数内存。
8. 一个最小实现长什么样
下面这个伪代码能把上面的公式对应起来：
import torch
def phi(x):
return torch.nn.functional.elu(x) + 1.0
def linear_attention(q, k, v, eps=1e-6):
# q: [B, H, N, R]
# k: [B, H, N, R]
# v: [B, H, N, D]
q = phi(q)
k = phi(k)
kv = torch.einsum("bhnr,bhnd->bhrd", k, v)
z = k.sum(dim=-2)
numerator = torch.einsum("bhnr,bhrd->bhnd", q, kv)
denominator = torch.einsum("bhnr,bhr->bhn", q, z)
return numerator / (denominator[..., None] + eps)这段代码的结构非常清楚：
kv对应 \(\sum_j \phi(k_j)v_j^\top\)z对应 \(\sum_j \phi(k_j)\)- 分子和分母都只做一次聚合读出
如果换成 causal 版本，就把 kv 和 z 改成逐步累积即可。
9. 两条不同路线：换一个核，还是近似 softmax
linear attention 的成败，很大程度上取决于 \(\phi(\cdot)\) 怎么选，但常见方法其实有两种不同目标。
直接定义可分解的注意力核
Katharopoulos 等人的 Linear Transformer 使用可分解的相似度，常见特征映射是：
\[ \phi(x)=\operatorname{ELU}(x)+1 \]它得到的是一种新的 kernel attention，通常不等于标准 softmax attention。正值映射可以让分母更容易保持稳定，但核的选择也直接限制了模型能表达怎样的相似度。
用随机特征近似 softmax 核
softmax 分子对应指数点积核：
\[ \exp(q^\top k/\sqrt{d_k}) \]Performer 使用 FAVOR+ 正交随机特征近似这个核，使它也能写成 \(\phi(q)^\top\phi(k)\)。它的目标是近似常规 full-rank softmax attention，而不是随意换一种相似度。
所以更准确的总结是：Linear Attention 是一类依靠可分解特征映射实现线性计算的机制；其中有些方法重新定义注意力核，有些方法专门近似 softmax 核。
10. 与 KV cache 的差异
标准 Transformer 做自回归推理时，需要保存每层所有历史 token 的 key 和 value：
\[ K_{\le t}\in\mathbb{R}^{t\times d_k},\qquad V_{\le t}\in\mathbb{R}^{t\times d_v} \]缓存大小随上下文长度 \(t\) 线性增长。Linear Attention 保存的则是：
\[ S_t\in\mathbb{R}^{r\times d_v},\qquad z_t\in\mathbb{R}^{r} \]只要 \(r\) 和 \(d_v\) 固定，这部分状态相对于序列长度就是常数级。
但“历史被压缩”既是优点，也是能力瓶颈。标准 Attention 可以针对当前 query 重新访问某个具体历史 token；Linear Attention 只能从聚合状态中读出历史。一旦不同 token 的信息在有限状态中互相干扰，后续 query 就很难把它们重新分开。
11. 一个三个 token 的直觉例子
假设某个 query \(q\) 对三个历史 token 的核权重分别是 \(w_1,w_2,w_3\)。普通的归一化加权输出为：
\[ o=\frac{w_1v_1+w_2v_2+w_3v_3}{w_1+w_2+w_3} \]如果 \(w_j=\phi(q)^\top\phi(k_j)\)，分子就能写成：
\[ w_1v_1+w_2v_2+w_3v_3=\phi(q)^\top\left[\phi(k_1)v_1^\top+\phi(k_2)v_2^\top+\phi(k_3)v_3^\top\right] \]括号里的内容与当前 query 无关，因此历史 token 可以先把信息写入同一个状态：
\[ S=\sum_j\phi(k_j)v_j^\top \]query 到来后只需读取一次 \(S\)，不再重新逐个扫描 \(k_1,k_2,k_3,\ldots,k_n\)。这就是“每个 query 与所有 key 交互”变成“query 读取累积状态”的最小例子。
12. 它为什么有用，又为什么不完全等价
linear attention 的优点很直接：
- 不需要显式构造 \(n \times n\) 注意力矩阵
- 长序列时更省显存、更省时间
- causal 场景下可以递推，适合流式生成
但它也有明显局限：
- 历史信息被压缩成固定状态，表达力弱于完整 softmax attention
- 重新定义核的方法不等于 softmax；随机特征方法则只是在近似 softmax
- 如果任务特别依赖“精确地回看某个历史 token”，linear attention 可能不如标准 attention
- 理论复杂度更低不保证短序列上墙钟时间更快，实际表现还取决于特征维度和硬件算子
所以它更像一种取舍：
你用更低的代价，换取更好的长序列可扩展性。
13. 和“先算 \(K^\top V\)”的直觉到底差在哪
这个问题值得单独收一下。
你的直觉其实抓到了“先聚合后读取”的骨架。
但标准 attention 里，不能直接聚合原始 \(K\) 和 \(V\)，因为 softmax 会把每个 query 对历史 token 的偏好重新归一化。
linear attention 做的事情是：
- 不保留逐对的相似度矩阵；
- 选择或近似一个能写成 \(\phi(q)^\top\phi(k)\) 的核；
- 再利用结合律，把历史项压成状态 \(S\) 和 \(z\)。
所以它本质上不是“把 KV 提前算掉”，而是“把 attention 改成可累积的核回归”。
14. 一句话总结
如果你只记一句话，那就是：
linear attention 不是简单换一下矩阵乘法顺序，而是选择或近似一个可分解的注意力核，从而把原本显式的 \(n \times n\) 交互，改成固定状态的累积与读取。
它适合长序列、流式生成、低显存推理；
但它也不是标准 attention 的完全替代，因为它牺牲了一部分精确的全局匹配能力。