《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》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Aug. 06, 2021). 《Transformer升级之路:5、作为无限维的线性Attention 》[Blog post]. Retrieved from https://kexue.fm/archives/8601