PaperlayerView exact revision on arXiv

Source-linked AI summary

An Image Patch is a Wave: Phase-Aware Vision MLP

Yehui Tang, Kai Han, Jianyuan Guo, Chang Xu, Yanxi Li, Chao Xu, Yunhe Wang

arXiv:2111.12294v5cs.CV

TL;DR

Existing vision MLPs aggregate tokens with fixed weights despite varying semantic content across images. Wave-MLP represents tokens with amplitude and dynamically generated phase to modulate aggregation, and experiments report strong results across classification and dense prediction tasks.

  • Problem

    Existing vision MLPs use fixed token-aggregation weights that overlook semantic differences among tokens from different input images.

  • Method

    Wave-MLP represents each token with real-valued amplitude and dynamically estimated phase, then aggregates tokens using phase-aware token mixing.

  • Results

    Wave-MLP is reported to outperform existing MLP architectures across image classification, object detection, and semantic segmentation; Wave-MLP-S reaches 82.6% top-1 ImageNet accuracy at 4.5G FLOPs versus Swin-T’s 81.3% at 4.5G FLOPs.

  • Takeaways & Limitations

    Dynamic phase information enables token aggregation according to varying image contents, supporting Wave-MLP as a backbone for dense prediction tasks.

  • Takeaways & Limitations

    The aggregation weights are set to 1 for simplicity without affecting the reported conclusion.

Abstract

from arXiv · show

In the field of computer vision, recent works show that a pure MLP architecture mainly stacked by fully-connected layers can achieve competing performance with CNN and transformer. An input image of vision MLP is usually split into multiple tokens (patches), while the existing MLP models directly aggregate them with fixed weights, neglecting the varying semantic information of tokens from different images. To dynamically aggregate tokens, we propose to represent each token as a wave function with two parts, amplitude and phase. Amplitude is the original feature and the phase term is a complex value changing according to the semantic contents of input images. Introducing the phase term can dynamically modulate the relationship between tokens and fixed weights in MLP. Based on the wave-like token representation, we establish a novel Wave-MLP architecture for vision tasks. Extensive experiments demonstrate that the proposed Wave-MLP is superior to the state-of-the-art MLP architectures on various vision tasks such as image classification, object detection and semantic segmentation. The source code is available at https://github.com/huawei-noah/CV-Backbones/tree/master/wavemlp_pytorch and https://gitee.com/mindspore/models/tree/master/research/cv/wave_mlp.

1. Introduction

Vision MLPs use fixed token-mixing weights, which overlook image-dependent semantic differences. Wave-MLP represents tokens with amplitude and dynamically estimated phase to enable phase-aware aggregation, with strong classification and dense-prediction results.

  • Vision MLPs alternate channel-mixing and token-mixing MLPs to transform individual token features and aggregate information across image patches.
  • Fixed token-mixing weights are shared across input images, neglecting token-semantic variation and limiting aggregation quality.
  • Wave-MLP represents each token with real-valued amplitude and a unit-complex phase that modulates relationships with fixed MLP weights.
  • 82.6% top-1 accuracy on ImageNet with 4.5G FLOPs surpasses Swin-T’s 81.3% accuracy with the same 4.5G FLOPs.
  • Wave-MLP also delivers strong performance on object detection and semantic segmentation.

2. Related Work

Prior computer-vision architectures center on CNNs, Transformers, and vision MLPs, with MLP designs separating channel and token processing. Related work also explores normalization changes, gating, and positional-information preservation.

  • CNN-based Architectures: CNN architectures rely on convolution and its variants, while vision Transformers brought a different computing paradigm to computer vision.
  • Vision MLP Architectures: MLP-Mixer separates channel-mixing feature extraction from token-mixing spatial-information capture.
  • Vision MLP Architectures: Related MLP research investigates affine transformation instead of Layer Normalization, gating, and methods for preserving positional information.

3. Method

Wave-MLP replaces conventional fixed token aggregation with phase-aware mixing: tokens are represented as amplitude–phase waves, dynamically estimated phases modulate interactions, and local spatial blocks support vision tasks.

  • 3.1. Preliminaries: Vision MLPs use channel-FC layers for independent token feature extraction and token-FC layers for mixing information across tokens.
  • 3.2. Phase-Aware Token Mixing: A wave-like token is represented as a real-valued amplitude combined with a periodic complex phase term.
  • 3.2. Phase-Aware Token Mixing: Phase differences determine whether token superposition is enhanced, weakened, or combined in a more complex way.
  • 3.2. Phase-Aware Token Mixing: A phase-estimation module generates phase from each token’s input features, allowing phases to vary with image content.
  • 3.2. Phase-Aware Token Mixing: PATM unfolds complex tokens into real and imaginary components, aggregates them with token-FC weights, and projects the result back to real-valued features.
  • 3.3. Wave-MLP Block: Wave-MLP blocks alternate channel-mixing MLPs with PATM-based token mixing while preserving 2D feature maps and restricting connections to local windows.

4. Experiments

Experiments evaluate Wave-MLP across image classification, object detection, instance segmentation, and semantic segmentation, comparing it with MLP, CNN, and Transformer architectures. Wave-MLP consistently achieves strong performance, while ablations support the contribution of dynamic phase information.

  • Experimental setup: Wave-MLP is evaluated on ImageNet, COCO, and ADE20K across classification, detection, instance segmentation, and semantic segmentation tasks.The experiments use ImageNet for classification, COCO with RetinaNet and Mask R-CNN, and ADE20K with Semantic FPN.
  • Image classification: 83.4% top-1 accuracy with 7.9G FLOPs enables Wave-MLP-M to outperform ResMLP-B24 at 81.0% accuracy with 23.0G FLOPs on ImageNet.Wave-MLP-T also reaches 80.6% accuracy versus 78.9% for CycleMLP-B1 with similar parameters and FLOPs.
  • Image classification: 82.6% top-1 accuracy at 4.5G FLOPs lets Wave-MLP-S surpass Swin-T at 81.3% accuracy with the same computational cost.The comparison also reports a favorable accuracy–computational-cost trade-off against typical CNN architectures.
  • Object detection and instance segmentation: Wave-MLP-T achieves 40.4% AP with RetinaNet 1×, exceeding CycleMLP-B1 by 1.8 AP at similar model size.The proposed backbone also improves results when used with Mask R-CNN.
  • Semantic segmentation: 44.4% mIoU for Wave-MLP-S exceeds PVT-Tiny at 39.8% mIoU by 4.6 percentage points on ADE20K semantic segmentation.Wave-MLP consistently outperforms existing models across parameter and computational-cost settings, including CycleMLP-B2 and Swin-S.
  • Ablation studies: Dynamic phase improves Wave-MLP-T* from 78.8% top-1 accuracy without phase information to 80.1% top-1 accuracy.Identity projection reaches 79.3%, whereas depth-wise convolution and channel-FC provide larger accuracy improvements than the baseline.
  • Ablation studies: Phase differences vary with image contents, and tokens with similar contents tend to have close phases and enhance one another.The visualization illustrates this behavior for tokens describing a house and the sky, while phase relationships differ across input images.

5. Conclusion

Wave-MLP represents each token with amplitude and phase, dynamically aggregating tokens according to their varying contents. Experiments show it suppresses existing MLP-like architectures and serves as a backbone for object detection and semantic segmentation.

  • Wave-MLP represents each token as a wave with amplitude and phase information.Amplitude is the original real-value feature, while phase modulates relationships between varying tokens and fixed MLP weights.
  • Dynamically produced phase enables token aggregation according to varying contents from different input images.
  • Wave-MLP suppresses existing MLP-like architectures in extensive experiments.
  • Wave-MLP can serve as a backbone for object detection and semantic segmentation.

6. Detailed Architectures

The Wave-MLP architecture is organized into four stages that produce hierarchical features while progressively reducing feature-map size. Its model family varies depth and width to provide four parameter and compute configurations.

  • Wave-MLP uses four stages to obtain hierarchical features and reduces feature-map size stage-wise.
  • The Wave-MLP family includes Wave-MLP-T, Wave-MLP-S, Wave-MLP-M, and Wave-MLP-B.
  • The four variants differ in parameter and computational costs through adjusted architectural depths and widths.

7. More Experiments

On COCO object detection and instance segmentation, Wave-MLP backbones achieve higher reported performance than the compared backbones under a 3× schedule with multi-scale training. Wave-MLP-T also uses fewer parameters and FLOPs than PVT-Tiny in the cited comparison.

  • 44.1 box AP and 40.1 mask AP are achieved by Wave-MLP-T with 25.3M parameters and 196.3G FLOPs.
  • Wave-MLP-T is reported as significantly superior to PVT-Tiny in the COCO comparison.
  • 39.8 box AP and 37.4 mask AP are reported for PVT-Tiny with 32.9M parameters and 208.1G FLOPs.
Loading 2111.12294v5…