- 谷歌大脑的新作,还是和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 ∈ R n × d X \in \mathbb{R}^{n \times d} X∈Rn×d 代表一串长度为 n n n 的token,维度为 d d d,那每个 block 的定义如下:
Z = σ ( X U ) Z ~ = s ( Z ) Y = Z ~ V \begin{array}{l}Z=\sigma(X U) \\\tilde{Z}=s(Z) \\Y=\tilde{Z} V\end{array} Z=σ(XU)Z~=s(Z)Y=Z~V
- U 和 V U 和 V U和V 分别是在通道方向上的线性映射,和 Transformer 中的前向层一样(在 BERT base 中,它们的尺寸分别是 768 × 3072 和 3072 × 768 768 \times 3072 和 3072 \times 768 768×3072和3072×768),为了表达方便,像短接(shortcut)、正则化和偏置都暂时忽略。
- 上式中的一个重要内容是 s ( ⋅ ) s(\cdot) s(⋅),它是用来捕捉空间交互信息的层。当 s s s是恒等映射的时候,上式就是一个普通的前向层,且每个 token 单独处理,之间没有 token 间的信息交流。因此我们主要的事情就是要设计一个良好的 s 来捕捉 token 之间复杂的空间信息。这一整个 block 的灵感来自于反向的 bottleneck 结构,期中 s 充当空间深度方向上的卷积(spatial depthwise convolution)。
- 模型的输入和输出与BERT和ViT的完全相同。
Spatial Gating Unit
- 为了让 token 之间能有信息交流, s ( ⋅ ) s(\cdot) s(⋅) 需要包含能在空间维度上压缩的操作。最简单的选择便是如下的线性映射:
f W , b ( Z ) = W Z + b f_{W, b}(Z)=W Z+b fW,b(Z)=WZ+b
W ∈ R n × n W \in \mathbb{R}^{n \times n} W∈Rn×n 的尺寸 n n n是序列的长度, b b b 是偏置,既可以是矩阵,也可以是个标量。比如 W W W是 128 × 128 128 \times 128 128×128。
我们把空间交互模块(spatial interaction unit)定义为它的原始输入与空间转换后的输入的乘积:
s ( Z ) = Z ⊙ f W , b ( Z ) s(Z)=Z \odot f_{W, b}(Z) s(Z)=Z⊙fW,b(Z)
⊙ \odot ⊙ 指代元素间的点积。
我们发现将 W W W初始化为接近0的值,将 b b b初始化为1,非常重要,也就是上式在训练开始时,可以近似为一个恒等映射。这种初始化能确保每个gMLP模块在训练开始时,都像是一个常规的前向层,每个 token 都分别独立地处理,并渐渐地注入 token 之间的空间信息。
这个乘法门可以被看成是一种用空间信息来对每个 token 的表征进行调制(modulate)的机制,也就是说,每个元素 Z Z Z的大小都能根据门函数 f W , b ( ⋅ ) f_{W, b}(\cdot) fW,b(⋅)迅速地调节。
我们还发现,如果把 Z Z Z在通道方向上拆分成两个独立的部分 ( Z 1 , Z 2 ) (Z_1, Z_2) (Z1,Z2)会更有效,像Gated Linear Unit中那样(也很像注意力机制中的多头机制 multi-head attention):
s ( Z ) = Z 1 ⊙ f W , b ( Z 2 ) s(Z)=Z_{1} \odot f_{W, b}\left(Z_{2}\right) s(Z)=Z1⊙fW,b(Z2)
此外,我们还会将输入正则化一下,根据经验,这样也能提高模型的稳定性。
总结
作者提出的SGU,和LSTM中的GLU很像,主要的一个区别在于,SGU的计算是在空间(cross-token)维度上,而不是通道(per-token)维度的。此外,SGU也能捕捉高次方的关系,比如二次方 z i z j z_{i} z_{j} zizj。SGU的复杂度是 n 2 e / 2 n^{2} e / 2 n2e/2,与点积的注意力相当( 2 n 2 d 2 n^{2} d 2n2d)。不过二者都与输入通道数成线性相关,与序列长度 n n n成二次方相关。