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 spatialagnostic. 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.

ConvMLP: Hierarchical Convolutional MLPs for Vision

“…Following that, gMLP [22] proposed a Spatial Gating Unit to process spatial features. S 2 -MLP [39] adopts shifted spatial feature maps to augment information communication.…”

PyramidTNT: Improved Transformer-in-Transformer Baselines with Pyramid Architecture

“…Recently, various variants are developed to achieve a better trade-off between accuracy and computational cost. For example, shift operation is introduced in S 2 -MLP [37] and AS-MLP [18] to exchange information across different tokens. Hire-MLP [8] present a hierarchical rearrangement operation, where the inner-region rearrangement and cross-region rearrangement capture local information and global context, respectively.…”

Are we ready for a new paradigm shift? A Survey on Visual Deep MLP

“…Research first propose a spatial-shift MLP architecture for vision, called S2MLP [93]. The actual practice is quite simple.…”

“…As shown in the leftmost subdiagram in Figure 13. Models of this design also include FeedForward [37], ResMLP [40], gMLP [88], S2MLP [93], CCS [106], RaftMLP [92], Sparse-MLP(MoE) [103]. Due to the limited computing resources, the patch partition during patch embedding of the single-stage model is usually large, e.g.…”