本站内容由 AI 生成,可能存在错误。如发现问题,欢迎到 GitHub Issues 反馈。

Attention 的低秩结构与 Efficient Attention

Attention 的低秩结构与 Efficient Attention

更新于 2026-04-22

上一篇我们看到了”利用低秩性”的第一个方向:LoRA 将微调增量 ΔW\Delta W 参数化为低秩矩阵 BABA,把”权重矩阵经验上低秩”这个观察转化为参数效率。现在我们转向第二个方向——注意力矩阵本身的低秩结构。

标准 Transformer(Vaswani et al., 2017)的 self-attention 核心计算是一个 n×nn \times n 矩阵:

Attn(Q,K,V)=softmax ⁣(QKTdk)⏟n×n 注意力矩阵V\text{Attn}(Q, K, V) = \underbrace{\text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)}_{n \times n \text{ 注意力矩阵}} V

其中 Q,K,V∈Rn×dkQ, K, V \in \mathbb{R}^{n \times d_k} 分别是 query、key、value 矩阵(nn 是序列长度,dkd_k 是 head 维度),softmax\text{softmax} 按行操作。这个 n×nn \times n 矩阵的计算和存储需要 O(n2)O(n^2) 的复杂度——当 n=4096n = 4096(长文档)或 n=16384n = 16384(长上下文模型)时,n2n^2 成为瓶颈。

一个关键的经验观察:这个 n×nn \times n 注意力矩阵在实践中接近低秩——它的大部分”能量”集中在少数几个奇异值上。这不是偶然的,而是有清晰的数学原因。

本文将回答三个问题:

  1. 为什么低秩? 从 QKTQK^T 的秩上界和 softmax 的”锐化”效应出发
  2. 怎么利用? 两种代表性方案——Linformer(显式低秩投影)和 Performer(随机特征近似 softmax kernel)
  3. 连接 Kernel? QKTQK^T 就是 query-key 的 Gram 矩阵,softmax 是一种 kernel——这连回 Art. 20 Kernel

为什么注意力矩阵接近低秩

原因一:QKTQK^T 的秩上界

注意力得分矩阵(softmax 之前)是:

S=QKTdk∈Rn×nS = \frac{QK^T}{\sqrt{d_k}} \in \mathbb{R}^{n \times n}

其中 Q∈Rn×dkQ \in \mathbb{R}^{n \times d_k},K∈Rn×dkK \in \mathbb{R}^{n \times d_k}。

由矩阵秩的性质(Art. 3 SVD),矩阵乘积的秩不超过任一因子的秩:

rank(QKT)≤min⁡ ⁣(rank(Q),  rank(K))≤min⁡(n,  dk)\text{rank}(QK^T) \leq \min\!\big(\text{rank}(Q),\; \text{rank}(K)\big) \leq \min(n,\; d_k)

在标准 Transformer 中,dk=dmodel/hd_k = d_{\text{model}} / h,其中 dmodeld_{\text{model}} 是模型维度,hh 是注意力头数。以 GPT-2 为例,dmodel=768d_{\text{model}} = 768,h=12h = 12,所以 dk=64d_k = 64。而序列长度 nn 通常为 512 到 4096。

当 dk≪nd_k \ll n 时(这是实践中的常见情况),我们有:

rank(QKT)≤dk≪n\text{rank}(QK^T) \leq d_k \ll n

一个 n×nn \times n 矩阵的秩最多为 nn,但 QKTQK^T 的秩被限制在 dkd_k——远小于 nn。这意味着 QKTQK^T(softmax 之前的注意力得分矩阵)天然就是低秩的,它最多只有 dkd_k 个非零奇异值。

逐项理解这个上界:

  • Q∈Rn×dkQ \in \mathbb{R}^{n \times d_k}:nn 个 query 向量,每个 dkd_k 维。QQ 的列空间至多 dkd_k 维。
  • K∈Rn×dkK \in \mathbb{R}^{n \times d_k}:nn 个 key 向量,每个 dkd_k 维。KK 的列空间至多 dkd_k 维。
  • QKTQK^T 的列空间 ⊆Q\subseteq Q 的列空间(因为 QKT=Q(KT)QK^T = Q(K^T),每一列都是 QQ 的列的线性组合)。
  • 因此 rank(QKT)≤rank(Q)≤dk\text{rank}(QK^T) \leq \text{rank}(Q) \leq d_k。

这是一个纯代数的结论,不依赖于数据分布或训练过程——只要 dk<nd_k < n,QKTQK^T 就不可能是满秩的。

原因二:softmax 的”锐化”效应

QKTQK^T 的低秩性还只是开始。应用 softmax 之后,注意力矩阵的有效秩通常进一步降低。

softmax 按行操作,将每行变成概率分布:

Aij=exp⁡(Sij)∑l=1nexp⁡(Sil)A_{ij} = \frac{\exp(S_{ij})}{\sum_{l=1}^{n}\exp(S_{il})}

其中 S=QKT/dkS = QK^T / \sqrt{d_k} 是缩放后的得分矩阵。

softmax 的指数函数会放大大值、压缩小值。如果 SS 的某一行中有一个或几个特别大的元素,softmax 后这些位置会接近 1,其余位置接近 0。这使得注意力矩阵的每一行接近一个 one-hot 向量——而 one-hot 向量的集合张成的空间维度很低。

直觉:想象一个长度为 1024 的序列,每个 token 主要关注 3-5 个其他 token。那么注意力矩阵的每行约有 3-5 个显著非零值,其余接近零。这样的矩阵显然可以被低秩矩阵很好地近似——它的”有效模式”远少于 1024。

定量地说:设注意力矩阵 A=softmax(S)A = \text{softmax}(S) 的 SVD 为 A=∑i=1rσiuiviTA = \sum_{i=1}^{r}\sigma_i \mathbf{u}_i\mathbf{v}_i^T(r=rank(A)r = \text{rank}(A))。“有效秩”可以用能量集中度衡量:

∑i=1kσi2∑i=1rσi2≥0.90即使 k≪r\frac{\sum_{i=1}^{k}\sigma_i^2}{\sum_{i=1}^{r}\sigma_i^2} \geq 0.90 \quad \text{即使 } k \ll r

在实践中,即使 r=dk=64r = d_k = 64,通常 k=5–10k = 5\text{–}10 个奇异值就能捕获 90% 以上的能量。softmax 的锐化效应让有效秩远低于理论秩上界 dkd_k。

可视化:attention 矩阵的奇异值衰减

下面的交互组件用合成数据演示了 attention 矩阵 softmax(QKT/dk)\text{softmax}(QK^T/\sqrt{d_k}) 的奇异值衰减。QQ 和 KK 的每个元素从标准正态分布 N(0,1)\mathcal{N}(0, 1) 独立采样。

调整 nn(序列长度)和 dkd_k(head 维度),观察:

  • dk≪nd_k \ll n 时(如 n=64,dk=4n = 64, d_k = 4):奇异值在第 dkd_k 个之后骤降到零——这是 rank(QKT)≤dk\text{rank}(QK^T) \leq d_k 的直接体现
  • dk≥nd_k \geq n 时(如 n=32,dk=64n = 32, d_k = 64):QKTQK^T 可以满秩,但 softmax 的锐化效应仍使奇异值快速衰减
  • 勾选”显示累积能量曲线”,看前 kk 个奇异值集中了多少能量
序列长度 n
Head 维度 d_k
0.000.250.500.751.0018162432奇异值序号 iσᵢ / σ₁(归一化)rank ≤ 4Attention 矩阵的奇异值衰减softmax(QKᵀ/√d_k),Q、K ~ N(0,1) (n=32, d_k=4)
rank(QKᵀ) ≤ min(n, d_k) = 4,softmax 后有效秩更低
d_k ≪ n 时,attention 矩阵天然低秩

读图要点:

  • 红色虚线标注理论秩上界 rank≤min⁡(n,dk)\text{rank} \leq \min(n, d_k)。当 dk<nd_k < n 时,奇异值在此处截断。
  • 即使在 dk≥nd_k \geq n 的情况下(没有红色虚线),奇异值仍然快速衰减——这是 softmax 锐化效应的体现。
  • 在累积能量视图中,少数几个奇异值就集中了绝大部分能量——这就是 efficient attention 方法的理论基础:既然注意力矩阵接近低秩,就不需要精确计算完整的 n×nn \times n 矩阵。

Gram 矩阵视角:连接 Kernel

在进入 Linformer 和 Performer 之前,先建立一个重要的概念连接。

在 Art. 20 Kernel 中,我们学到:给定 nn 个数据点 x1,…,xn\mathbf{x}_1, \ldots, \mathbf{x}_n,Gram 矩阵(Gram matrix)定义为 Gij=xiTxjG_{ij} = \mathbf{x}_i^T \mathbf{x}_j,即所有数据点两两的内积。Kernel 矩阵 Kij=k(xi,xj)K_{ij} = k(\mathbf{x}_i, \mathbf{x}_j) 是 Gram 矩阵在特征空间中的推广。

现在回看 attention 得分矩阵:

Sij=qiTkjdkS_{ij} = \frac{\mathbf{q}_i^T \mathbf{k}_j}{\sqrt{d_k}}

其中 qi∈Rdk\mathbf{q}_i \in \mathbb{R}^{d_k} 是第 ii 个 query 向量(QQ 的第 ii 行),kj∈Rdk\mathbf{k}_j \in \mathbb{R}^{d_k} 是第 jj 个 key 向量(KK 的第 jj 行)。

QKTQK^T 就是 query 和 key 的”交叉 Gram 矩阵”——第 (i,j)(i, j) 元素是 qi\mathbf{q}_i 和 kj\mathbf{k}_j 的内积。如果 Q=KQ = K(self-attention 中它们来自同一输入的不同投影,但数值不同),QKTQK^T 类似于 Gram 矩阵 XXTXX^T。

进一步,应用 softmax 之后:

Aij=exp⁡(qiTkj/dk)∑lexp⁡(qiTkl/dk)A_{ij} = \frac{\exp(\mathbf{q}_i^T \mathbf{k}_j / \sqrt{d_k})}{\sum_{l}\exp(\mathbf{q}_i^T \mathbf{k}_l / \sqrt{d_k})}

分子 exp⁡(qiTkj/dk)\exp(\mathbf{q}_i^T \mathbf{k}_j / \sqrt{d_k}) 可以看成一种核函数:它将 query-key 内积通过指数函数映射为非负的”相似度”。具体来说,定义:

kSM(q,k)=exp⁡ ⁣(qTkdk)k_{\text{SM}}(\mathbf{q}, \mathbf{k}) = \exp\!\left(\frac{\mathbf{q}^T \mathbf{k}}{\sqrt{d_k}}\right)

这就是所谓的 softmax kernel——它是一个正定核(可以证明它对应一个无穷维的特征映射)。注意力矩阵的分子 exp⁡(QKT/dk)\exp(QK^T / \sqrt{d_k}) 是 softmax kernel 的 kernel 矩阵,每一行再做归一化(除以行和)得到注意力权重。

这个视角的意义:

  • 在 Art. 20 Kernel 中我们看到,kernel 矩阵 K=ΦΦTK = \Phi\Phi^T,其中 Φ\Phi 是数据在特征空间中的表示。如果我们能找到一个显式的有限维特征映射 ϕ\phi 使得 ϕ(q)Tϕ(k)≈kSM(q,k)\phi(\mathbf{q})^T\phi(\mathbf{k}) \approx k_{\text{SM}}(\mathbf{q}, \mathbf{k}),就可以避免计算 n×nn \times n 的 kernel 矩阵——这正是 Performer 的核心思路。
  • Kernel 矩阵的低秩性由核函数的特征值衰减速度决定(Mercer 展开中 λi\lambda_i 的衰减)。softmax kernel 的特征值快速衰减,对应了注意力矩阵的低秩性。

Art. 20 的预告兑现

在 Art. 20 Kernel 的”前瞻”部分,我们写道:“Transformer 的注意力矩阵本质上是一个 kernel 矩阵——query 和 key 的内积经过 softmax 变成非负的相似度。线性注意力的核心思想正是用显式特征映射 ϕ\phi 替代 softmax,使计算从 O(n2)O(n^2) 降到 O(n)O(n)——这是 kernel trick 的逆向应用。” 现在我们来兑现这个承诺。

Linformer:显式低秩投影

核心思想

Wang et al. (2020) 的 Linformer 方法直截了当:既然注意力矩阵是低秩的,就用投影矩阵将 key 和 value 的序列维度从 nn 降到 kk(k≪nk \ll n),从而避免构造完整的 n×nn \times n 矩阵。

标准 attention(复杂度 O(n2dk)O(n^2 d_k)):

Attn(Q,K,V)=softmax ⁣(QKTdk)V\text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V

Linformer(复杂度 O(nkdk)O(nk d_k)):

Attn~(Q,K,V)=softmax ⁣(Q(EK)Tdk)(FV)\widetilde{\text{Attn}}(Q, K, V) = \text{softmax}\!\left(\frac{Q(EK)^T}{\sqrt{d_k}}\right) (FV)

其中 E,F∈Rk×nE, F \in \mathbb{R}^{k \times n} 是两个投影矩阵,kk 是投影维度(超参数,k≪nk \ll n)。

逐项理解

让我们仔细拆解每一步。

EK∈Rk×dkEK \in \mathbb{R}^{k \times d_k}:EE 将 key 矩阵 K∈Rn×dkK \in \mathbb{R}^{n \times d_k} 的序列维度从 nn 压缩到 kk。原来有 nn 个 key 向量,现在变成 kk 个”汇总 key”(每个是原始 key 的线性组合)。

Q(EK)T∈Rn×kQ(EK)^T \in \mathbb{R}^{n \times k}:注意力得分矩阵从 n×nn \times n 缩小为 n×kn \times k。每个 query 只需要和 kk 个汇总 key 计算内积,而不是 nn 个原始 key。

softmax(⋅)∈Rn×k\text{softmax}(\cdot) \in \mathbb{R}^{n \times k}:softmax 按行操作,每行是一个 kk 维的概率分布。

FV∈Rk×dkFV \in \mathbb{R}^{k \times d_k}:FF 将 value 矩阵 V∈Rn×dkV \in \mathbb{R}^{n \times d_k} 的序列维度从 nn 压缩到 kk。kk 个”汇总 value”。

最终结果 ∈Rn×dk\in \mathbb{R}^{n \times d_k}:n×kn \times k 的注意力权重乘以 k×dkk \times d_k 的汇总 value,得到 n×dkn \times d_k 的输出。

复杂度分析:

步骤标准 AttentionLinformer
注意力得分QKTQ K^T:O(n2dk)O(n^2 d_k)Q(EK)TQ(EK)^T:O(nkdk)O(nk d_k)
softmaxO(n2)O(n^2)O(nk)O(nk)
加权求和A⋅VA \cdot V:O(n2dk)O(n^2 d_k)A~⋅FV\tilde{A} \cdot FV:O(nkdk)O(nk d_k)
总计O(n2dk)O(n^2 d_k)O(nkdk)O(nk d_k)

当 k=O(1)k = O(1)(不随 nn 增长)时,Linformer 将 attention 的复杂度从 O(n2)O(n^2) 降到 O(n)O(n)。

理论保证

Wang et al. (2020) 证明了以下近似界(Theorem 1,Linformer 论文第 3 节):

对随机投影矩阵 EE(如从正态分布采样再归一化),如果投影维度 k=Ω(dklog⁡dk/ε2)k = \Omega(d_k \log d_k / \varepsilon^2),则以高概率有:

∥softmax ⁣(QKTdk)−softmax ⁣(Q(EK)Tdk)∥≤ε\left\|\text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) - \text{softmax}\!\left(\frac{Q(EK)^T}{\sqrt{d_k}}\right)\right\| \leq \varepsilon

关键洞察:kk 的需求只与 dkd_k 有关,不依赖于 nn。这正是因为 QKTQK^T 的秩不超过 dkd_k——无论序列多长,注意力矩阵的”内在维度”都被 dkd_k 限制。这与 Art. 23 学习算子 中讨论的”内禀维度”(intrinsic dimension)本质相同。

实践考量

Linformer 的投影矩阵 E,FE, F 有几种选择:

  • 随机高斯投影:Eij∼N(0,1/k)E_{ij} \sim \mathcal{N}(0, 1/k),固定不训练。理论保证最强,但实践中不一定最优。
  • 可学习投影:将 E,FE, F 作为可训练参数。通常比随机投影效果更好。
  • 共享投影:令 E=FE = F,或所有层共享同一对 (E,F)(E, F)。减少参数量。

实验表明(Wang et al., 2020, Table 2),即使 kk 取到 n/4n/4 甚至 n/8n/8,模型性能与标准 attention 几乎无差别。

Performer:随机特征近似 softmax kernel

从 kernel 视角出发

Linformer 直接压缩 KK 和 VV 的序列维度。Performer(Choromanski et al., 2021)采取了不同的路线:它从 kernel 视角出发,用随机特征映射(random feature map)近似 softmax kernel,从而改变计算的结合律,避免显式构造 n×nn \times n 矩阵。

回忆 softmax attention 的计算(去掉归一化因子,先看分子):

Attnunnorm=exp⁡ ⁣(QKTdk)V\text{Attn}_{\text{unnorm}} = \exp\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V

其中 exp⁡(QKT/dk)\exp(QK^T / \sqrt{d_k}) 的第 (i,j)(i, j) 元素是 kSM(qi,kj)=exp⁡(qiTkj/dk)k_{\text{SM}}(\mathbf{q}_i, \mathbf{k}_j) = \exp(\mathbf{q}_i^T\mathbf{k}_j / \sqrt{d_k})。

如果存在一个特征映射 ϕ:Rdk→Rm\phi: \mathbb{R}^{d_k} \to \mathbb{R}^m 使得:

kSM(q,k)=exp⁡ ⁣(qTkdk)≈ϕ(q)Tϕ(k)k_{\text{SM}}(\mathbf{q}, \mathbf{k}) = \exp\!\left(\frac{\mathbf{q}^T\mathbf{k}}{\sqrt{d_k}}\right) \approx \phi(\mathbf{q})^T \phi(\mathbf{k})

那么 kernel 矩阵可以分解为:

exp⁡ ⁣(QKTdk)≈ϕ(Q)⋅ϕ(K)T\exp\!\left(\frac{QK^T}{\sqrt{d_k}}\right) \approx \phi(Q) \cdot \phi(K)^T

其中 ϕ(Q)∈Rn×m\phi(Q) \in \mathbb{R}^{n \times m} 表示对 QQ 的每一行分别应用 ϕ\phi,ϕ(K)∈Rn×m\phi(K) \in \mathbb{R}^{n \times m} 类似。

关键的结合律变换:标准计算是 (ϕ(Q)⋅ϕ(K)T)⋅V\big(\phi(Q) \cdot \phi(K)^T\big) \cdot V——先算 n×nn \times n 矩阵,再乘 VV,复杂度 O(n2m)O(n^2 m)。但矩阵乘法满足结合律,我们可以改为:

ϕ(Q)⋅(ϕ(K)T⋅V)\phi(Q) \cdot \big(\phi(K)^T \cdot V\big)

先算 ϕ(K)TV∈Rm×dk\phi(K)^T V \in \mathbb{R}^{m \times d_k}(复杂度 O(nmdk)O(nmd_k)),再乘 ϕ(Q)\phi(Q)(复杂度 O(nmdk)O(nmd_k))。总复杂度 O(nmdk)O(nmd_k)——如果 mm 与 nn 无关,就从 O(n2)O(n^2) 降到了 O(n)O(n)。

FAVOR+ 机制

Performer 使用的特征映射称为 FAVOR+(Fast Attention Via positive Orthogonal Random features)。具体构造为:

ϕ(x)=exp⁡(−∥x∥2/2)m(exp⁡(ω1Tx)exp⁡(ω2Tx)⋮exp⁡(ωmTx))\phi(\mathbf{x}) = \frac{\exp(-\|\mathbf{x}\|^2 / 2)}{\sqrt{m}} \begin{pmatrix} \exp(\boldsymbol{\omega}_1^T \mathbf{x}) \\ \exp(\boldsymbol{\omega}_2^T \mathbf{x}) \\ \vdots \\ \exp(\boldsymbol{\omega}_m^T \mathbf{x}) \end{pmatrix}

其中 ω1,…,ωm∈Rdk\boldsymbol{\omega}_1, \ldots, \boldsymbol{\omega}_m \in \mathbb{R}^{d_k} 是随机采样的特征方向。

为什么这个 ϕ\phi 能近似 softmax kernel? 利用高斯积分恒等式(Choromanski et al., 2021, Lemma 1):

exp⁡(qTk)=exp⁡ ⁣(∥q∥22)exp⁡ ⁣(∥k∥22)⋅Eω∼N(0,I) ⁣[exp⁡(ωTq)⋅exp⁡(ωTk)]\exp(\mathbf{q}^T\mathbf{k}) = \exp\!\left(\frac{\|\mathbf{q}\|^2}{2}\right) \exp\!\left(\frac{\|\mathbf{k}\|^2}{2}\right) \cdot \mathbb{E}_{\boldsymbol{\omega} \sim \mathcal{N}(\mathbf{0}, I)}\!\left[\exp(\boldsymbol{\omega}^T\mathbf{q}) \cdot \exp(\boldsymbol{\omega}^T\mathbf{k})\right]

逐步推导这个恒等式:

  1. exp⁡(qTk)\exp(\mathbf{q}^T\mathbf{k}) 可以改写为一个关于高斯随机向量的期望。考虑 ω∼N(0,Idk)\boldsymbol{\omega} \sim \mathcal{N}(\mathbf{0}, I_{d_k})。
  2. E[exp⁡(ωTq)]=exp⁡(∥q∥2/2)\mathbb{E}[\exp(\boldsymbol{\omega}^T\mathbf{q})] = \exp(\|\mathbf{q}\|^2/2)——这是高斯分布的矩生成函数。
  3. 更一般地,对联合分布中的两个线性函数:E[exp⁡(ωTq+ωTk)]=exp⁡(∥q+k∥2/2)=exp⁡(∥q∥2/2)exp⁡(∥k∥2/2)exp⁡(qTk)\mathbb{E}[\exp(\boldsymbol{\omega}^T\mathbf{q} + \boldsymbol{\omega}^T\mathbf{k})] = \exp(\|\mathbf{q} + \mathbf{k}\|^2/2) = \exp(\|\mathbf{q}\|^2/2) \exp(\|\mathbf{k}\|^2/2) \exp(\mathbf{q}^T\mathbf{k})。
  4. 因此 exp⁡(qTk)=exp⁡(−∥q∥2/2)exp⁡(−∥k∥2/2)⋅E[exp⁡(ωT(q+k))]\exp(\mathbf{q}^T\mathbf{k}) = \exp(-\|\mathbf{q}\|^2/2) \exp(-\|\mathbf{k}\|^2/2) \cdot \mathbb{E}[\exp(\boldsymbol{\omega}^T(\mathbf{q} + \mathbf{k}))]。
  5. 将 exp⁡(ωT(q+k))\exp(\boldsymbol{\omega}^T(\mathbf{q}+\mathbf{k})) 拆成 exp⁡(ωTq)⋅exp⁡(ωTk)\exp(\boldsymbol{\omega}^T\mathbf{q}) \cdot \exp(\boldsymbol{\omega}^T\mathbf{k}),得到上述恒等式。

用 mm 个独立采样 ω1,…,ωm\boldsymbol{\omega}_1, \ldots, \boldsymbol{\omega}_m 做 Monte Carlo 近似:

exp⁡(qTk)≈1m∑j=1m[exp⁡ ⁣(−∥q∥22)exp⁡(ωjTq)]⋅[exp⁡ ⁣(−∥k∥22)exp⁡(ωjTk)]\exp(\mathbf{q}^T\mathbf{k}) \approx \frac{1}{m}\sum_{j=1}^{m} \left[\exp\!\left(-\frac{\|\mathbf{q}\|^2}{2}\right)\exp(\boldsymbol{\omega}_j^T\mathbf{q})\right] \cdot \left[\exp\!\left(-\frac{\|\mathbf{k}\|^2}{2}\right)\exp(\boldsymbol{\omega}_j^T\mathbf{k})\right]

这正是 ϕ(q)Tϕ(k)\phi(\mathbf{q})^T \phi(\mathbf{k}),其中 ϕ\phi 的第 jj 个分量为 1mexp⁡(−∥x∥2/2)exp⁡(ωjTx)\frac{1}{\sqrt{m}}\exp(-\|\mathbf{x}\|^2/2)\exp(\boldsymbol{\omega}_j^T\mathbf{x})。

正交随机特征:Choromanski et al. 进一步发现,如果 ω1,…,ωm\boldsymbol{\omega}_1, \ldots, \boldsymbol{\omega}_m 不是独立高斯采样,而是从正交矩阵的行中采样(正交随机特征,Orthogonal Random Features),近似的方差更低。这就是 FAVOR+ 中”+“的含义。

Performer 的完整计算流程

加入缩放因子 1/dk1/\sqrt{d_k} 和行归一化后,Performer 的 attention 计算为:

Attn~i=ϕ(qi/dk4)T∑j=1nϕ(kj/dk4)vjTϕ(qi/dk4)T∑j=1nϕ(kj/dk4)\widetilde{\text{Attn}}_i = \frac{\phi(\mathbf{q}_i / \sqrt[4]{d_k})^T \sum_{j=1}^{n}\phi(\mathbf{k}_j / \sqrt[4]{d_k})\mathbf{v}_j^T}{\phi(\mathbf{q}_i / \sqrt[4]{d_k})^T \sum_{j=1}^{n}\phi(\mathbf{k}_j / \sqrt[4]{d_k})}

分子和分母中的求和 ∑jϕ(kj)vjT\sum_j \phi(\mathbf{k}_j) \mathbf{v}_j^T 和 ∑jϕ(kj)\sum_j \phi(\mathbf{k}_j) 可以预先计算(对所有 query 共享),每个 query 只需一次 O(mdk)O(md_k) 的向量乘法。

标准 AttentionPerformer
核心计算softmax(QKT/dk)⋅V\text{softmax}(QK^T / \sqrt{d_k}) \cdot Vϕ(Q)⋅[ϕ(K)TV]\phi(Q) \cdot [\phi(K)^T V]
复杂度O(n2dk)O(n^2 d_k)O(nmdk)O(nmd_k)
存储O(n2)O(n^2)(注意力矩阵)O(nm+mdk)O(nm + md_k)(特征 + 累积矩阵)
近似质量精确依赖 mm,mm 越大越精确

当 m=O(dklog⁡dk)m = O(d_k \log d_k)(即 mm 与 nn 无关)时,Performer 将 attention 的复杂度从 O(n2)O(n^2) 降到 O(n)O(n)。

数值例子:softmax kernel 的随机特征近似

让我们用一个简单例子体会随机特征近似的精度。

取 dk=2d_k = 2,q=(1,0)T\mathbf{q} = (1, 0)^T,k=(0.5,0.5)T\mathbf{k} = (0.5, 0.5)^T。

精确值:

kSM(q,k)=exp⁡(qTk)=exp⁡(1×0.5+0×0.5)=exp⁡(0.5)≈1.6487k_{\text{SM}}(\mathbf{q}, \mathbf{k}) = \exp(\mathbf{q}^T\mathbf{k}) = \exp(1 \times 0.5 + 0 \times 0.5) = \exp(0.5) \approx 1.6487

随机特征近似(m=2m = 2,使用 ω1=(1,0)T,ω2=(0,1)T\boldsymbol{\omega}_1 = (1, 0)^T, \boldsymbol{\omega}_2 = (0, 1)^T 作为正交特征):

ϕ(q)=e−∥q∥2/22(eω1Tqeω2Tq)=e−0.52(e1e0)=12(e0.5e−0.5)≈12(1.64870.6065)\phi(\mathbf{q}) = \frac{e^{-\|\mathbf{q}\|^2/2}}{\sqrt{2}} \begin{pmatrix} e^{\boldsymbol{\omega}_1^T\mathbf{q}} \\ e^{\boldsymbol{\omega}_2^T\mathbf{q}} \end{pmatrix} = \frac{e^{-0.5}}{\sqrt{2}} \begin{pmatrix} e^1 \\ e^0 \end{pmatrix} = \frac{1}{\sqrt{2}} \begin{pmatrix} e^{0.5} \\ e^{-0.5} \end{pmatrix} \approx \frac{1}{\sqrt{2}}\begin{pmatrix} 1.6487 \\ 0.6065 \end{pmatrix}

ϕ(k)=e−∥k∥2/22(eω1Tkeω2Tk)=e−0.252(e0.5e0.5)=e0.252(11)≈1.28402(11)\phi(\mathbf{k}) = \frac{e^{-\|\mathbf{k}\|^2/2}}{\sqrt{2}} \begin{pmatrix} e^{\boldsymbol{\omega}_1^T\mathbf{k}} \\ e^{\boldsymbol{\omega}_2^T\mathbf{k}} \end{pmatrix} = \frac{e^{-0.25}}{\sqrt{2}} \begin{pmatrix} e^{0.5} \\ e^{0.5} \end{pmatrix} = \frac{e^{0.25}}{\sqrt{2}} \begin{pmatrix} 1 \\ 1 \end{pmatrix} \approx \frac{1.2840}{\sqrt{2}} \begin{pmatrix} 1 \\ 1 \end{pmatrix}

ϕ(q)Tϕ(k)=12(1.6487×1.2840+0.6065×1.2840)=1.28402(1.6487+0.6065)=1.2840×2.25522≈1.4478\phi(\mathbf{q})^T\phi(\mathbf{k}) = \frac{1}{2} \left(1.6487 \times 1.2840 + 0.6065 \times 1.2840\right) = \frac{1.2840}{2}(1.6487 + 0.6065) = \frac{1.2840 \times 2.2552}{2} \approx 1.4478

近似值 1.44781.4478 vs 精确值 1.64871.6487,误差约 12%。增加 mm(更多随机特征)会显著减小误差——Choromanski et al. 的实验表明,m=2dkm = 2d_k 到 4dk4d_k 通常足以在实际任务中保持模型质量。

Linformer vs Performer:两种路线的对比

两种方法从不同角度利用了注意力矩阵的低秩性:

LinformerPerformer
切入点直接压缩 K,VK, V 的序列维度近似 softmax kernel 的特征映射
数学工具低秩投影(Johnson-Lindenstrauss 引理)随机特征(Random Fourier Features 的推广)
近似对象KK 和 VV(投影后再做 softmax)softmax 核函数(先近似再利用结合律)
因果掩码需要额外处理自然支持(累积和可增量更新)
投影矩阵E,F∈Rk×nE, F \in \mathbb{R}^{k \times n}(与 nn 相关)ω1,…,ωm\boldsymbol{\omega}_1, \ldots, \boldsymbol{\omega}_m(与 nn 无关)
理论保证基于 JL 引理的逐点近似界基于高斯恒等式的无偏估计

关键区别:Linformer 的投影矩阵维度为 k×nk \times n,当 nn 变化时需要重新定义(或插值),因此对可变长度的适应性较弱。Performer 的随机特征只依赖 dkd_k,不依赖 nn,对任意长度序列都适用。

在实践中,两种方法都比标准 attention 快(尤其是长序列),但近似质量不如精确 attention——目前工业界更多使用 FlashAttention(不改变数学,只优化计算硬件利用率)来加速标准 attention,而 Linformer/Performer 开创的低秩思路则深刻影响了后续研究方向。

从低秩到更高效的架构

注意力矩阵的低秩性不仅催生了 Linformer 和 Performer,还启发了一个更深远的问题:

如果注意力矩阵本质上是低秩的,那么花 O(n2)O(n^2) 计算一个低秩矩阵是否太浪费了?能否设计一种本质上就是线性复杂度的序列建模机制?

这个问题的一个回答就是状态空间模型(State Space Model, SSM)及其变体 Mamba(Gu & Dao, 2023)。SSM 不构造 n×nn \times n 的注意力矩阵,而是用一个状态向量在时间步之间传递信息——每一步的更新是一次矩阵-向量乘法 O(Nd)O(Nd)(NN 是状态维度),总复杂度天然为 O(nNd)O(nNd),线性于序列长度。

下一篇我们将进入 Part 3 的第三个方向:SSM/Mamba——从矩阵指数 eAΔe^{A\Delta}(Art. 17 Kalman 中的给定算子)到可学习的对角状态矩阵,看特征分解和对角化(Art. 2 特征分解)如何在序列建模中达到极致应用。

总结

本文围绕 attention 矩阵的低秩结构展开,建立了从观察到利用的完整链条:

  • 低秩的数学原因:rank(QKT)≤dk\text{rank}(QK^T) \leq d_k,当 dk≪nd_k \ll n 时注意力矩阵天然低秩;softmax 的指数函数进一步”锐化”注意力分布,使有效秩更低
  • Gram 矩阵视角:QKTQK^T 是 query-key 的交叉 Gram 矩阵,softmax 是一种正定核函数——注意力机制本质上在做 kernel 回归,连接回 Art. 20 Kernel
  • Linformer(Wang et al., 2020):用投影矩阵 E,F∈Rk×nE, F \in \mathbb{R}^{k \times n} 将 K,VK, V 的序列维度从 nn 压缩到 kk,复杂度从 O(n2dk)O(n^2 d_k) 降到 O(nkdk)O(nkd_k)
  • Performer(Choromanski et al., 2021):用随机特征 ϕ(q)Tϕ(k)≈exp⁡(qTk/dk)\phi(\mathbf{q})^T\phi(\mathbf{k}) \approx \exp(\mathbf{q}^T\mathbf{k}/\sqrt{d_k}) 近似 softmax kernel,利用结合律将复杂度从 O(n2)O(n^2) 降到 O(n)O(n)——这是 kernel trick 的逆向应用
  • 弧线位置:在 Part 3 “汇——学习算子”的三个方向中,本文是第二个——利用注意力矩阵的近似低秩性加速推理。LoRA(Art. 24 LoRA)利用的是权重增量的低秩性(训练阶段),而这里利用的是计算中间结果的低秩性(推理阶段)

下一篇:SSM/Mamba——从给定算子到学习算子