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

Tan Yu; Xu Li; Yunfeng Cai; Mingming Sun; Ping Li

Source-linked AI summary

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

Tan Yu, Xu Li, Yunfeng Cai, Mingming Sun, Ping Li

arXiv:2106.07477v2Jun 14, 2021cs.CVcs.AIcs.LG

TL;DR

MLP-Mixer’s token-mixing design leaves an accuracy gap on medium-scale datasets and can over-fit because it has global, spatial-specific mixing. The paper proposes S2-MLP, which uses channel-mixing MLPs with parameter-free spatial shifts for local, spatial-agnostic communication. On ImageNet-1K, S2-MLP outperforms MLP-Mixer and matches ViT’s accuracy with fewer parameters and FLOPs.

Problem
MLP-Mixer performs less strongly than ViT on medium-scale ImageNet-1K and ImageNet-21K datasets, motivating a reconsideration of its token-mixing MLP.
Method
S2-MLP uses only channel-mixing MLPs and a parameter-free spatial-shift operation with a local reception field for communication between patches.
Results
S2-MLP achieves higher ImageNet-1K recognition accuracy than MLP-Mixer and comparable accuracy to ViT with fewer parameters and FLOPs.
Takeaways & Limitations
Spatial shifting provides a simpler pure-MLP alternative that is competitive with ViT on ImageNet-1K while improving over MLP-Mixer.
Takeaways & Limitations
FLOPs counting considers only multiplication operations between floating-point numbers.

Abstract

from arXiv · show

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.

1. Introduction

The paper identifies MLP-Mixer’s medium-scale-data accuracy gap as a consequence of its globally receptive, spatial-specific token-mixing MLP, then proposes S2-MLP with spatial-shift communication. S2-MLP uses only channel-mixing MLPs and achieves stronger ImageNet-1K results than MLP-Mixer while matching ViT with lower complexity.

Motivation: MLP-Mixer combines channel-mixing and token-mixing MLPs, but retains an accuracy gap versus ViT on ImageNet1K and ImageNet-21K.Its token-mixing MLP communicates across spatial locations, while the channel-mixing MLP communicates across channels.
Motivation: Token-mixing MLP applies a learned matrix to patch features, producing weighted sums whose weights are input-agnostic, unlike Transformer attention.The operation resembles depthwise convolution but uses a global reception field and spatial-specific weights.
Proposed architecture: S2-MLP processes non-overlap image patches through stacked blocks, global average pooling, and a fully connected classifier.Each block contains fully connected layers, GELU activations, layer normalization, skip connections, and spatial shifting.
Motivation: Global reception and spatial-specific weights make token-mixing MLP more flexible but prone to over-fitting, increasing its dependence on huge training sets.Depthwise convolution instead has a local reception field and shared kernels across spatial locations.
Proposed architecture: S2-MLP replaces token-mixing MLP with a parameter-free spatial-shift operation while retaining only channel-mixing MLPs.Spatial shifting communicates between locations using a local reception field and spatial-agnostic channel groups.
Results: S2-MLP attains higher ImageNet1K recognition accuracy than MLP-Mixer and comparable accuracy to ViT with a simpler architecture and fewer parameters and FLOPs.The reported comparison emphasizes both recognition accuracy and computational efficiency.

2. Related Work

Related work progresses from Transformer-based visual backbones to architectures using only MLP layers. These approaches differ mainly in how they model spatial communication and in their dependence on training scale.

Transformer-based vision models: ViT replaces CNN convolution with self-attention to model relations between patches at different spatial locations.Its standard Transformer blocks reduce hand-crafted manipulation and inductive biases relative to CNN kernel design.
Transformer-based vision models: Subsequent Transformer-based models improve visual representation through changes including pyramidal computation, progressive tokenization, local patch modeling, positional encoding, global tokens, pooling, and hierarchy.The cited examples represent multiple architectural directions for improving ViT.
MLP-based vision models: Other MLP-based models modify MLP-Mixer with elements such as a global CLS token, positional embedding, or alternative pure-MLP designs.The related-work passage identifies FF and Res-MLP as examples.

3. Method

S2-MLP replaces MLP-Mixer’s token-mixing MLP with parameter-free spatial shifts while retaining channel-mixing MLPs. Its blocks combine local spatial communication with channel projections, and the operation is equivalent to depthwise convolution with fixed group-specific kernels.

Backbone: The backbone comprises patch embedding, N S2-MLP blocks, and a classification fully connected layer.Images are divided into w × h non-overlapping patches, each projected to an embedding and normalized before entering the blocks.
S2-MLP block: S2-MLP contains only channel-mixing MLPs and uses spatial shifting to communicate between patches.The architecture stacks N same-sized blocks; all fully connected layers mix channels, while the spatial-shift module handles spatial communication.
Spatial-shift operation: Each block splits channels into groups and shifts each group in a different spatial direction.The default operation uses four groups shifted by ±1 along width and height, allowing patches to absorb content from adjoining patches.
S2-MLP block: Each S2-MLP block uses four fully connected layers, two layer-normalization layers, two GELU layers, and two skip-connections.The channel dimension is preserved by the first two fully connected layers, expanded to hidden size c̄ = rc, then projected back to c.
Spatial-shift operation: Spatial shifting is parameter-free, spatial-agnostic, and maintains a local reception field.After shifting, each patch receives visual content from neighboring patches without introducing token-mixing parameters.
Relations with depthwise convolution: Spatial shifting is equivalent to depthwise convolution with fixed, group-specific kernels that share weights within channel groups.This connects S2-MLP to depthwise convolution while distinguishing it from MLP-Mixer’s global, spatial-specific token mixing.

4. Experiments

Experiments on ImageNet-1K and ImageNet100 compare S2-MLP with prior models and evaluate architectural choices. S2-MLP improves accuracy-efficiency trade-offs, while ablations show costs and accuracy depend on depth, width, expansion ratio, shifting, input scale, and patch size.

4.1. Main results: S2-MLP matches ViT accuracy on ImageNet-1K with a simpler structure and fewer parameters and FLOPs.The paper also reports that it does not reach the accuracy of state-of-the-art Transformer models such as CaiT, Swin-B, and Nest-B.
4.1. Main results: S2-MLP-wide reaches 80.0% top-1 accuracy, compared with 76.4% for MLP-Mixer and 74.9% for FF.All three models use 12 blocks with hidden size 768.
4.2. Ablation study: Increasing depth from N = 1 to N = 12 raises top-1 accuracy from 56.7% to 87.1%, while parameters rise from 6.5M to 71M and FLOPs from 1.3B to 14B.Accuracy drops when depth increases further from 12 to 16, possibly because ImageNet100 is small and over-fitting occurs.
4.2. Ablation study: Increasing hidden size c from 192 to 768 raises top-1 accuracy from 79.7% to 87.1%, while parameters increase from 4.3M to 71M and FLOPs from 0.9b to 14B.Accuracy saturates beyond c = 768, which is selected as the default setting.
4.2. Ablation study: Increasing expansion ratio r from 1 to 3 raises top-1 accuracy from 86.1% to 87.0% and parameters from 29M to 57M, after which accuracy saturates or worsens.The authors attribute the degradation at larger r possibly to over-fitting on ImageNet100.
4.2. Ablation study: Four-direction shifting is the default: removing shifting performs poorly, two-dimensional shifting outperforms one-dimensional shifting, and four directions suffice.The comparison also finds horizontal shifting more useful than vertical shifting.
4.2. Ablation study: Increasing input scale from 112 × 112 to 336 × 336 raises top-1 accuracy from 80.6% to 88.2%, but increases FLOPs from 3.5B to 31B without changing parameters.The paper does not recommend larger inputs when computing resources are limited.
4.2. Ablation study: Increasing patch size from 16 to 32 reduces FLOPs from 14B to 3.5B but lowers top-1 accuracy from 87.1% to 81.0%.Larger patches are recommended when fast inference is more important than fine-level detail capture.

5. Conclusion

S^2-MLP is a pure MLP vision architecture that replaces convolution and self-attention with parameter-free spatial shifting and channel-mixing MLPs. On ImageNet-1K, it outperforms MLP-Mixer and ResMLP while using fewer parameters and FLOPs than ViT.

S^2-MLP uses a pure MLP structure without convolution or self-attention, with spatial shifting for communication between spatial locations.The spatial-shift operation is simple, parameter-free, and efficient.