Day 11: 谷歌大脑新作：Pay Attention to MLPs

• 谷歌大脑的新作，还是和MLP相关，文章仍然是质疑 self-attention 在 Transformer 中的重要性，以及是否可以被取代。
• 提出一种基于MLP的，无注意力机制的，由通道映射、空间映射和门控制组成的结构，来替代 Transformer。
• 新结构的实验得到了很好的效果，在ImageNet上，gMLP达到了和DeiT接近的效果。和MLP-Mixer相比，它的参数量减少了66%，且准确率还有3%的提升。
• 实验中表明，模型的复杂度只与模型的能力有关，而与是否有注意力模块不太相关。当模型的能力提升的时候，我们观察到不论是预训练还是精调，它们的指标都提升得和 Transformer 一样块。这一点非常值得注意，因为这说明 gMLP 尽管没有注意力模块，但也能像 Transformer 一样缩放得非常好（意思是模型可大可小），这样，任何模型效果上的差距，都可以通过训练一个更大的模型、增加数据量和计算复杂度来填补。
• 当在 BERT 上做精调的时候，Transformer 可能比 gMLP 更实用，因为这个时候需要句子间的对齐（cross-sentence alignment）。不过这个情况可以通过用一个3倍 Transformer 大小的 gMLP 模型来解决。一个更实际点的方法是稍微加入一丢丢注意力——一个最大尺寸为128的单头注意力就足够让 gMLP 的效果在所有NLP的任务上超过 Transformer，甚至还有更好的参数效率。
• 最后，结果显示，自注意力模块并不是在将模型扩大（scaling up）的时候必须的东西，当有足够多的数据和计算量的时候，只有一些简单的空间交互机制的模型，比如gMLP，也能和 Transformer 一样强大，因此自注意力可以被直接移除，或者大规模地减少。

模型细节

• gMLP模型由一连串完全相同且搭在一起的 block 组成，让$X\in {\mathbb{R}}^{n\times d}$ 代表一串长度为 $n$ 的token，维度为 $d$，那每个 block 的定义如下：

$$\begin{array}{l}Z=\sigma (XU)\\ \tilde{Z}=s(Z)\\ Y=\tilde{Z}V\end{array}$$

• $U和V$ 分别是在通道方向上的线性映射，和 Transformer 中的前向层一样（在 BERT base 中，它们的尺寸分别是 $768\times 3072和3072\times 768$），为了表达方便，像短接（shortcut）、正则化和偏置都暂时忽略。
• 上式中的一个重要内容是$s(\cdot )$，它是用来捕捉空间交互信息的层。当$s$是恒等映射的时候，上式就是一个普通的前向层，且每个 token 单独处理，之间没有 token 间的信息交流。因此我们主要的事情就是要设计一个良好的 s 来捕捉 token 之间复杂的空间信息。这一整个 block 的灵感来自于反向的 bottleneck 结构，期中 s 充当空间深度方向上的卷积（spatial depthwise convolution）。
• 模型的输入和输出与BERT和ViT的完全相同。

Spatial Gating Unit

• 为了让 token 之间能有信息交流，$s(\cdot )$ 需要包含能在空间维度上压缩的操作。最简单的选择便是如下的线性映射：

$$f_{W,b}(Z)=WZ+b$$

$W\in {\mathbb{R}}^{n\times n}$ 的尺寸$n$是序列的长度，$b$ 是偏置，既可以是矩阵，也可以是个标量。比如$W$是$128\times 128$。

我们把空间交互模块（spatial interaction unit）定义为它的原始输入与空间转换后的输入的乘积：

$$s(Z)=Z\odot f_{W,b}(Z)$$

$\odot$ 指代元素间的点积。

我们发现将$W$初始化为接近0的值，将$b$初始化为1，非常重要，也就是上式在训练开始时，可以近似为一个恒等映射。这种初始化能确保每个gMLP模块在训练开始时，都像是一个常规的前向层，每个 token 都分别独立地处理，并渐渐地注入 token 之间的空间信息。

这个乘法门可以被看成是一种用空间信息来对每个 token 的表征进行调制（modulate）的机制，也就是说，每个元素$Z$的大小都能根据门函数$f_{W,b}(\cdot )$迅速地调节。

我们还发现，如果把$Z$在通道方向上拆分成两个独立的部分$(Z_1,Z_2)$会更有效，像Gated Linear Unit中那样（也很像注意力机制中的多头机制 multi-head attention）：

$$s(Z)=Z_1\odot f_{W,b}\left(Z_2\right)$$

此外，我们还会将输入正则化一下，根据经验，这样也能提高模型的稳定性。

总结

作者提出的SGU，和LSTM中的GLU很像，主要的一个区别在于，SGU的计算是在空间（cross-token）维度上，而不是通道（per-token）维度的。此外，SGU也能捕捉高次方的关系，比如二次方$z_i{z}_j$。SGU的复杂度是$n^{2}e/2$，与点积的注意力相当（$2n^{2}d$）。不过二者都与输入通道数成线性相关，与序列长度$n$成二次方相关。