Transformer升级之路：5、作为无限维的线性Attention

在《Performer：用随机投影将Attention的复杂度线性化》中我们了解到Google提出的Performer模型，它提出了一种随机投影方案，可以将标准Attention转化为线性Attention，并保持一定的近似。理论上来说，只要投影的维度足够大，那么可以足够近似标准Attention。换句话说，标准Attention可以视作一个无限维的线性Attention。

本文将介绍笔者构思的另外两种将标准Attention转换为无限维线性Attention的思路，不同于Performer的随机投影，笔者构思的这两种方案都是确定性的，并且能比较方便地感知近似程度。

简要介绍 #

关于标准Attention和线性Attention，这里就不多做介绍了，还不了解的读者可以参考笔者之前的文章《线性Attention的探索：Attention必须有个Softmax吗？》和《Transformer升级之路：3、从Performer到线性Attention》。简单来说，标准Attention的计算方式为
\begin{equation}a_{i,j}=\frac{e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}{\sum\limits_j e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}\end{equation}
而线性Attention的计算方式为
\begin{equation}a_{i,j}=\frac{\phi(\boldsymbol{q}_i)\cdot \varphi(\boldsymbol{k}_j)}{\sum\limits_j \phi(\boldsymbol{q}_i)\cdot \varphi(\boldsymbol{k}_j)}\end{equation}
所以说，要将标准Attention（近似地）变换为线性Attention，那么一般情况下就要找到变换$\phi,\varphi$，使得有近似
\begin{equation}\phi(\boldsymbol{q})\cdot \varphi(\boldsymbol{k})\approx e^{\boldsymbol{q}\cdot \boldsymbol{k}}\end{equation}
这时候$e^{\boldsymbol{q}\cdot \boldsymbol{k}}$也就是核方法中的“核函数”。

随机投影 #

Performer找到了第一个比较实用的随机投影变换方案，本质上来说，它基于以下积分：
\begin{equation}\begin{aligned}
e^{\boldsymbol{q}\cdot \boldsymbol{k}} =&\,\frac{1}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}-\boldsymbol{q}-\boldsymbol{k}\Vert^2 / 2 + \boldsymbol{q}\cdot \boldsymbol{k}}d\boldsymbol{\omega}\\
=&\,\frac{1}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}\Vert^2 / 2}\times e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}d\boldsymbol{\omega}
\\
\end{aligned}\end{equation}
得到
\begin{equation}\begin{aligned}
e^{\boldsymbol{q}\cdot \boldsymbol{k}}&=\mathbb{E}_{\boldsymbol{\omega}\sim \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d)}\left[e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\right]\\[6pt]
&\approx\underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \\
e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2}\\
\vdots\\
e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \end{pmatrix}}_{\phi(\boldsymbol{q})}
\cdot \underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \\
e^{\boldsymbol{\omega}_2\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\\
\vdots\\
e^{\boldsymbol{\omega}_m\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \end{pmatrix}}_{\varphi(\boldsymbol{k})}
\end{aligned}\end{equation}
其中$\boldsymbol{\omega}_1,\boldsymbol{\omega}_2,\cdots,\boldsymbol{\omega}_m\sim \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d)$。这样我们就通过随机投影的思想，将两个$d$维向量的内积指数，近似地转化为了两个$m$维向量的内积，并且$m\to\infty$时，两者理论上是相等的。

上述随机投影的方案还是比较巧妙的，不容易想到。下面介绍笔者构思的两种方案，相对来说更容易理解一些，尤其是对于熟悉核函数的读者来说，可能扫一眼就能理解了。

泰勒展开 #

笔者的第一个思路，是基于泰勒展开的：
\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} = \sum_{m=1}^{\infty} \frac{(\boldsymbol{q}\cdot \boldsymbol{k})^m}{m!}\end{equation}
截断到前$n+1$项，那么就得到关于$\boldsymbol{q}\cdot \boldsymbol{k}$的一个$n$次多项式：
\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} \approx 1 + \boldsymbol{q}\cdot \boldsymbol{k} + \frac{1}{2}(\boldsymbol{q}\cdot \boldsymbol{k})^2 + \cdots + \frac{1}{n!}(\boldsymbol{q}\cdot \boldsymbol{k})^n\end{equation}
这其实就是一个“多项式核函数”，注意到我们有：
\begin{equation}\begin{aligned}
(\boldsymbol{q}\cdot \boldsymbol{k})^m =&\, \left(\sum_i q_i k_i\right)^m = \left(\sum_{i_1} q_{i_1} k_{i_1}\right)\cdots\left(\sum_{i_m} q_{i_m} k_{i_m}\right) \\
=&\, \sum_{i_1,\cdots,i_m} (q_{i_1}\cdots q_{i_m}) (k_{i_1}\cdots k_{i_m})
\end{aligned}\end{equation}
如果我们将$q_{i_1}\cdots q_{i_m},k_{i_1}\cdots k_{i_m}$分别看成一个$d^m$维的大向量，那么$(\boldsymbol{q}\cdot \boldsymbol{k})^m$就是这两个大向量的内积。事实上，由若干个向量得到“大向量”的这步运算，我们称为向量的“外积”，也叫“张量积”，一般也记为$\otimes$。此时
\begin{equation}
\frac{1}{m!}(\boldsymbol{q}\cdot \boldsymbol{k})^m = \frac{1}{m!}\underbrace{(\boldsymbol{q}\otimes\cdots\otimes\boldsymbol{q})}_{m\text{个}\boldsymbol{q}}\cdot\underbrace{(\boldsymbol{k}\otimes\cdots\otimes\boldsymbol{k})}_{m\text{个}\boldsymbol{k}} = \left(\frac{\otimes^m\boldsymbol{q}}{\sqrt{m!}}\right)\cdot\left(\frac{\otimes^m\boldsymbol{k}}{\sqrt{m!}}\right)
\end{equation}
这里$\otimes^m\boldsymbol{q},\otimes^m\boldsymbol{k}$是$m$个$\boldsymbol{q},\boldsymbol{k}$连着外积（外积的$m$次幂）的简写。利用这个结果，我们有
\begin{equation}
e^{\boldsymbol{q}\cdot \boldsymbol{k}}\approx \sum_{m=0}^n \left(\frac{\otimes^m\boldsymbol{q}}{\sqrt{m!}}\right)\cdot\left(\frac{\otimes^m\boldsymbol{k}}{\sqrt{m!}}\right) =\underbrace{\begin{pmatrix} 1 \\
\boldsymbol{q}\\
\frac{\otimes^2\boldsymbol{q}}{\sqrt{2}} \\
\vdots\\
\frac{\otimes^n\boldsymbol{q}}{\sqrt{n!}}\end{pmatrix}}_{\phi(\boldsymbol{q})}
\cdot \underbrace{\begin{pmatrix} 1 \\
\boldsymbol{k}\\
\frac{\otimes^2\boldsymbol{k}}{\sqrt{2}} \\
\vdots\\
\frac{\otimes^n\boldsymbol{k}}{\sqrt{n!}}\end{pmatrix}}_{\varphi(\boldsymbol{k})}
\end{equation}
这就完成了标准Attention到线性Attention的转换。

指数定义 #

相比Performer的随机投影，上述基于泰勒展开的思路应该说更好理解。不过还有一种比泰勒展开更简单直接的思路，那就是利用自然指数的定义式：
\begin{equation}e^x = \lim_{n\to\infty} \left(1+\frac{x}{n}\right)^n\end{equation}
因此，选取适当的$n$，我们就有
\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} \approx \left(1+\frac{{\boldsymbol{q}\cdot \boldsymbol{k}}}{n}\right)^n = \left(\begin{pmatrix} 1 \\ \frac{\boldsymbol{q}}{\sqrt{n}}\end{pmatrix} \cdot \begin{pmatrix}1 \\ \frac{\boldsymbol{k}}{\sqrt{n}}\end{pmatrix}\right)^n \end{equation}
结合前一节的多项式核函数的转化结果，我们有
\begin{equation}e^{\boldsymbol{q}\cdot \boldsymbol{k}} \approx \underbrace{\left(\otimes^n\begin{pmatrix} 1 \\ \frac{\boldsymbol{q}}{\sqrt{n}}\end{pmatrix}\right)}_{\phi(\boldsymbol{q})} \cdot \underbrace{\left(\otimes^n\begin{pmatrix}1 \\ \frac{\boldsymbol{k}}{\sqrt{n}}\end{pmatrix}\right)}_{\varphi(\boldsymbol{k})}\end{equation}
这可能是将标准Attention转换为线性Attention的最简单直接的方案。

结果分析 #

要说实用价值，后两个确定性的方案远不如Performer的随机投影方案，因为随机投影的输出维度可以比较灵活地控制，而两个确定性的方案输出纬度则是$d^n$级别的，这个通常都远远大于序列长度本身了，所以用它们来做线性Attention效率基本比标准的Attention还差。

不过，从理论上来讲，后两种方案提供了更为简明便捷的思路，让我们将标准Attention跟无限维的线性Attention等价起来。这种等价性通常能帮助我们更好地理解Attention机制，其中最直接的便是关于Attention的秩的理解。

做过线性Attention研究的读者应该知道，如果用线性Attention做双向注意力任务（比如MLM），那么效果下降会很明显，这是因为线性Attention的$\phi(\boldsymbol{Q}),\varphi(\boldsymbol{K})\in\mathbb{R}^{n\times d}$（$d$是每个head的head_size），一般有$n \gg d$，所以$\phi(\boldsymbol{Q})\varphi(\boldsymbol{K})^{\top}$得到的$n\times n$的Attention矩阵的秩顶多为$d$。这就是线性Attention的低秩问题，低秩限制了线性Attention的表达能力。

相反，前述介绍的三种变换，都告诉我们标准Attention可以视为无限维的线性Attention，所以标准Attention的秩理论上就不受限于$d$，因此同样参数量的标准Attention表现往往表现得比线性Attention好。在《Transformer升级之路：3、从Performer到线性Attention》中我们也说过，如果要将标准Attention切换为线型Attention，那么$d$也要相应地进行放大，才能要效果保持一定程度上的近似。

文章小结 #

本文介绍了三种将标准Attention视为无限维线性Attention的理解，这些不同的视角能让我们将标准Attention与线性Attention联系起来，从多个角度更全面地理解Attention机制。

转载到请包括本文地址：https://kexue.fm/archives/8601

更详细的转载事宜请参考：《科学空间FAQ》