PaperlayerView exact revision on arXiv

Source-linked AI summary

AS-MLP: An Axial Shifted MLP Architecture for Vision

Dongze Lian, Zehao Yu, Xing Sun, Shenghua Gao

arXiv:2107.08391v2cs.CV

TL;DR

Vision architectures have progressed from convolutional and attention-based models, while MLP-Mixer's global token mixing leaves local feature interaction underused and windowed MLPs face capacity or input-size constraints. The paper proposes AS-MLP, which uses horizontal and vertical axial shifts in a pure MLP architecture to model local dependencies and CNN-like receptive fields. AS-MLP reaches competitive classification, detection, and segmentation results, including 83.3% Top-1 accuracy on ImageNet-1K with 88M parameters and 15.2 GFLOPs.

  • Problem

    MLP-Mixer emphasizes global token mixing but underuses local information, while windowed token mixing creates capacity or variable-input-size limitations.

  • Method

    AS-MLP introduces horizontal and vertical axial feature shifts to aggregate local features without fixed windows and enable CNN-like receptive-field and dilation design.

  • Results

    AS-MLP achieves competitive performance across image classification, object detection, and semantic segmentation, including 83.3% Top-1 accuracy on ImageNet-1K with 88M parameters and 15.2 GFLOPs.

  • Takeaways & Limitations

    AS-MLP establishes a strong MLP-based vision baseline that transfers to downstream tasks and performs competitively against transformer-based architectures.

  • Takeaways & Limitations

    The paper identifies natural-language processing and further downstream-task evaluation as future work.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper motivates locality-focused MLPs because global token mixing underuses local information, while windowed alternatives face capacity or input-size constraints. AS-MLP addresses this with axial feature shifts that aggregate neighboring positions without fixed windows and supports CNN-like receptive-field design.

  • Motivation: MLP-Mixer captures global receptive fields through matrix transposition and token mixing but rarely exploits local information important for low-level feature extraction.The motivation contrasts global long-range dependency modeling with locality emphasized in CNN-like architectures.
  • Problem: Windowed token mixing introduces a 49 × 49 parameter matrix for a 7 × 7 window, limiting capacity when shared across windows.Unshared window-specific weights instead cannot adapt to downstream tasks with varying input sizes.
  • Method: AS-MLP spatially shifts features horizontally and vertically to aggregate surrounding positions without fixing a window or window size.The axial shift divides channels into k groups rather than k^2 groups for a k × k receptive field.
  • Method: Axial shifting enables AS-MLP to design receptive-field size and dilation in the same spirit as convolutional kernels.This preserves a pure MLP architecture while allowing convolution-like structural choices.
  • Results: 83.3% Top-1 accuracy is achieved on ImageNet-1K with 88M parameters and 15.2 GFLOPs, without extra training data.The paper reports competitive performance against transformer-based architectures and downstream transfer to object detection and semantic segmentation.

2 RELATED WORK

The related work situates AS-MLP among CNN, transformer, and MLP-based vision architectures. These lines emphasize the progression from convolutional feature extraction to attention and pure-MLP alternatives.

  • CNN-based Architectures: CNN-based architectures established convolutional image-feature extraction, with VGG using stacked 3 × 3 convolutions and ResNet using residual connections.The supplied passage introduces these architectures as representative CNN developments.
  • Transformer-based Architectures: Transformer-based vision architectures use attention to model relationships between spatial features, with ViT abandoning convolution layers.The passage also identifies BERT, ViT, and DeiT as milestones connecting transformer progress in NLP and vision.
  • MLP-based Architectures: MLP-based architectures transmit information between spatial features using MLP operations, with MLP-Mixer, FF, and Res-MLP providing concise pure-MLP frameworks.The supplied related-work passage presents these as preceding MLP-based approaches.

3 THE AS-MLP ARCHITECTURE

AS-MLP builds a hierarchical vision architecture from axial feature shifts and pure MLP blocks, introducing local spatial interactions without windows. Its axial design supports configurable receptive fields and lower complexity than Swin Transformer.

  • AS-MLP block: Each AS-MLP block combines normalization, axial shift, an MLP, and a residual connection.Channel projection plus vertical and horizontal shifts translate features along spatial directions.
  • AS-MLP block: A shift size of 3 divides channels into three parts shifted by −1, 0, and 1 units, using zero padding before channel projection.The shifted features are extracted from the resulting region for subsequent processing.
  • Complexity: Axial shifting requires no multiplication or addition operations, giving AS-MLP slightly lower complexity than Swin Transformer.The stated complexities are Ω(MSA) = 4hwC^2 + 2(hw)^2C, Ω(W-MSA) = 4hwC^2 + 2M^2hwC, and Ω(AS) = 4hwC^2.
  • Comparisons: Unlike MLP-Mixer’s global token mixing, AS-MLP models local dependencies through horizontal and vertical axial shifts with channel projection.Compared with original shift, axial operation divides channels into k rather than k^2 groups for a k × k receptive field, reducing complexity.

4 EXPERIMENTS

Experiments evaluate AS-MLP on image classification, mobile models, block configurations, object detection, instance segmentation, semantic segmentation, and feature visualization. The results show competitive or superior performance across these settings, while ablations examine design choices such as shift configuration and connection type.

  • Image classification: 76.05% vs. 75.11% shows AS-MLP significantly exceeds Swin Transformer in the mobile setting with models of about 10M parameters.
  • AS-MLP block ablations: Parallel connection consistently outperforms serial connection across different shift sizes, except that serial is better at shift size 1.At shift size 1, only channel-mixing MLP is used, so the authors describe that case as unrepresentative.
  • AS-MLP block ablations: Shift size (5, 5) outperforms Global-MLP, Axial-MLP, Window-MLP, and single-direction shift baselines.
  • Downstream tasks: 46.5 vs. 45.8 MS mIoU shows AS-MLP-T outperforming Swin-T on ADE20K with slightly lower FLOPs.For large models, AS-MLP-B obtains 49.5 MS mIoU with 121M parameters and 1166 GFLOPs, versus Swin-B's 49.7 MS mIoU with the same parameters and 1188 GFLOPs.
  • Visualization: AS-MLP feature heatmaps better focus on object regions than Swin Transformer, while horizontal and vertical shifts emphasize corresponding object orientations.Horizontal shifts better focus on vertical object parts, whereas vertical shifts better focus on horizontal parts.

5 CONCLUSION AND FUTURE WORK

The conclusion presents AS-MLP as a local-feature-focused axial-shifted MLP architecture with competitive classification and downstream-task results. Future work will examine natural language processing and further downstream-task performance.

  • Conclusion: AS-MLP uses axial feature shifting to emphasize local feature extraction and channel interaction between spatial positions.
  • Conclusion: 83.3% Top-1 accuracy with 88M parameters and 15.2 GFLOPs is reported on ImageNet-1K.
  • Conclusion: AS-MLP is applied to object detection and semantic segmentation, with results competitive or better than transformer-based architectures.
  • Future work: Future work will investigate AS-MLP in natural language processing and further explore its downstream-task performance.

A.1 THE DETAILED CONFIGURATIONS OF DIFFERENT ARCHITECTURES

The architecture configurations assume 224 × 224 input images and specify stage output sizes and shift settings for comparing different architectures.

  • 224 × 224 is the assumed input image size for the detailed architecture configurations.
  • The configuration table records each stage's output size and uses “Concat n × n” for concatenating n × n neighboring patch features.
  • “shift size (5, 5)” denotes shifts of size 5 in both horizontal and vertical directions.

A.2 THE COMPUTATIONAL COMPLEXITY OF AS-MLP ARCHITECTURE

AS-MLP’s computational analysis specifies its input, architectural, and complexity variables, while experiments also cover mobile variants, downstream settings, and related shift-based baselines.

  • Computational complexity: The complexity analysis defines the input image, patch size, stage block counts, channel dimension, and MLP ratio before presenting layer costs.Only convolution operations are counted in the reported computational complexity.
  • Mobile setting: AS-MLP and Swin are both configured for mobile settings with channel dimension C = 64 and four stages containing {2, 2, 2, 2} blocks.The mobile Swin configuration additionally specifies heads {2, 4, 8, 16}, whereas mobile AS-MLP does not use attention heads.
  • Baseline configurations: Global-MLP and Axial-MLP use fixed-size global or axial token-mixing operations, while Window-MLP performs MLP operations within a fixed 7 × 7 window.The fixed-dimension token-mixing weights of Global-MLP and Axial-MLP cannot adapt to downstream tasks with varying input sizes.
  • Shift-based comparisons: AS-MLP uses spatial axial shifts for image classification, object detection, and segmentation, reducing shift-split complexity relative to temporal-shift video modeling.The paper contrasts AS-MLP’s spatial focus with TSM’s temporal shift and reports that shifting more channels does not significantly degrade pure-MLP performance.

B.1 MORE EXPERIMENTAL RESULTS ON COCO

Additional COCO experiments compare AS-MLP-T with Swin-T under two detection frameworks and report ImageNet throughput evaluation settings.

  • COCO detection and segmentation: 44.0 vs. 43.7 APb under Mask R-CNN and 48.4 vs. 48.1 APb under Cascade Mask R-CNN favor AS-MLP-T over Swin-T.These comparisons use the 3x schedule, corresponding to 36 epochs.
  • Throughput evaluation: Throughput evaluation for ImageNet-1K reports measurements at batch sizes 1, 4, 8, 16, 32, 64, and 128.Throughput is measured on a single 32GB V100 GPU with batch size 64 in the table comparison.

B.3 EVALUATION ACCURACY

Training curves compare AS-MLP with Swin across ImageNet, COCO, and ADE20K, showing similar final ImageNet accuracy and stronger downstream trajectories for AS-MLP-T.

  • ImageNet evaluation: 81.3 vs. 81.3 final accuracy shows AS-MLP-T and Swin-T converge to similar ImageNet performance.AS-MLP-T keeps pace with Swin-T at each ImageNet training epoch.
  • Downstream evaluation: AS-MLP-T achieves better performance than Swin-T in the early stages of COCO object detection and ADE20K semantic segmentation training.The reported advantage continues throughout the training process.
  • Throughput: Figure 5 organizes throughput curves by batch sizes 1, 4, 8, 16, 32, 64, and 128.The figure caption identifies the batch-size conditions but does not state a numerical outcome.
  • Training curves: Figure 6 presents evaluation accuracy for AS-MLP and Swin Transformer during training on ImageNet, COCO, and ADE20K.The figure spans image classification, object detection, and semantic segmentation datasets.

C THE VISUALIZATION OF RESULTS ON COCO AND ADE20K

The paper visualizes AS-MLP-T outputs for COCO detection and instance segmentation and ADE20K semantic segmentation using established task-specific frameworks.

  • COCO visualization: Figure 7 visualizes object detection and instance segmentation on COCO using Cascade Mask R-CNN with an AS-MLP-T backbone.The paper states that the object is detected and segmented correctly.
  • ADE20K visualization: Figure 8 visualizes semantic segmentation on ADE20K using UperNet with an AS-MLP-T backbone.The visualization is part of the paper’s qualitative downstream evaluation.
Loading 2107.08391v2…