深度学习——Pay Attention to MLPs

前言

前不久入职某大厂了，有点怀念无忧无虑的学生时代。入职后很快收到了第一个优化任务，算法岗的高不确定性确实会让人有一丝焦虑。目前体感来看，现有的深度学习模型性能非常依赖于数据质量，在数据质量足够的前提下，才有模型上的一系列操作。

本文将总结类ViT网络结构Gmlp，论文题为Pay Attention to MLPs

本文为个人总结，如有错误，欢迎指出。本文默认读者具有ViT相关知识

Gmlp

输入与输出

Gmlp是类ViT的结构，其输入仍为若干图像块（即将一张图像切割成若干图像块），输出为若干个向量（token）堆叠组成的矩阵，例如token的维度为$L$，个数为$N$，则输出为$N\ast L$的矩阵。通过池化等操作转换为最终的特征向量。

结构

对比ViT，Gmlp取消了Position Embedding以及Self Attention，由若干个基本构成单元堆叠而成，基本构成单元（unit）的结构如下图所示

设输入矩阵（即图中的Input Embeddedings）为$n\ast d$的矩阵$X$，则Gmlp的unit结构可以简化为（省略了Norm等操作）：
$$\begin{array}{rl}Z& =\delta (XU)\\ \hat{Z}& =s(Z)\\ Y& =\hat{Z}V+X\end{array}$$

$U、V$为可学习的矩阵（就是FC层，维度自己定），$\delta$为激活函数，$s(z)$为图中的Spatial Gating Unit，其结构可以表示为

$$\begin{array}{rl}{f}_{W,b}(Z)& =WZ+b\\ s(Z)& =Z\odot f_{W,b}(Z)\end{array}$$
设$Z$为$n\ast d$的矩阵，则$W$为$n\ast n$的矩阵(不会改变输入矩阵的维度)，$b$为$n$维向量（$WZ+b$表示$WZ$的第一行元素与b的第一维元素相加），为了保证训练的稳定性，$W$初始化值接近于0（貌似用[-1,1]的均匀分布初始化），$b$的初始值为1，此时$f_{W,b}(Z)\approx 1$，此时Spatial Gating Unit相当于Identity Mapping。

更进一步的作者发现将$Z$沿着channel维度切割成$Z_1、Z_2$（$Z_1、Z_2$维度分别为$n\ast d_1$，$n\ast d_2$，$d_1+d_2=n$）两个部分更为有效，此时$s(Z)$操作变为
$$s(Z)=Z_1\odot f_{W,b}(Z_2)$$

Spatial Gating Unit的输出在经过矩阵$V$调整维度后与unit的输入进行相加。

个人理解

这里总结几点我认为比较重要的点，首先是脱离Gmlp的大方向上，在看完ViT的结构后，个人一直感觉Self Attention不是其性能的主要来源，gMLP、MLP-Mixer，以及接下来要总结的MetaFormer都验证了这一点。更进一步的，这几个工作都有一个共性，即类ViT结构的模型都会先用FC层混合单个token中的信息，接着用一个特定的操作来混合多个token的信息（MLP-Mixer通过token mixer来混和多个token的信息，MetaFormer通过池化操作，gMLP通过Spatial Gating Unit），这两个操作可看成对CNN中卷积操作的拆解，在CNN中，多个特征图输入到卷积核后，卷积核会糅合多个特征图中的空间（Spatial）和通道（channel）信息。而在类ViT结构中，若将token看成特征图，则其首先利用FC层糅合单个特征图（token）的空间信息，接着糅合多个特征图（token）的信息（通道信息）。此外，类ViT结构中的矩阵操作可以等价于卷积操作（只不过卷积核的分辨率大小等于特征图大小）。综合上述分析，类ViT结构更像是CNN模型的变种，我们似乎发现了一种更为有效的CNN结构。

回到Gmlp上，Spatial Gating Unit形式上很想Attention，但作者引入其的本意并不是挑选哪类特征是重要的，而是混合多个token的信息，不过这种混合的方式比较诡异，应该存在更具解释性的方式。