[论文解读] Sparse MLP for Image Recognition: Is Self-Attention Really Necessary?

Chuanxin Tang, Yucheng Zhao+4 more|arXiv (Cornell University)|Sep 12, 2021

Advanced Neural Network Applications被引用 13

一句话总结

该论文提出sMLPNet，一种无需注意力机制的视觉模型，通过使用轴向全局建模和参数共享的稀疏MLP（sMLP）模块替代自注意力机制。通过减少参数量和FLOPs并避免过拟合，sMLPNet在ImageNet-1K上实现了83.4%的top-1准确率，仅使用65.9M参数——模型规模仅为Swin Transformer的25%，FLOPs减少10%，性能相当，证明自注意力并非实现最先进图像识别的必要条件。

ABSTRACT

Transformers have sprung up in the field of computer vision. In this work, we explore whether the core self-attention module in Transformer is the key to achieving excellent performance in image recognition. To this end, we build an attention-free network called sMLPNet based on the existing MLP-based vision models. Specifically, we replace the MLP module in the token-mixing step with a novel sparse MLP (sMLP) module. For 2D image tokens, sMLP applies 1D MLP along the axial directions and the parameters are shared among rows or columns. By sparse connection and weight sharing, sMLP module significantly reduces the number of model parameters and computational complexity, avoiding the common over-fitting problem that plagues the performance of MLP-like models. When only trained on the ImageNet-1K dataset, the proposed sMLPNet achieves 81.9% top-1 accuracy with only 24M parameters, which is much better than most CNNs and vision Transformers under the same model size constraint. When scaling up to 66M parameters, sMLPNet achieves 83.4% top-1 accuracy, which is on par with the state-of-the-art Swin Transformer. The success of sMLPNet suggests that the self-attention mechanism is not necessarily a silver bullet in computer vision. The code and models are publicly available at https://github.com/microsoft/SPACH

研究动机与目标

• 挑战视觉Transformer中自注意力对实现最先进性能必不可少的假设。
• 探究MLP-based模型是否可通过解决其关键局限性——过拟合和高参数量——在不使用自注意力的情况下实现有竞争力的性能。
• 设计一种新型稀疏MLP（sMLP）模块，通过轴向计算和参数共享实现高效的全局特征建模。
• 构建一种无需注意力机制的网络（sMLPNet），其性能可与Swin Transformer等最先进视觉Transformer模型相匹配或超越。
• 在标准训练协议下（无需数据增强或额外预训练）验证sMLPNet在ImageNet-1K上的有效性。

提出的方法

• sMLP模块在2D图像标记的行和列轴上应用1D MLP，实现轴向全局依赖建模。
• 所有行（或列）共享参数，显著降低模型规模和计算复杂度。
• 通过限制每个标记仅与行和列邻居交互，引入稀疏性，避免全连接的二次方注意力。
• sMLP模块被集成到类似ViT的架构中，取代标记混合模块中的自注意力机制。
• 在标记混合模块中使用深度可分离卷积以保留局部性偏置，而通道混合部分保持与标准MLP相同。
• 通过调整宽度和深度实现模型缩放，实验在224×224分辨率的ImageNet-1K上进行。

实验结果

研究问题

• RQ1 基于MLP的视觉模型是否能在不使用自注意力的情况下实现最先进性能？
• RQ2 视觉Transformer中的自注意力机制是否相对于其他全局建模机制具有根本性优势？
• RQ3 MLP中的稀疏连接和参数共享是否能缓解过拟合并降低计算成本，同时保持性能？
• RQ4 通过结合局部性偏置和轴向全局建模，是否可能实现更小模型规模下的有竞争力准确率？
• RQ5 Swin Transformer的成功源于局部性还是自注意力机制？是否可能构建一个无需注意力机制的模型来匹配其性能？

主要发现

• sMLPNet在仅24.1M参数下于ImageNet-1K上达到81.9%的top-1准确率，优于相同参数量约束下的大多数CNN和视觉Transformer。
• 当参数量为65.9M时，sMLPNet-B达到83.4%的top-1准确率，与Swin-B性能相当，但参数量减少25%，FLOPs减少10%。
• 即使在66M参数下，该模型也未表现出过拟合迹象，而MLP-Mixer及其他MLP-based模型则严重过拟合。
• sMLPNet-T（19.2M参数）达到81.3%的top-1准确率，是参数量低于5B FLOPs的模型中最高者。
• 与DeiT-S和DeiT-B相比，sMLPNet-T在显著更少的参数量和FLOPs下实现了更高的准确率。
• sMLP模块将计算复杂度从序列长度的二次方降低为线性，从而实现对高分辨率输入的高效处理。