从核函数角度深度理解线性注意力Linear Attention

原创已于 2025-12-19 19:52:23 修改 · 305 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能 #注意力机制

于 2025-12-19 19:26:58 首次发布

1、映射函数

当我们比较x与y的相似性时，通常将其从低维度映射到高维特征空间，以便在其中用简单的点积来衡量原始空间中复杂的相似性

映射函数为 $\phi\left ( \cdot \right )$ ：

$\phi : R^d \rightarrow R ^D \left ( D\geq d \right )$

例如：对于数据 $X = \left ( x_1,x_2 \right ) \in R^2$

$\phi\left ( X \right ) = \phi\left ( x_1,x_2 \right ) = \left (x_1^2,x_2^2,\sqrt{2}x_1x_2,\sqrt{2}x_1,\sqrt{2}x_2,x_1x_2, 1 \right ) \in R^6$

就将2维数据变成了6维数据了

2、核函数

常规而言，计算 $x$ 和 $y$ 的相似度，应该先分别计算 $\phi\left ( x \right )$ 和 $\phi\left ( y \right )$ ，再计算 $\phi\left ( x \right ) \cdot \phi\left ( y \right )$ 就是 $x$ 和 $y$ 的高维相似度了。但是这样计算太麻烦了，我们希望直接把值带入进去就可以直接知道 $\phi\left ( x \right ) \cdot \phi\left ( y \right )$ 的值。

换句话说，我们不需要显式知道 $\phi\left ( x \right )$ 和 $\phi\left ( y \right )$ 到底等于什么，只需要一个能计算最终相似度的公式就可以了。

于是核函数为 $\kappa \left (x \cdot y\right )$ ：

$\kappa \left (x \cdot y\right ) = <\phi\left ( x \right ) \cdot \phi\left ( y \right ) >$

例如：如果有映射函数 $\phi\left ( \cdot \right )$ ：

$\phi\left ( x \right ) = \left ( x^2,\sqrt{2}x,1 \right )$

将1维数据升为3维，我们可以计算对于 $x$ 和 $y$ 在映射空间中的相似度：

$\phi\left ( x \right ) \cdot \phi\left (y \right )^T = \begin{pmatrix} x^2,\sqrt{2}x, 1 \end{pmatrix}\cdot \begin{pmatrix} y^2\\ \sqrt{2}y \\ 1 \end{pmatrix} = x^2y^2+2xy+1 =(xy+1)^2$

所以可得该情况下的核函数：

$\kappa \left (x \cdot y\right ) =(xy+1)^2$

以上是显式可推导而出的核函数公式，也就是直接带入 $x$ 和 $y$ 就可以。妙处在于，我们无需知道显式的映射 $\phi\left ( x \right )$ 和 $\phi\left ( y \right )$ 到底等于什么，也无需在高维空间中进行计算，就能得到高维空间的内积结果。

3、注意力公式与核函数

对于注意力公式 $Attn = softmax(\frac {Q\cdot K^T}{\sqrt{d_{model}}})V$ 而言，我们先省略缩放参数 $d_{model}$ ，取第i个查询 $Q_i$ ，则原公式为：

$Attn_i = softmax(Q_i\cdot K^T)V$

进一步我们对所有的 $K,V$ 做softmax的展开为：

$Attn_i = \frac{\sum _{j=1}^{N} {exp(Q_i\cdot K_j^T)\cdot V_j}}{\sum _{j=1}^{N} {exp(Q_i\cdot K_j^T)}}$

当我们审视标准注意力公式时，会发现其核心是计算查询 $Q_i$ 与键 $K_j$ 的相似度，我们可以将 $exp(Q\cdot K^T)$ 视为一种特殊的核函数。

然而，这个核函数对应的特征映射是复杂且隐式的，虽然可以直接带入 $Q_i$ 与 $K_j$ 求出来，但是直接 $Q\cdot K^T$ 是进行矩阵乘法，复杂度 $o(n^2)$ 随着 $n$ 增大暴增，所以我们回归核函数定义的本质，将核函数还原为 $\kappa \left (q \cdot k\right ) = <\phi\left ( q \right ) \cdot \phi\left ( k \right ) >$ ，即先分别计算 $\phi\left ( q \right )$ 和 $\phi\left ( k \right )$ ，再计算 $\phi\left ( q \right ) \cdot \phi\left ( k \right )$ ，改为多个 $o(n)$ 的串行计算得到 $exp(Q\cdot K^T)$ ，从而避免复杂度 $o(n^2)$ 的矩阵乘法

线性注意力的核心思想是：能否设计一种新的、更简单的核函数 $\phi(Q\cdot K^T)$ ，并找到其对应的显式特征映射 $\phi(\cdot )$ ，使得相似度计算可以写成 $\phi(Q\cdot K^T)\approx \phi(Q)\cdot \phi(K)^T$

4、进一步公式推导

首先在这里， $exp(Q\cdot K^T)$ 不能直接进行分解，因为是指数函数

举例为：

$\phi(xy)=exp(xy)=(exp(x))^y=(\phi(x))^y \neq \phi(x)\cdot \phi(y)^T$

所以我们只能去另外找近似的新的核函数（已经不是exp了），使得：

$sim(Q\cdot K^T)\approx \phi(Q)\cdot \phi(K)^T$

由此，原式子可以化为：

$LinearAttn_i = \frac{\sum _{j=1}^{N} {\phi(Q_i)\phi( K_j)^T\cdot V_j}} {\sum _{j=1}^{N} {\phi(Q_i)\phi( K_j)^T}}$

同时，因为 $Q_i$ 只与i有关，求和符号只与j有关，所以可以将 $\phi(Q_i)$ 提出求和符号，放在前面：

$LinearAttn_i = \frac{\phi(Q_i)\sum _{j=1}^{N} {\phi( K_j)^T\cdot V_j}} {\phi(Q_i)\sum _{j=1}^{N} {\phi( K_j)^T}}$

我们可以将其进一步写作向量的形式，结合矩阵乘法的交换律优先计算 $\phi(K)^TV$ ：

$(\phi(Q)\phi(K)^T)V=\phi(Q)(\phi(K)^TV)$

5、总结

可以看到，线性注意力的核心就是为注意力机制选择了一个易于分解的‘核函数’ $\phi(Q\cdot K^T)\approx \phi(Q)\cdot \phi(K)^T$ 。

不同的线性注意力变体（如Linear Transformer, Performer, FLASH等）的区别主要就在于特征映射函数 $\phi(\cdot )$ 的设计上（例如，Performer使用exp和随机投影来近似softmax核，而最早的Linear Transformer使用elu(x)+1等简单函数）