线性注意力从头到尾讲清楚：从注意力矩阵到可递推状态
关键不是少算几次乘法，而是把 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^T 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^T 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^T V 原封不动提前算掉，而是把“核化后的 key/value 聚合”提前算掉。
这里有一条很重要的概念边界：有些 Linear Attention 直接选择新的可分解核，例如 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^\topz就是\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^T 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 的完全替代，因为它牺牲了一部分精确的全局匹配能力。