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

Introduction

In computer vision, convolutional neural networks (CNNs) have been the mainstream architectures for a long time . It is challenged by the recent works , in which a standard Transformer model can also work well on various computer vision tasks, such as image classification, object detection and semantic segmentation . Considering the high complexity of self-attention modules in the vision transformer, more simple architectures (e.g., MLP-Mixer , ResMLP) stacking only multi-layer perceptrons (MLPs) have attracted much attention. Compared with CNNs and Transformers, these vision MLP architectures involve less inductive bias and have potential to be applied on more diverse tasks.

Taking a sequence of image patches (tokens) as input, MLP-like models mainly contain two separable blocks, i.e., channel-mixing MLP and token-mixing MLP, both composing of full-connected layers and activation functions. The channel-mixing MLP transforms feature of each token and the token-mixing MLP tries to aggregate information from different tokens. By stacking these two types of MLP block alternatively, the simple MLP architecture could have sufficient capacity to extract features and achieve good performance on vision tasks.

However, the performance of MLP architecture is still inferior to that of SOTA Transformer and CNN architectures. We point out that one of the bottlenecks for vision MLP lies in its manner of aggregating different tokens, i.e., mixing different tokens with fixed weights of fully-connected layers. Recall that Transformer aggregates tokens with weights dynamically adjusted by the attention mechanism. The inner products between different tokens are calculated and tokens with higher similarities tend to have larger weights in the aggregation process of each other. However, the existing vision MLP models aggregate different tokens with fixed weights. The same weights are used for tokens from different input images, neglecting differences in semantic information of various tokens, which may not aggregate tokens well for all the input images.

Different from Transformer that delicately designs the attention mechanism, we aim to improve the representation way of tokens for dynamically aggregating them according to their semantic contents. Actually, in quantum mechanics, an entity (e.g., electron, photon) is usually represented by a wave function (e.g., de Broglie wave) containing both amplitude and phase . The amplitude part measures the maximum intensity of a wave and the phase part modulates the intensity by indicating the location of a point in the wave period. Inspired by the quantum mechanics, we describe each token as a wave to realize the dynamic aggregation procedure of tokens.

In this paper, we present a novel vision MLP architecture (dubbed as Wave-MLP), which takes each token as a wave with both amplitude and phase. The amplitude is the real-value feature representing the content of each token, while the phase term is a unit complex value modulating the relationship between tokens and fixed weights in MLP. The phase difference between these wave-like tokens affects their aggregated output and tokens with close phases tend to enhance each other. Considering that tokens from different input images contain diverse semantic contents, we use a simple module to dynamically estimate the phase for each token. With tokens equipped with amplitude and phase information, we introduce a phase-aware token mixing module (PATM in Figure 2) to aggregate these tokens. The whole Wave-MLP architecture is constructed by stacking the PATM module and channel-mixing MLP, alternately.

The proposed Wave-MLP architecture shows a large superiority to the existing architectures (shown in Figure 1). For example, the proposed Wave-MLP-S model achieves 82.6% top-1 accuracy on ImageNet with 4.5G FLOPs, which significantly surpasses Swin-T with 81.3% accuracy and 4.5G FLOPs. Besides, Wave-MLP also achieves strong performance on the dense prediction tasks such as object detection and semantic segmentation.

The paper is organized as follows: Section 2 briefly reviews the existing works about designing model architectures, and Section 3 discusses the proposed Wave-MLP architecture detailedly. In Section 4, we empirically investigate the method’s effectiveness on multiple vision tasks and make conclusions in Section 6.

Related Work

CNN-based Architectures. Convolutional neural networks (CNNs) have been the mainstream in computer vision for a long time. The prototype of CNN model is presented in for the document recognition task, where convolution is the core operation. Beginning with great success of AlexNet in ILSVRC 2012, various architectures such as GoogleNet , VGGNet , ResNet , RegNet are developed. Though the model architectures become more complex for pursuing high performance, the core operations have always the convolution and its variants. The occurrence of new computing paradigm such as vision Transformer , vision MLP bring new blood to the area of architecture design in computer vision.

Transformer-based Architectures. Transformer is originally proposed for the natural language processing (NLP) tasks such as language modeling and machine translation. Dosovitskiy et al. introduce it to computer vision and achieve excellent performance on image classification tasks especially when training data are extremely sufficient. Touvron et al. refine the training recipe and present a teacher-student strategy specific to transformers, which produce competitive transformer models trained on ImageNet from scratch. Then many works explore the architecture design of vision transformers . For example, Han et al. present a nested transformer architecture to capture global and local information simultaneously. To be compatible with the dense prediction task such as object detection and semantic segmentation, hierarchical architectures are adopted in , which splits the whole architecture into multiple stages and reduce the spatial resolution stage-wisely. Swin Transformer extract representation with shifted windows and limit the self-attention in local regions. Compared with the self-attention in connecting all the tokens in a layer, the shifted window operation is more efficient.

MLP-based Architectures. Recently, MLP-like architectures composing of fully connected layers and non-linear activation functions have been paid much attention . Though they have more simple architectures and introduce less inductive bias, their performances are still comparable with SOTA models. The MLP-Mixer model uses two type of MLP layers, i.e., channel-mixing MLP and token-mixing MLPs. The channel-MLP extract features for each tokens while the token-mixing MLPs capture the spatial information. Touvron et al. present a similar architecture and replace the Layer Normalization with the simpler affine transformation. Liu et al. empirically validate that MLP architectures with gating can achieve similar performance with Transformers in both language and vision tasks. To preserve the positional information of input images, Hou et al. keep the 2D shape of the input image and extract features by permuting them along width and height, respectively. Based on MLP-Mixer, Yu et al. replace the token-mixing MLP with a spatial shift operation for capturing the local spatial information, which is also computationally efficient. Currently, Lian et al. propose to shift tokens along two orthogonal directions to obtain an axial receptive field. Chen et al. propose a cycle fully-connected layer, which mixes information along the spatial and channel dimensions simultaneously and can cope with variable input image scales. Different from them, we explore how to represent the tokens in vision MLP and take each token as a wave with both amplitude and phase. Empirically, we find that our Wave-MLP architecture achieves a better trade-off between accuracy and computational cost compared with the existing architectures.

Method

In this section, we discuss the proposed Wave-MLP models detailedly. After introducing the vision MLP architecture briefly, we present the phase-aware token mixing module (PATM), which represents each token as a wave and aggregate them by considering amplitude and phase simultaneously. At last, we describe the blocks in Wave-MLP and architecture variants with different computational costs.

A MLP-like model is a neural architecture mainly composed of full-connected layers and non-linear activation functions. For the vision MLP, it first splits an image into multiple patches (also referred to as tokens) and then extract their features with two components, i.e., channel-FC and token-FC described as following.

Denote the intermediate feature containing $n$ tokens as $Z=[{\bm{z}}_{1},{\bm{z}}_{2},\cdots,{\bm{z}}_{n}]$, where each token ${\bm{z}}_{j}$ is a $d$-dimension vector. The channel-FC is formulated as:

where $W^{c}$ is the weight with learnable parameters. The channel-FC operates on each token independently to extract their features. To enhance the transformation ability, multiple channel-FC layers are usually stacked together with the non-linear activation function, which constructs a channel-mixing MLP.

To aggregate information from different tokens, the token-FC operation is required, i.e.,

2 Phase-Aware Token Mixing

To dynamically modulate the relationship between tokens and fixed weights in MLP for aggregating tokens more properly, we take each token as a wave with both amplitude and phase. We firstly discuss the wave-like representation of a token and then present the phase-aware token mixing module (PATM) for aggregating tokens.

where ${\rm atan2}(x,y)$ is the two-argument arctangent function. As shown in the above equations, the phase difference $|{\bm{\theta}}_{j}-{\bm{\theta}}_{i}|$ between two tokens has a large impact on the amplitude of aggregated result ${\bm{z}}_{r}$. An intuitive diagram is shown in Figure 3. The left is the superposition of two waves in the complex-value domain, while the right shows how their projections along the real axis varies w.r.t. the phase. When two tokens have the same phase (${\bm{\theta}}_{j}={\bm{\theta}}_{i}+2\bm{\pi}*m,m\in[0,\pm 2,\pm 4,\cdots]$), they will be enhanced by each other, i.e., $|{\bm{z}}_{r}|=|{\bm{z}}_{i}|+|{\bm{z}}_{j}|$ (Figure 3 (b)). For the opposite phase (${\bm{\theta}}_{j}={\bm{\theta}}_{i}+\bm{\pi}*m,m\in[\pm 1,\pm 3,\cdots]$), the resultant wave will be weakened ($|{\bm{z}}_{r}|=||{\bm{z}}_{i}|-|{\bm{z}}_{j}||$). In other cases, their interaction is more complex but whether they will be enhanced or weakened also depends on the phase difference (Figure 3 (a)). Note that the classical representation strategy with only real-value feature is a special case of Eq 3, whose phase ${\bm{\theta}}_{j}$ is only the integer multiple of $\bm{\pi}$.

Amplitude. To get the wave-like tokens in Eq. 3, both amplitude and phase information are required. The amplitude $|{\bm{z}}_{i}|$ is similar to the real-value feature in the traditional model, expect for an absolute operation. Actually, the element-wisely absolute operation can be absorbed into the phase term, i.e., $|z_{j,t}|e^{i\theta_{j,t}}=z_{j,t}e^{i\theta_{j,t}}$ if $z_{j,t}>0$, and $|z_{j,t}|e^{i\theta_{j,t}}=z_{j,t}e^{i(\theta_{j,t}+\pi)}$ otherwise, where $z_{j,t}$ and $\theta_{j,t}$ denote the $t$-th element in ${\bm{z}}_{j}$ and ${\bm{\theta}}_{j}$. Thus we remove the absolute operation in practical implementation for simplicity. Denoting $X=[{\bm{x}}_{1},{\bm{x}}_{2},\dots,{\bm{x}}_{n}]$ as the input of a block, we get the token’s amplitude ${\bm{z}}_{j}$ by a plain channel-FC operation, i.e.,

Phase. Recalling that the phase indicates the current location of token in a period of wave, we discuss different strategy to generate phases as following. The simplest strategy (‘static phase’) is to represent the phase ${\bm{\theta}}_{j}$ of each tokens with fixed parameters, which can be learned in the training process. Though the static phase can distinguish different tokens, it neglects the diversity of different input images either.

To capture the particular attributes for each input respectively, we use an estimation module $\Theta$ to generate the phase information according to input features ${\bm{x}}_{j}$, i.e., $\bm{\theta}_{j}=\Theta({\bm{x}}_{j},W^{\theta})$, where $W^{\theta}$ denotes the learnable parameters. Considering that simplicity is an important characteristic of MLP-like architectures, complex operations are undesirable. Thus we also adopt the simple channel-FC in Eq. 1 as the phase estimation module. The estimation module can also be constructed with other formulations, whose impact on the model performance is empirically investigated in Table 6 of Section 4.4.

Token aggregation. In Eq. 3, the wave-like tokens are represented in the complex domain. To embed it in a general MLP-like architecture, we unfold it with Euler’s formula and represent it with real part and imaginary part, i.e.,

where $W^{t}$, $W^{i}$ are both learnable weights. In the above equation, the phase $\bm{\theta}_{k}$ adjusts dynamically according to the semantic content of input data. Besides the fixed weights, the phases also modulate the aggregating process of different tokens.

3 Wave-MLP Block

A basic unit in the proposed Wave-MLP mainly contains two blocks, channel-mixing MLP and phase-aware token-mixing block (Figure 2). The channel-mixing MLP is stack by two channel-FC layers (Eq. 1) and non-linear activation functions, which extracts features for each token. The token-mixing block composes of the proposed PATM modules, aggregating different tokens by considering both amplitude and phase information.

To be more compatible with computer vision tasks, we preserve the 2D spatial shape of input image by using feature maps with shape $H\times W\times C$, which $H$, $W$, $C$ are the height, width and channel’s number, respectively. This is a successful practice widely used in recent vision transformer architectures (e.g., PVT , Swin-Transformer ). There are two parallel PATM modules, which aggregate spatial information along high and width dimensions, respectively. Similar to , different branches are summed with a re-weighting module. In the traditional MLP-Mixer , each token-FC layer connects all tokens together, whose dimension depends on specific input size. Thus it is not compatible with the dense prediction tasks (e.g., object detection and semantic segmentation) with varying sizes of input images. To address this issue, we use a simple strategy that restricts the FC layers only connect tokens within a local window. The empirical investigation of the window size is shown in Table 7 of Section 4.4. Besides the PATM modules, another channel-FC connecting the input and output directly is also used to preserve the original information. The final output of the block is the summation of these three branches.

The whole model is constructed by stacking phase-aware token-mixing blocks, channel-mixing MLPs, and normalization layers, alternately. To produce hierarchical features, we split the architecture into 4 stages, which reduces the size of feature maps and increases the number of channels stage-wisely. By varying the width and depth of model, we develop 4 models with different parameters and computational costs, denoted as Wave-MLP-T, Wave-MLP-S, Wave-MLP-M, Wave-MLP-B, sequentially. The detailed configures of these models can be found in the supplemental material.

Experiments

In this section, we empirically investigate the proposed Wave-MLP architecture on multiple tasks, containing image classification, object detection and semantic segmentation. Wave-MLP is firstly compared with the existing vision MLPs, vision Transformers and CNNs on ImageNet for image classification. Then it is used as the backbone of two detectors (RetinaNet and Mask R-CNN ) for object detection and instance segmentation on COCO dataset . As for semantic segmentation, the widely used semantic FPN on ADE20K is adopted. Finally, ablation studies are conducted to verify the effectiveness of each component.

Settings. We conduct image classification experiments on the benchmark dataset ImageNet , which contains 1.28M training images and 50k validation images from 1000 classes. For a fair comparison, we use the same training strategy as . Specially, the model is trained for 300 epochs with AdamW optimizer, whose learning rate is initialized as 0.001 and declines with a cosine decay strategy. The batchsize and weight decay are set to 1024 and 0.05, respectively. We use the common data augmentation strategies following , containing Mixup , CutMix and Rand-Augment . At the inference phase, the top-1 accuracy on a single crop is reported. To be compatible with the downstream tasks, we use a local window for token-FC and set the window size to 7 empirically. By adjusting the architecture configures, four models (T, S ,M, B) with different parameters and computational costs are developed. Besides, by replacing the FC layer of phase estimation module with a depth-wise convolution, an more efficient architecture is developed and denoted as Wave-MLP-T*. All the experiments are conducted with PyTorch and MindSpore on NVIDIA V100 GPUs.

Comparison with the existing MLP-like architectures. Table 1 compares the proposed Wave-MLP with existing vision MLP models proposed recently or currently. Throughput is measured on a V100 GPU following Note that AS-MLP reports throughput under the mixed precision mode (mixed FP16 and FP32). For a fair comparison with the existing models, we remeasure it with pure FP32 following . The family of Wave-MLP achieves a better trade-off between the computational cost and accuracies than the existing methods. For example, our Wave-MLP-M model achieves 83.4% top-1 accuracy with only 7.9G FLOPs, which shows a large superiority to ResMLP-B24 (81.0% accuracy with 23.0G FLOPs). Compared with the SOTA MLP architecture CycleMLP , Wave-MLP also achieves higher accuracies with similar parameters and FLOPs, e.g., Wave-MLP-T achieve a accuracy of 80.6%, much higher than that of CycleMLP-B1 with 78.9% accuracy. It shows that equipping each token with the phase information can well capture the relationship between varying tokens and fixed weights to improve the performance of MLP architecture.

Comparison with SOTA models on ImageNet. We further compare the proposed Wave-MLP with typical CNN and transformer architectures on ImageNet in Table 2. Compared with Swin Transformer , our Wave-MLP achieves higher performance with fewer parameters and computational costs. For example, with 4.5G FLOPs, Wave-MLP-S achieves 82.6% top-1 accuracy, wihch significantly superior to Swin-T with 81.3% accuracy. Its trade-off between computational cost and accuracy also suppresses the typical CNN architectures such as RegNetY and ResNet18. The superiority of Wave-MLP implies that the simple MLP architecture has a large potential and modulating the token aggregating process with phase term can exploit it adequately.

2 Object Detection on COCO

Settings. We further investigate the proposed Wave-MLP architecture on the object detection and instance segmentation tasks. The experiments are conducted on the COCO 2017 dataset , which contains 118k training images and 5k validation images. Wave-MLP is used as the backbone and embedded into two prevalent detectors, RetinaNet and Mask R-CNN . For a fair comparison, we follow the training recipe in and train the model with AdamW optimizer for 12 epochs (1$\times$ training scheduler). The batchsize is set to 16 and initial learning to 0.0001. The backbones are initialized with the pre-trained weights on ImageNet while other layers are initialized with Xavier .

Results. Table 3 compares the object detection results with different architectures as the backbone. For both RetinaNet and Mask R-CNN, the proposed Wave-MLP achieves obviously higher performance compared with the existing models. For example, With RetinaNet 1$\times$, Wave-MLP-T achieves 40.4% AP with only 25.3M parameters and 196.3G FlOPs, which is higher than CycleMLP-B1 (38.6 AP) with similar model size by 1.8 AP. When using Mask R-CNN as the detector, the performance improvements are also significant. Compared with Swin-T of 42.2 box AP and 39.1 mask AP with 47.8M parameters and 264.0G FLOPs, our Wave-MLP-S achieves significantly higher performance (44.0 box AP and 40.0 mask AP) with fewer parameters (47.0M) and lower computational cost (250.3G).

3 Semantic Segmentation on ADE20K

Settings. The experiments for the semantic segmentation task are conducted on the challenging ADE20K dataset , which contains 25k images from 150 semantic categories, 20k for training, 2k for validation and 3k for testing. Following , we combine the proposed Wave-MLP architecture with the widely used Semantic FPN approach. With the pre-trained weights on ImageNet, the model is fine-tuned for 40k iterations with AdamW optimizer and the batchsize is set to 32. The initial learning rate is 0.0001 and decays with the polynomial schedule (a power of 0.9). The images are randomly resized and cropped to $512\times 512$ for training and rescaled to have a shorter side of 512 for testing. The FLOPs are tested with 2048$\times$512 input.

Results. The results of different models for semantic segmentation are shown in Table 4. Under different configures of parameters and computational costs, Wave-MLP outperforms the existing models consistently. Compared with the transformer-based model such as PVT, the model show a large superiority, e.g., 4.6% mIoU gap between Wave-MLP-S (44.4% mIoU) and PVT-Tiny (39.8% mIoU). It also suppress the CycleMLP-B2 model with 42.4% mIoU and Swin-S with 41.5%. We infer that modulating the aggregating process of different tokens with the phase term can capture more detailed information and thus enhance the semantic segmentation results.

4 Ablation Studies

For better understanding the proposed method, we investigate the effectiveness of each component via ablation studies. The experiments are conducted on ImageNet with the Wave-MLP-T* model.

The effectiveness of phase information. The phase plays a vital role in aggregating the information of different tokens, whose effectiveness is investigated in Table 5. Without the phase information (‘No phase’), the model’s performance is obviously inferior compared with others, with only 78.8% top-1 accuracy. The proposed ’dynamic phase’ flexibly generates phases and modulate the aggregating process for each input instance, which achieves much better performance (e.g., 80.1% top-1 accuracy).

The formulation of phase estimation module. The phase estimation module generates phases for different inputs, which can be implemented with different formulations. We investigate three simple formulations, depth-wise convolution, channel-FC and identity projection. The identity projection directly copies the input feature instead of estimating the phase, incurring poor performance (i.e., 79.3%). The depth-wise convolution and channel-FC can achieve high accuracy improvement compared with the baseline (e.g., 1.3% and 1.8%), implying they can capture the phase information well for aggregating tokens. Using channel-FC achieves higher performance than the depth-wise convolution, but also increase the computational cost slightly.

The size of window for aggregating tokens. To be compatible with dense prediction tasks (e.g., object detection and semantic segmentation) with varying sizes of input images, we restrict that the token-FC only aggregates features within a local window, and Table 7 investigates the impact of window size. Changing window size from 3 to 7, the top-1 accuracies increase accordingly. ‘All’ denotes that the token-FC connects all the tokens in a layer, which achieves similar performance with window size 7. However, its parameter configure is corrected to the size of input image, and thus is infeasible in dense prediction tasks such as object detection and semantic segmentation.

Visualization. The phase difference between two tokens ($|{\bm{\theta}}_{j}-{\bm{\theta}}_{i}|$) directly affects the aggregating process as analyzed in Section 3.2 (Eq. 4, 5). In order to have an intuitive understanding, we show the cosine value of phase difference of the 3rd and 4th stages in Figure 4. Take the visualized figure of the 1st image and the 4th stage for example, the $7\times 7$ values in the $(i,j)$-th patch denote the phase differences between the $(i,j)$-th token and all the $7\times 7$ tokens. From the figure, we can see that tokens with similar contents tend to have close phases and then enhanced by each other. For example, in the first image, a token describing the ’house’ has a closer phase with another token of the ‘house’ than that of the sky (magnifying parts in the figure). The phase difference of different tokens also varies w.r.t. different input images depending on the image contents.

Conclusion

This paper proposes a Wave-MLP architecture for vision tasks, which takes each token as a wave with both amplitude and phase information. Amplitude is the original real-value feature and the phase modulates relationship between the varying tokens and fixed weights in MLP. With the dynamically produced phase, the tokens are aggregated according to their varying contents from different input images. Extensive experiments show that the proposed Wave-MLP suppresses the existing MLP-like architectures and can also be used as a strong backbone for the dense prediction tasks such as object detection and semantic segmentation. In the future, we will further explore the potential of MLP-like architectures on more diverse tasks.

Acknowledgment. This work is supported by National Natural Science Foundation of China under Grant No.61876007, Australian Research Council under Project DP210101859 and the University of Sydney SOAR Prize.

References

Detailed Architectures

Table 8 shows the detailed specifications of the proposed Wave-MLP architecture. To get hierarchical features, we split the whole model into four stages, and reduce the size of feature map stage-wisely. The Wave-MLP family contains four models with different parameters and computational costs by adjusting the depths and widths of architecture specifications, which are denoted as Wave-MLP-T, Wave-MLP-S, Wave-MLP-M, and Wave-MLP-B, sequentially. From Wave-MLP-T to Wave-MLP-B, the number of parameters varies from 17M to 63M, and FLOPs varies from 2.4G to 10.2G.

More Experiments

For the object detection and instance segmentation tasks on COCO , we further train Mask R-CNN models with 3$\times$ schedule and multi-scale training strategy . The results of different backbone are shown in Table 9. Compared with other backbones, the proposed Wave-MLP achieves much higher performance. For example, our Wave-MLP-T achieves 44.1 box AP and 40.1 mask AP with 25.3M parameters and 196.3G FLOPs, which is significantly superior to the PVT-Tiny model with 39.8 box AP, 37.4 mask AP, 32.9M parameters and 208.1G FLOPs.