Hire-MLP: Vision MLP via Hierarchical Rearrangement

TL;DR

Hire-MLP introduces a hierarchical rearrangement approach for vision MLPs, capturing local and global information effectively, leading to competitive performance across multiple vision tasks with improved flexibility and efficiency.

Contribution

The paper proposes Hire-MLP, a novel vision MLP architecture with hierarchical rearrangements for better spatial information capture and versatility as a backbone for various vision tasks.

Findings

Achieves 83.8% top-1 accuracy on ImageNet.

Surpasses previous models in object detection and segmentation.

Demonstrates effective global and local feature integration.

Abstract

Previous vision MLPs such as MLP-Mixer and ResMLP accept linearly flattened image patches as input, making them inflexible for different input sizes and hard to capture spatial information. Such approach withholds MLPs from getting comparable performance with their transformer-based counterparts and prevents them from becoming a general backbone for computer vision. This paper presents Hire-MLP, a simple yet competitive vision MLP architecture via \textbf{Hi}erarchical \textbf{re}arrangement, which contains two levels of rearrangements. Specifically, the inner-region rearrangement is proposed to capture local information inside a spatial region, and the cross-region rearrangement is proposed to enable information communication between different regions and capture global context by circularly shifting all tokens along spatial directions. Extensive experiments demonstrate the effectiveness of Hire-MLP as a versatile backbone for various vision tasks. In particular, Hire-MLP achieves competitive results on image classification, object detection and semantic segmentation tasks, e.g., 83.8% top-1 accuracy on ImageNet, 51.7% box AP and 44.8% mask AP on COCO val2017, and 49.9% mIoU on ADE20K, surpassing previous transformer-based and MLP-based models with better trade-off for accuracy and throughput. Code is available at https://github.com/ggjy/Hire-Wave-MLP.pytorch.