【原】Self-Attention真的是必要的吗？微软&中科大提出Sparse MLP，降低计算量的同时提升性能！

我爱计算机视觉 2021-09-24

展开全文

▊ 写在前面

Transformer由于其强大的建模能力，目前在计算机视觉领域占据了重要的地位。在这项工作中，作者探究了Transformer的自注意（Self-Attention）模块是否是其实现图像识别SOTA性能的关键 。为此，作者基于现有的基于MLP的视觉模型，建立了一个无注意力网络sMLPNet。

具体来说，作者将以往工作中用于token混合的MLP模块替换为一个稀疏MLP(sMLP)模块。对于二维图像token，sMLP沿轴向（横向或者纵向）应用一维MLP，参数在行、列维度共享。通过稀疏连接 和权重共享 ，sMLP模块显著降低了模型参数的数量和计算复杂度，避免了MLP模型的内在问题（如过拟合、参数量大、计算量大）。

当仅在ImageNet-1K数据集上训练时，sMLPNet在只有24M参数下达到81.9%的Top-1精度，比相同模型大小约束下的大多数CNN和视觉Transformer要好得多。当扩展到66M参数时，sMLPNet达到了83.4%的Top-1精度，这与SOTA的 Swin Transformer相当。

▊ 1. 论文和代码地址

Sparse MLP for Image Recognition: Is Self-Attention Really Necessary?

论文地址：https://arxiv.org/abs/2109.05422代码地址：未开源sMLP Block复现代码：https://github.com/xmu-xiaoma666/External-Attention-pytorch#5-sMLP-Usage

▊ 2. Motivation

自AlexNet提出以来，卷积神经网络(CNN)一直是计算机视觉的主导范式。随着Vision Transformer的提出，这种情况发生了改变。ViT将一个图像被划分为不重叠的patch，并用线性层将这些patch转换为token，然后输入到Transformer中进行处理。

Transformer编码器由多头自注意网络(Multi-Head Self-Attention)和前馈网络(FFN)组成，来实现空间信息混合和通道信息混合。当在一个非常大的数据上进行预训练时，ViT在图像识别任务上表现得非常好。接着DeiT进一步证明了只在ImageNet-1K上训练时，通过适当的数据增强和正则化技术，无卷积的Vision Transformer也可以实现SOTA的图像识别精度。

目前，无卷积的Vision Transformer主要存在两个核心的思想：首先，全局依赖性建模很重要 。不仅如此，它甚至可以取代卷积操作的局部建模。第二，自注意很重要 。尽管ViT和DeiT表现良好，学术界并没有完全接受这两种观点。

一方面，研究人员挑战了用全局建模代替局部建模的必要性 。既然局部偏置在自然图像中是有效的，为什么要通过全局自注意模块来学习它，而不是直接将它注入到网络中呢 ？此外，全局自注意对于输入token的数量具有二次计算复杂度。因此，网络结构不有利于高分辨率输入，对金字塔结构并不友好。

基于这一点，Swin Transformer通过限制局部窗口内的自注意操作，将局部偏置注入回网络中。这种设置还控制了计算的复杂度，并允许使用金字塔结构和多阶段处理。Swin Transformer的优越性能表明了局部偏置和多阶段处理的有效性。

另一方面，研究人员也挑战了Self-Attention的必要性 。MLP-Mixer也建模了全局依赖关系，但它采用了一个MLP块，而不是一个自注意模块来实现。MLP-Mixer的整体架构与ViT相似。输入图像被分成多个patch，然后线性层将patch映射到token中。该编码器包含用于空间混合和通道混合的交替层。

唯一的主要区别是，空间混合模块是由一个MLP块实现的。MLP-Mixer继承了ViT的所有缺点，且由于参数数量过多，容易发生过拟合 。因此，MLP-Mixer和SOTA模型的性能还是存在一定差距，尤其是在不预训练的情况下。因此，作者在本文中探究了：在解决了所有的缺陷后，一个无注意力的网络是否有可能实现在图像识别上的SOTA性能？

因此，作者设计了一个无注意力的网络，称为sMLPNet，它只使用卷积和MLP作为构建块。sMLPNet采用了与ViT和MLP-Mixer类似的体系结构，且通道混合模块与他们完全相同。在每个token混合模块中，采用深度卷积来利用局部偏置，并使用改进的MLP来建模全局依赖关系 。

具体来说，作者提出了具有轴向（即横向和纵向）全局依赖建模特征的稀疏MLP(sMLP)模块（如上图右所示）。sMLP显著降低了计算的复杂度，并允许采用金字塔结构进行多阶段处理。因此，sMLPNet能够在更小的模型上实现与Swin Transformer相同的图像识别性能。

在本文中，作者研究了Transformer的关键组成部分（即Self-Attention）是否是图像理解的真正关键因素 。基于过去视觉模型的设计思想，作者采用了在设计时采用了局部偏置和金字塔结构。此外，作者也采用了全局依赖建模的思想，但使用稀疏MLP模块来实现。

基于以上思想，作者建立了一个名为sMLPNet的无注意力网络，实现了SOTA图像识别性能。本文表明，自注意力可能不是视觉模型设计的核心组成部分。相反，正确使用局部偏置、金字塔结构和对计算复杂度的控制是设计高性能视觉模型的关键 。

▊ 3. 方法

3.1. Design Guidelines

在这项工作中，作者保留了CNN使用的一些重要的设计理念，并添加了受Transformer启发的新组件。设计指南如下：

1.采用类似于ViT、MLP-Mixer和Swin Transformer的架构，以确保一个公平的比较。

2.显式地将局部偏置注入到网络中。

3.探索不使用自注意模块的全局依赖关系。

4.在金字塔结构中进行多阶段处理。

3.2. Overall Architecture

上图展示了本文网络的整体架构。与ViT、MLP-Mixer和Swin Transformer类似，空间分辨率为H×W的输入图像被分割为不重叠的patch。作者在网络中采用了4×4的patch大小，每个patch被reshape成一个48维的向量，然后由一个线性层映射到一个c维embedding，整张图像可以表示为的tensor。

整个网络由四个阶段组成。除第一阶段从线性embedding层开始外，其他阶段从patch合并层开始，将空间维数减少2×2，将通道维数增加2倍。patch合并层由一个线性层实现，它以每个2×2个相邻patch的concat特征作为输入，输出合并后的patch的特征。然后，将新的图像token输入到token混合模块和通道混合模块中。