S${}^2$-MLP: Spatial-Shift MLP Architecture for Vision

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.

References in corpus (12)

• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
• Decoupled Weight Decay Regularization
• Gaussian Error Linear Units (GELUs)
• MLP-Mixer: An all-MLP Architecture for Vision
• Transformer in Transformer
• Conditional Positional Encodings for Vision Transformers
• CvT: Introducing Convolutions to Vision Transformers
• Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks
• Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet
• Rethinking Spatial Dimensions of Vision Transformers
• Container: Context Aggregation Network
• Pay Attention to MLPs

Cited by in corpus (13)

• TSMixer: Lightweight MLP-Mixer Model for Multivariate Time Series Forecasting
• CycleMLP: A MLP-like Architecture for Dense Prediction
• AS-MLP: An Axial Shifted MLP Architecture for Vision
• S${}^2$-MLPv2: Improved Spatial-Shift MLP Architecture for Vision
• ConvMLP: Hierarchical Convolutional MLPs for Vision
• Rethinking Token-Mixing MLP for MLP-based Vision Backbone
• Hire-MLP: Vision MLP via Hierarchical Rearrangement
• MSCA-Net:Multi-Scale Context Aggregation Network for Infrared Small Target Detection
• An Image Patch is a Wave: Phase-Aware Vision MLP
• MOI-Mixer: Improving MLP-Mixer with Multi Order Interactions in Sequential Recommendation
• RaftMLP: How Much Can Be Done Without Attention and with Less Spatial Locality?
• PointMixer: MLP-Mixer for Point Cloud Understanding
• Global Interaction Modelling in Vision Transformer via Super Tokens