Transformer升级之路：5、作为无限维的线性Attention
原文链接 | 2021-08-06 | 苏剑林
在《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=0}^{\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机制。