MetaFormer Is Actually What You Need for Vision

Weihao Yu; Mi Luo; Pan Zhou; Chenyang Si; Yichen Zhou; Xinchao Wang; Jiashi Feng; Shuicheng Yan

Source-linked AI summary

MetaFormer Is Actually What You Need for Vision

Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, Shuicheng Yan

arXiv:2111.11418v3Nov 22, 2021cs.CVcs.AIcs.LG

TL;DR

The paper asks whether Transformer success comes from attention-specific token mixing or from a broader architecture. It defines MetaFormer and tests it with simple pooling in PoolFormer, which achieves competitive results across vision tasks, including 82.1% ImageNet-1K accuracy with fewer resources than cited baselines.

Problem
The paper questions the common belief that attention-based token mixers are primarily responsible for Transformer success in vision.
Method
The authors abstract Transformers into MetaFormer and instantiate it with a simple non-parametric pooling operator to create PoolFormer.
Results
PoolFormer achieves competitive performance across multiple vision tasks; PoolFormer-M36 reaches 82.1% top-1 accuracy on ImageNet-1K while using fewer parameters and MACs than DeiT-B and ResMLP-B24.
Takeaways & Limitations
The results support focusing future MetaFormer research on the general architecture rather than a specific token mixer, with PoolFormer serving as a baseline.
Takeaways & Limitations
The authors identify future evaluation under self-supervised and transfer learning settings, and on NLP tasks, as still needed.

Abstract

from arXiv · show

Transformers have shown great potential in computer vision tasks. A common belief is their attention-based token mixer module contributes most to their competence. However, recent works show the attention-based module in Transformers can be replaced by spatial MLPs and the resulted models still perform quite well. Based on this observation, we hypothesize that the general architecture of the Transformers, instead of the specific token mixer module, is more essential to the model's performance. To verify this, we deliberately replace the attention module in Transformers with an embarrassingly simple spatial pooling operator to conduct only basic token mixing. Surprisingly, we observe that the derived model, termed as PoolFormer, achieves competitive performance on multiple computer vision tasks. For example, on ImageNet-1K, PoolFormer achieves 82.1% top-1 accuracy, surpassing well-tuned Vision Transformer/MLP-like baselines DeiT-B/ResMLP-B24 by 0.3%/1.1% accuracy with 35%/52% fewer parameters and 50%/62% fewer MACs. The effectiveness of PoolFormer verifies our hypothesis and urges us to initiate the concept of "MetaFormer", a general architecture abstracted from Transformers without specifying the token mixer. Based on the extensive experiments, we argue that MetaFormer is the key player in achieving superior results for recent Transformer and MLP-like models on vision tasks. This work calls for more future research dedicated to improving MetaFormer instead of focusing on the token mixer modules. Additionally, our proposed PoolFormer could serve as a starting baseline for future MetaFormer architecture design. Code is available at https://github.com/sail-sg/poolformer.

1. Introduction

The paper argues that MetaFormer’s general architecture, rather than a specific token mixer such as attention, is central to strong vision-model performance. PoolFormer tests this claim with simple pooling and remains competitive across vision tasks.

Motivation: MetaFormer abstracts Transformers into a general architecture while leaving the token mixer unspecified.It retains components such as channel MLPs and residual connections around a replaceable token-mixing module.
Motivation: Spatial MLPs and other token mixers can approach Transformer performance, challenging attention’s presumed dominance.Fourier Transform token mixing, for example, reaches around 97% of vanilla Transformer accuracy.
Evidence: 82.1% top-1 accuracy: PoolFormer-M36 surpasses DeiT-B/ResMLP-B24 by 0.3%/1.1% on ImageNet-1K.It uses 35%/52% fewer parameters and 50%/62% fewer MACs than those baselines, respectively.
Evidence: PoolFormer consistently outperforms well-tuned Transformer and MLP-like models despite using pooling for only basic token mixing.The result supports the claim that the MetaFormer architecture is more essential than any specific token mixer.
Contributions: PoolFormer achieves competitive performance on image classification, object detection, instance segmentation, and semantic segmentation.The authors present it as a baseline for future MetaFormer architecture design.

2. Related work

The related-work discussion frames the paper as a response to attention-centered research and asks what fundamentally drives the success of Transformer variants. The authors investigate this question through the general architecture rather than designing another complex token mixer.

Transformer background: Transformers originated in translation and became successful across NLP before being adapted to vision tasks.Vision applications include image modeling and self-supervised pixel prediction.
Research focus: Prior work focused heavily on which attention-based token mixer is better, while paying less attention to the general architecture.The paper explicitly distinguishes its goal from designing more complicated token mixers.
Research focus: The paper’s central question is what is truly responsible for the success of Transformers and their variants.Its answer is MetaFormer, probed by using pooling as a basic token mixer.

3. Method

The method defines MetaFormer as a Transformer-like architecture with an unspecified token mixer, then instantiates it as PoolFormer using simple pooling. A hierarchical design makes this lightweight mixer practical for vision inputs.

3.1. MetaFormer: MetaFormer keeps Transformer components while allowing different concrete token mixers.Attention yields Transformer models, while spatial MLPs yield MLP-like models.
3.1. MetaFormer: Each MetaFormer block contains residual token-mixing and two-layer MLP sub-blocks.Normalization precedes token mixing, while the MLP uses learnable weights and a nonlinear activation.
3.2. PoolFormer: Pooling performs parameter-free local token mixing by averaging nearby token features.The operator is deliberately simple and is used to test whether the general architecture itself supports competitive performance.
3.2. PoolFormer: PoolFormer uses a four-stage hierarchical architecture with stage depths L/6, L/6, L/2, and L/6.This structure follows CNNs and hierarchical Transformer variants.
3.2. PoolFormer: Pooling has linear complexity in sequence length and no learnable parameters, unlike quadratic-complexity self-attention and spatial MLPs.This motivates the hierarchical design used by PoolFormer.

4. Experiments

Experiments show that PoolFormer remains competitive across image classification, object detection, instance segmentation, and semantic segmentation despite using simple pooling token mixing. Ablations further show that alternative token mixers and pooling choices preserve strong performance within limits.

Image classification: 80.3 top-1 accuracy is achieved by PoolFormer-S24 versus 75.5 for RSB-ResNet-34, despite similar model size and computation.RSB-ResNet-34 uses ∼22M parameters and 3.7G MACs, while PoolFormer-S24 performs better with pooling-based spatial modeling.
Object detection and instance segmentation: 36.2 AP is achieved by PoolFormer-S12 with RetinaNet versus 31.8 AP for ResNet-18, while Mask R-CNN results likewise favor PoolFormer-S12.For Mask R-CNN, PoolFormer-S12 reaches 37.3 bounding box AP and 34.6 mask AP, versus 34.0 and 31.2 for ResNet-18.
Semantic segmentation: 37.1 mIoU is achieved by PoolFormer-12 on ADE20K, exceeding ResNet-18 by 4.3 and PVT-Tiny by 1.5.The models use Semantic FPN backbones for semantic segmentation.
Ablation studies: 78.1% top-1 accuracy is achieved when pooling is replaced with depthwise convolution, 0.9% higher than PoolFormer-S12.The result is attributed to depthwise convolution’s better local spatial modeling ability while retaining competitive performance across token mixers.
Ablation studies: 0.5% performance drops when pooling size increases to 9, while pooling sizes 3, 5, and 7 show similar performance.Pooling size 3 is therefore adopted as the default for PoolFormer.

5. Conclusion and future work

The paper abstracts Transformers into MetaFormer, leaving the token mixer unspecified, and uses simple pooling to show that PoolFormer remains competitive across vision tasks. It identifies broader evaluation settings and NLP transfer as future tests.

MetaFormer abstracts the Transformer architecture while leaving its token mixer unspecified.
PoolFormer uses extremely simple pooling as its token mixer yet achieves competitive performance on different vision tasks.
The authors argue that MetaFormer is more essential for reasonable performance than any specific token mixer.
Future work will evaluate PoolFormer under self-supervised and transfer learning settings and investigate whether it extends to NLP tasks.

A. Detailed hyper-parameters on ImageNet-1K

The ImageNet-1K experiments use specified training hyper-parameters, including batch size, learning rate, and stochastic-depth scheduling. Hybrid pooling-attention models use a separate training configuration.

PoolFormer ImageNet-1K training uses batch size 4096 and learning rate 4 × 10^-3.These values are set based on the reported batch-size and learning-rate relation.
Stochastic-depth probability increases linearly from 0.0 in the bottom block to dr in the top block.
Hybrid models with pooling and attention use batch size 1024, learning rate 10^-3, and Layer Normalization.

B. Training for longer epochs

PoolFormer-S12 improves with substantially longer ImageNet-1K training, reaching saturated performance after around 2000 epochs, although the default comparison uses 300 epochs.

1.8% top-1 accuracy improvement is observed for PoolFormer-S12 after training for around 2000 epochs.The model’s performance is reported to saturate after approximately 2000 epochs.
300 epochs remains the default PoolFormer training duration for fair comparison with other ViT/MLP-like models.

C. Qualitative results

Grad-CAM visualizations show that all evaluated models locate semantic objects, while their activation distributions differ. PoolFormer and RSB-ResNet produce more gathered activations than DeiT and ResMLP.

All models can locate semantic objects in ImageNet-1K Grad-CAM visualizations, although ResMLP also activates irrelevant parts.
DeiT and ResMLP activations are more scattered, whereas RSB-ResNet and PoolFormer activations are more gathered.

D. Comparison between Layer Normalization and Modified Layer Normalization

Modified Layer Normalization extends normalization across both spatial and channel dimensions while preserving Layer Normalization’s affine-parameter shape.

Modified Layer Normalization computes mean and variance along spatial and channel dimensions, unlike vanilla Layer Normalization, which uses only the channel dimension.
Modified Layer Normalization keeps learnable affine parameters shaped as R_C, matching Layer Normalization.
Modified Layer Normalization can be implemented with PyTorch’s GroupNorm API by setting the group number to 1.

E. Code in PyTorch

The PyTorch implementation provides modules for normalization, pooling, MLP processing, and the PoolFormer block, with accompanying training and evaluation references.

Reference materials: The document includes PyTorch-like algorithms for the normalization modules and PoolFormer block, plus ImageNet-1K hyper-parameter and epoch-schedule tables.
Normalization: Modified Layer Normalization normalizes across channel and spatial dimensions, whereas LayerNormChannel normalizes vectors along the channel dimension.
Modules: Pooling performs token mixing by returning the pooled feature map minus the original input, accounting for the block’s residual connection.
Modules: The PoolFormer implementation defines Modified Layer Normalization, pooling, and an MLP with 1×1 convolutions as reusable modules.The MLP applies two 1×1 convolutions with GELU activation and dropout.
PoolFormer block: A PoolFormer block applies normalization, pooling-based token mixing, residual connections, stochastic depth, optional layer scaling, and an MLP.The block uses separate normalized token-mixer and MLP branches.