[Paper Review] AS-MLP: An Axial Shifted MLP Architecture for Vision

Dongze Lian, Zehao Yu+2 more|arXiv (Cornell University)|Jul 18, 2021

Advanced Neural Network Applications47 references85 citations

TL;DR

AS-MLP introduces axial channel shifts in an MLP framework to capture local dependencies, achieving competitive ImageNet performance and extending to downstream tasks like object detection and segmentation.

ABSTRACT

An Axial Shifted MLP architecture (AS-MLP) is proposed in this paper. Different from MLP-Mixer, where the global spatial feature is encoded for information flow through matrix transposition and one token-mixing MLP, we pay more attention to the local features interaction. By axially shifting channels of the feature map, AS-MLP is able to obtain the information flow from different axial directions, which captures the local dependencies. Such an operation enables us to utilize a pure MLP architecture to achieve the same local receptive field as CNN-like architecture. We can also design the receptive field size and dilation of blocks of AS-MLP, etc, in the same spirit of convolutional neural networks. With the proposed AS-MLP architecture, our model obtains 83.3% Top-1 accuracy with 88M parameters and 15.2 GFLOPs on the ImageNet-1K dataset. Such a simple yet effective architecture outperforms all MLP-based architectures and achieves competitive performance compared to the transformer-based architectures (e.g., Swin Transformer) even with slightly lower FLOPs. In addition, AS-MLP is also the first MLP-based architecture to be applied to the downstream tasks (e.g., object detection and semantic segmentation). The experimental results are also impressive. Our proposed AS-MLP obtains 51.5 mAP on the COCO validation set and 49.5 MS mIoU on the ADE20K dataset, which is competitive compared to the transformer-based architectures. Our AS-MLP establishes a strong baseline of MLP-based architecture. Code is available at https://github.com/svip-lab/AS-MLP.

Motivation & Objective

• Motivate the need for leveraging local feature interactions in MLP-based vision models rather than solely global token mixing.
• Propose a lightweight axial shift mechanism to enable local receptive fields within a pure MLP architecture.
• Design a scalable AS-MLP backbone with four stages and hierarchical feature merging.
• Demonstrate competitive performance on ImageNet-1K and competitive transfer to downstream tasks (COCO detection, ADE20K segmentation).
• Provide ablations to understand the impact of shift configurations, padding, dilation, and connection style.

Proposed method

• Introduce Axial Shifted MLP (AS-MLP) blocks that perform horizontal and vertical feature shifts followed by channel projection to enable local feature aggregation.
• Use Norm layers, residual connections, and MLP-based channel mixing to combine shifted features.
• Shift operations aggregate information from different spatial positions without relying on full attention, keeping complexity low.
• Adopt a four-stage Swin-like backbone with patch partitioning and patch merging to form hierarchical representations.
• Ablate shift sizes, padding methods, dilation rates, and serial vs parallel connections to identify effective configurations.

Experimental results

Research questions

• RQ1 Can axial (horizontal and vertical) feature shifts in an MLP-only backbone achieve competitive local receptive fields comparable to CNNs or window-based transformers?
• RQ2 What shift sizes, padding strategies, and connectivity (serial vs parallel) maximize accuracy while maintaining efficiency?
• RQ3 How well does AS-MLP transfer to downstream tasks such as object detection and semantic segmentation compared to transformer-based backbones?
• RQ4 What is the trade-off between model size, FLOPs, and accuracy for AS-MLP variants on ImageNet-1K?
• RQ5 Is AS-MLP capable of mobile-friendly performance relative to Swin Transformer under similar resource constraints?

Key findings

• AS-MLP achieves 83.3% Top-1 accuracy on ImageNet-1K with 88M parameters and 15.2 GFLOPs (AS-MLP-B, 224x224).
• AS-MLP-B at 384x384 attains 84.3% Top-1 with 88M params and 44.6 GFLOPs.
• AS-MLP-S reaches 83.1% Top-1 with 50M parameters and 8.5 GFLOPs.
• AS-MLP-T achieves 81.3% Top-1 with 28M parameters and 4.4 GFLOPs.
• In mobile settings, AS-MLP (mobile) outperforms Swin (mobile) in Top-1 (76.05% vs 75.11%).
• AS-MLP demonstrates competitive results on COCO object detection (e.g., AS-MLP-B 51.5 APb) and ADE20K segmentation (AS-MLP-B 49.5 MS mIoU) compared to transformer baselines.