Synthesizer：我们还不够了解自注意力

本文关注的是自注意力机制。直观上来看，自注意力机制算是解释性比较强的模型之一了，它通过自己于自己的 Attention 来自动捕捉 token 与 token 之间的关联，事实上在《Attention is All You Need》那篇论文中，就给出了如下的看上去挺合理的可视化效果：

但注意力机制真的是这样生效的吗？这种 "token" 对 "token" 的注意力是必须的吗？前不久 Google 的新论文《Synthesizer: Rethinking Self-Attention in Transformer Models》对自注意力机制做了一些 "异想天开" 的探索，里边的结果也许会颠覆我们对自注意力的认知

Self-Attention

Transformer 的核心是 query-key-value 的点积自注意力，点积自注意力的基本作用是学习自对齐（self-alignment），即确定单个 token 相对于序列中所有其他 token 的相对重要性。实际上 query、key 和 value 隐含着自注意力模拟一个基于内容的检索过程，而这个过程的核心是 pairwise 之间的交互。Self-Attention 的基础是 Scaled-Dot Attention，定义如下：

$$ \text{Attenion}(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = \text{Softmax}\left(\frac{\boldsymbol{Q}\boldsymbol{K}^T}{\sqrt{d_k}}\right)\cdot \boldsymbol{V}\tag{1} $$

其中 $\boldsymbol {Q}\in \mathbb {R}^{n\times d_k},\boldsymbol {K}\in \mathbb {R}^{m\times d_k},\boldsymbol {V}\in \mathbb {R}^{m\times d_v}$，softmax 则是在 $m$ 的那一维进行归一化。而自注意力，则是对于同一个 $\boldsymbol {X}\in \mathbb {R}^{n\times d}$，通过不同的投影矩阵 $\boldsymbol {W}_q,\boldsymbol {W}_k,\boldsymbol {W}_v\in \mathbb {R}^{d\times d'}$ 得到 $\boldsymbol {Q}=\boldsymbol {X}\boldsymbol {W}_q,\boldsymbol {K}=\boldsymbol {X}\boldsymbol {W}_k,\boldsymbol {V}=\boldsymbol {X}\boldsymbol {W}_v$，然后再做 Attention，即

$$ \begin{aligned} \text{Self-Attention}(\boldsymbol{X})=&\, \text{Attention}(\boldsymbol{X}\boldsymbol{W}_q,\boldsymbol{X}\boldsymbol{W}_k,\boldsymbol{X}\boldsymbol{W}_v)\\ =&\, \text{Softmax}\left(\frac{\boldsymbol{X}\boldsymbol{W}_q\boldsymbol{W}_k^T\boldsymbol{X}^T}{\sqrt{d_k}}\right) \end{aligned}\tag{2} $$

至于 Multi-Head Attention，则不过是 Attention 运算在不同的参数下重复多次然后将多个输出拼接起来，属于比较朴素的增强。而关于它的进一步推广，则可以参考《突破瓶颈，打造更强大的 Transformer》

天马行空

本质上来看，自注意力就是通过一个 $n\times n$ 的矩阵 $\boldsymbol {A}$ 和 $d\times d'$ 的矩阵 $\boldsymbol {W}_v$，将原本是 $n\times d$ 的矩阵 $\boldsymbol {X}$ 变成 $n\times d'$ 的矩阵 $\boldsymbol {A}\boldsymbol {X}\boldsymbol {W}_v$。其中 $\boldsymbol {A}$ 是动态生成的，即

$$ \begin{equation}\boldsymbol{A}=\text{Softmax}\left(\boldsymbol{B}\right),\quad\boldsymbol{B}=\frac{\boldsymbol{X}\boldsymbol{W}_q \boldsymbol{W}_k^{\top}\boldsymbol{X}^{\top}}{\sqrt{d_k}}\tag{3}\end{equation} $$

对于矩阵 $B$，本质上来说它就是 $\boldsymbol {X}$ 里边两两向量的内积组合，所以我们称它为 "token" 对 "token" 的 Attention

那么就到了前面提出的问题："token" 对 "token" 是必须的吗？能不能通过其他方式来生成这个矩阵 $\boldsymbol {B}$？Google 的这篇论文提出了几种 "天马行空" 的形式并做了实验，这些形式统称为 Synthesizer

Dense 形式

第一种形式在原论文中称为 Dense：$\boldsymbol {B}$ 需要是 $n\times n$ 大小的，而 $\boldsymbol {X}$ 是 $n\times d$ 的，所以只需要一个 $d\times n$ 的变换矩阵 $\boldsymbol {W}_a$ 就可以将它变成 $n\times n$ 了，即

$$ \boldsymbol{B}=\boldsymbol{X}\boldsymbol{W}_a\tag{4} $$

这其实就相当于把 $\boldsymbol {K}$ 固定为常数矩阵 $\boldsymbol {W}_a^T$ 了。当然，原论文还做得更复杂一些，用到了两层 Dense 层

$$ \boldsymbol{B}=\text{relu}(\boldsymbol{X}\boldsymbol{W}_1+\boldsymbol{b}_1)\boldsymbol{W}_2+\boldsymbol{b}_2\tag{5} $$

但思想上并没有什么变化

Random 形式

刚才说 Dense 形式相当于把 $\boldsymbol {K}$ 固定为常数矩阵，我们还能不能更 "异想天开" 一些：直接用一个随机初始化的矩阵 $\boldsymbol {R}\in \mathbb {R}^{n\times n}$ 当作 $\boldsymbol {B}$ 矩阵，即

$$ \boldsymbol{B}=\boldsymbol{R}\tag{6} $$

原论文中还真实验了这种形式，称之为 Random，顾名思义，就是 $\boldsymbol {B}$ 是随机初始化的，并且可以选择随训练更新或不更新。据原论文描述，固定形式的 Attention 首次出现在论文《Fixed Encoder Self-Attention Patterns in Transformer-Based Machine Translation》，不同点是那篇论文的 Attention 矩阵是由一个函数算出来的，而 Google 这篇论文则是完全随机初始化的。从形式上看，该方法不依赖 token 对之间的交互或者任何单个 token 的信息

低秩分解

上面两种新形式，往往会面临参数过多的问题，所以很自然地就想到通过低秩分解来降低参数量。对于 Dense 和 Random，原论文也提出并验证了对应的低秩分解形式，分别称为 Factorized Dense 和 Factorized Radom

Factorized Dense 通过 Dense 的方式，生成两个 $n\times a, n\times b$ 的矩阵 $\boldsymbol {B}_1,\boldsymbol {B}_2$，其中 $ab=n$；然后将 $\boldsymbol {B}_1$ 重复 $b$ 次、将 $\boldsymbol {B}_2$ 重复 $a$ 次，得到对应的 $n\times n$ 的矩阵 $\tilde {\boldsymbol {B}}_1,\tilde {\boldsymbol {B}}_2$，最后将它们逐位相乘，合成一个 $n\times n$ 的矩阵：

$$ \boldsymbol{B} = \tilde{\boldsymbol{B}}_1\otimes \tilde{\boldsymbol{B}}_2\tag{7} $$

至于 Factorized Random 就很好理解了，本来是一整个 $n\times n$ 的矩阵 $\boldsymbol {R}$，现在变成两个 $n\times k$ 的矩阵 $\boldsymbol {R}_1,\boldsymbol {R}_2$，然后

$$ \boldsymbol{B}=\boldsymbol{R}_1\boldsymbol{R}_2^T\tag{8} $$

混合模式

到目前为止，连同标准的自注意力，我们共有 5 种不同的生成矩阵 $\boldsymbol {B}$ 的方案，它们也可以混合起来，即

$$ \begin{equation}\boldsymbol{B}=\sum_{i=1}^N \alpha_i \boldsymbol{B}_i\end{equation} $$

其中 $\boldsymbol {B}_i$ 是不同形式的自注意力矩阵，而 $\sum\limits_{i=1}^N \alpha_i=1$ 是可学习参数

结果分析

前面介绍了统称为 Synthesizer 的几种新型自注意力形式，它们的共同特点是没有保持 "token" 对 "token" 的形式，尤其是 Random，完全抛弃了原有注意力的动态特点，变成了静态的矩阵

那么，这些新型自注意力的效果如何呢？它们又怎样冲击我们对自注意力机制的认识呢？

机器翻译

第一个评测任务是机器翻译，详细地比较了各种自注意力形式的效果

不知道读者怎么想，反正 Synthesizer 的这些结果是冲击了笔者对自注意力的认知的。表格显示，除了固定的 Random 外，所有的自注意力形式表现基本上都差不多，而且就算是固定的 Random 也有看得过去的效果，这表明我们以往对自注意力的认知和解释都太过片面了，并没有揭示自注意力生效的真正原因

摘要对话

接下来是摘要和对话生成任务上的效果

在自动摘要这个任务上，标准注意力效果比较好，但是对话生成这个任务上，结果则反过来：标准的自注意力是最差的，Dense（D）和 Random（R）是最好的，而当 Dense 和 Random 混合了标准的自注意力后（即 D+V 和 R+V），效果也变差了。这说明标准注意力并没有什么 "独占鳌头" 的优势，而几个 Synthesizer 看起来是标准注意力的 "退化"，但事实上它们互不从属，各有优势

预训练 + 微调

最后，对于我们这些普通读者来说，可能比较关心的是 "预训练 + 微调" 的效果怎么样，也就是说，将 BERT 之类模型的自注意力替换之后表现如何？原论文确实也做了这个实验，不过 Baseline 不是 BERT 而是 T5，结果如下：

在这个结果中，相比标准自注意力，Dense 和 Random 就显得逊色了，这表明 Dense 和 Random 也许会在单一任务上表现得比较好，但是迁移能力则比较弱。不过不能否认的是，像 Random 这样的自注意力，由于直接省去了 $\boldsymbol {Q}\boldsymbol {K}^T$ 这个矩阵运算，因此计算效率会有明显提升，所以如果能想办法解决这个迁移性问题，说不定 Transformer 模型家族会迎来大换血

References

• Google | 没有点积注意力就不香了吗？Transformer 中自注意力机制的反思
• Google 新作 Synthesizer：我们还不够了解自注意力