[论文解读] S$^2$-MLP: Spatial-Shift MLP Architecture for Vision

Tan Yu, Li Xu+3 more|arXiv (Cornell University)|Jun 14, 2021

Advanced Neural Network Applications参考文献 39被引用 29

一句话总结

S2-MLP 是一种纯 MLP 架构，使用无参数的空间位移操作进行跨补丁通信，在 FLOPs 与参数量均少于 ViT 与 MLP-Mixer 的情况下实现对 ImageNet-1K 的具有竞争力的性能。

ABSTRACT

Recently, visual Transformer (ViT) and its following works abandon the convolution and exploit the self-attention operation, attaining a comparable or even higher accuracy than CNNs. More recently, MLP-Mixer abandons both the convolution and the self-attention operation, proposing an architecture containing only MLP layers. To achieve cross-patch communications, it devises an additional token-mixing MLP besides the channel-mixing MLP. It achieves promising results when training on an extremely large-scale dataset. But it cannot achieve as outstanding performance as its CNN and ViT counterparts when training on medium-scale datasets such as ImageNet1K and ImageNet21K. The performance drop of MLP-Mixer motivates us to rethink the token-mixing MLP. We discover that the token-mixing MLP is a variant of the depthwise convolution with a global reception field and spatial-specific configuration. But the global reception field and the spatial-specific property make token-mixing MLP prone to over-fitting. In this paper, we propose a novel pure MLP architecture, spatial-shift MLP (S$^2$-MLP). Different from MLP-Mixer, our S$^2$-MLP only contains channel-mixing MLP. We utilize a spatial-shift operation for communications between patches. It has a local reception field and is spatial-agnostic. It is parameter-free and efficient for computation. The proposed S$^2$-MLP attains higher recognition accuracy than MLP-Mixer when training on ImageNet-1K dataset. Meanwhile, S$^2$-MLP accomplishes as excellent performance as ViT on ImageNet-1K dataset with considerably simpler architecture and fewer FLOPs and parameters.

研究动机与目标

• 在中等规模数据上说明在视觉骨干中替代卷积与自注意力的必要性。
• 提出一种带有无参数空间位移块用于补丁通信的纯 MLP 架构（S2-MLP）。
• 显示基于空间位移的通信在参数和 FLOPs 更少的情况下仍能获得与 ViT 和 MLP-Mixer 相当的精度。
• 在 ImageNet-1K 上评估 S2-MLP，并进行消融研究以理解深度、宽度、位移方向和输入尺度的影响。

提出的方法

• 引入对图像补丁的逐补丁全连接嵌入。
• 使用 N 层堆叠的 S2-MLP 块，每个块包含四个全连接层、两个 GELU 激活和两个层归一化。
• 用空间位移模块替代 token 混合，将通道分组并向一个方向移动每组，以实现局部跨补丁通信。
• 将空间位移定义为一个无参数的操作，相当于在相邻补丁之间的固定深度卷积式位移。
• 给出 PFL、S2-MLP 块和最终分类器层的复杂度分析。

实验结果

研究问题

• RQ1 一个纯 MLP 架构在没有 token 混合的情况下，是否能够在中等规模数据上达到有竞争力的 ImageNet-1K 精度？
• RQ2 一个无参数的空间位移机制是否提供足够的跨补丁通信以匹配或超过 MLP-Mixer，并接近 ViT 的性能？
• RQ3 深度（N）、隐藏大小（c）、扩张比率（r）、位移方向和输入尺度如何影响准确性与效率？
• RQ4 S2-MLP 架构对输入尺度是否不变，这与带 token 混合的 MLP 模型相比有何不同？

主要发现

模型	分辨率	Top-1	Top-5	参数量 (M)	FLOPs (B)
S2-MLP-wide	224×224	80.0	94.8	71	14.0
S2-MLP-deep	224×224	80.7	95.4	51	10.5
Mixer-B/16	224×224	76.4	-	59	11.6
FF	224×224	74.9	-	59	11.6
ResMLP-36	224×224	79.7	-	45	8.9
ViT-B/16	384×384	77.9	-	55.5	-

• S2-MLP-wide 在 ImageNet-1K 上达到 Top-1 80.0% 和 Top-5 94.8%，参数量 71M、FLOPs 14B，在规模相近的情况下优于 MLP-Mixer（Top-1 76.4%）。
• S2-MLP-deep 在 ImageNet-1K 上达到 Top-1 80.7% 和 Top-5 95.4%，参数量 51M、FLOPs 10.5B，在相似条件下超过 ResMLP-36。
• 在 ImageNet-1K 上，S2-MLP 性能与 ViT 相竞争，同时在所报道的配置中保持更简单的架构和更低的 FLOPs/参数量。
• 消融结果表明，将深度从 1 增加到 12 块可将准确率提升（ImageNet100 上 Top-1 从 56.7% 提升至 87.1%），在较小数据集上超过 12–16 块后则趋于饱和或略有下降，原因是过拟合。
• 将隐藏大小 c 增大在一定范围内提升准确性（c=768 时达到 87.1% Top-1），但更高的 c 会增加参数量与 FLOPs。
• 在四个方向上进行位移（默认设置）提供强跨补丁通信；增加方向在某点前提下提升性能，而移除位移则会极大降低准确性。