AS-MLP: AN AXIAL SHIFTED MLP ARCHITECTURE
Year: 2021

Abstract:
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 communication. 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, just like designing those of convolution kernels. 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.

heading: 1 INTRODUCTION
text: In the past decade, Convolutional Neural Networks (CNNs) (Krizhevsky et al., 2012; He et al., 2016) have received widespread attention and have become the de-facto standard for computer vision. Furthermore, with the in-depth exploration and research on self-attention, transformer-based architectures have also gradually emerged, and have surpassed CNN-based architectures in natural language processing (e.g., Bert (Devlin et al., 2018)) and vision understanding (e.g., ViT (Dosovitskiy et al., 2021), DeiT (Touvron et al., 2021b)) with amounts of training data. Recently, Tolstikhin et al. (2021) first propose MLP-based architecture, where almost all network parameters are learned from MLP (linear layer) and achieve amazing results, which is comparable with CNN-like models.
Such promising results drive our exploration of MLP-based architecture. In the MLP-Mixer (Tolstikhin et al., 2021), the model obtains the global receptive field through matrix transposition and token-mixing projection such that the long-range dependencies are covered. However, this rarely makes full use of the local information, which is very important in CNN-like architecture (Simonyan & Zisserman, 2015; He et al., 2016) because not all pixels need long-range dependencies, and the local information more focuses on extracting the low-level feature. In the transformer-based architectures, Swin Transformer (Liu et al., 2021b) computes the self-attention in a window (7×7) instead of the global receptive field, which is similar to directly using a convolution layer with a large kernel size (7 × 7) to cover the local receptive field. Some other papers have also already emphasized the advantages of local receptive fields, and introduced local information in the transformer, such as Localvit (Li et al., 2021), NesT (Zhang et al., 2021), etc. Driven by these ideas, we mainly explore the influence of locality on MLP-based architectures.
In order to introduce locality into the MLP-based architecture, one of the simplest and most intuitive ideas is to add a window to the MLP-Mixer, and then perform a token-mixing projection of the local information on the features within the window, just as done in Swin Transformer (Liu et al., 2021b)
and LINMAPPER (Fang et al., 2021). However, for the MLP-based architecture, if we divide the window (e.g., 7 × 7) and perform the token-mixing projection in the window, then the linear layer has the 49× 49 parameters shared between windows1, which greatly limits the model capacity and thus affects the learning of parameters and the final results.
Therefore, a more ideal way to introduce locality is to directly model the relationship between a feature point and its surrounding feature points at any position, without the need to set a fixed window (and window size) in advance. To aggregate the features of different spatial positions in the same position and model their relationships, inspired by (Wu et al., 2018; Lin et al., 2019; Wang et al., 2020; Ho et al., 2019), we propose an axial shift strategy for MLP-based architecture, where we spatially shift features in both horizontal and vertical directions. Such an approach not only aggregates features from different locations, but also makes the feature channel only need to be divided into k groups instead of k2 groups to obtain a receptive field of size k × k with the help of axial operation. After that, a channel-mixing MLP combines these features, enabling the model to obtain local dependencies. It also allows us to design MLP structure as the same as the convolutional kernel, for instance, to design the kernel size and dilation rate.
Based on the axial shift strategy, we design Axial Shifted MLP architecture, named AS-MLP. Our AS-MLP obtains 83.3% Top-1 accuracy with 88M parameters and 15.2 GFLOPs in the ImageNet1K dataset without any extra training data. Such a simple yet effective method outperforms all MLP-based architectures and achieves competitive performance compared to the transformer-based architectures. It is also worth noting that the model weights in MLP-Mixer trained with fixed image size cannot be adapted to downstream tasks with various input sizes because the token-mixing MLP has a fixed dimension. On the contrary, the AS-MLP architecture can be transferred to downstream tasks (e.g., object detection) due to the design of axial shift. As far as we know, it is also the first work to apply MLP-based architecture to the downstream task. With the pre-trained model in the ImageNet-1K dataset, 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.

heading: 2 RELATED WORK
text: CNN-based Architectures. Since AlexNet (Krizhevsky et al., 2012) won the ImageNet competition in 2012, the CNN-based architectures have gradually been utilized to automatically extract image features instead of hand-crafted features. Subsequently, the VGG network (Simonyan & Zisserman, 2015) is proposed, which purely uses a series of 3 × 3 convolution and fully connected layers, and obtains outstanding performance in image classification. Furthermore, ResNet (He et al., 2016) is proposed, which utilizes the residual connection to transfer features in different layers, thereby alleviating the problem of gradient vanishing and obtaining superior performance. After that, the residual module becomes an important component of the network design and is also employed in subsequent transformer-based architectures and MLP-based architectures. Some papers have made further improvements to the convolution operation in CNN-based architecture, such as dilated convolution (Yu & Koltun, 2016) and deformable convolution (Dai et al., 2017). EfficientNet (Tan & Le, 2019; 2021) introduces neural architecture search into CNN to search for a suitable structure. These architectures build the CNN family and are used extensively in computer vision tasks.
Transformer-based Architectures. Transformer is first proposed in (Vaswani et al., 2017), where the attention mechanism is utilized to model the relationship between features from the different spatial positions. Subsequently, the popularity of BERT (Devlin et al., 2018) in NLP also promotes the research on transformer in the field of vision. ViT (Dosovitskiy et al., 2021) uses a pure transformer framework to extract visual features, where an image is divided into 16× 16 patches and the convolution layer is completely abandoned. It shows that the transformer-based architecture can perform well in large-scale datasets (e.g., JFT-300M). After that, DeiT (Touvron et al., 2021b) carefully designs training strategies and data augmentation to further improve performance on the small datasets (e.g., ImageNet-1K). DeepViT (Zhou et al., 2021) and CaiT (Touvron et al., 2021c) consider the optimization problem when the network deepens, and train a deeper transformer network. CrossViT (Chen et al., 2021a) combines the local patch and global patch using two vision transformers. CPVT (Chu et al., 2021b) uses a conditional position encoding to effectively encode the spatial positions of
1If linear layer is not shared between windows, the model weights trained with fixed image size cannot be adapted to downstream tasks with various input sizes because unfixed input sizes cause a mismatch in the number of windows.
patches. LeViT (Graham et al., 2021) improves ViT from many aspects, including the convolution embedding, extra non-linear projection and batch normalization, etc. Transformer-LS (Zhu et al., 2021) proposes a long-range attention and a short-term attention to model long sequences for both language and vision tasks. Some papers also design hierarchical backbone to extract spatial features at different scales, such as PVT (Wang et al., 2021), Swin Transformer (Liu et al., 2021b), Twins (Chu et al., 2021a) and NesT (Zhang et al., 2021), which can be applied to downstream tasks (e.g., object detection and semantic segmentation).
MLP-based Architectures. MLP-Mixer (Tolstikhin et al., 2021) has designed a very concise framework that utilizes matrix transposition and MLP to transmit information between spatial features. Resort to MLP, skip connection between layers and normalization layer, MLP-Mixer obtains promising experimental results. The concurrent work FF (Melas-Kyriazi, 2021) also applies a similar network architecture and reaches similar conclusions. Such experimental results are surprising, which shows that the MLP-based architecture also achieves comparable performance with CNN-based architectures and transformer-based architectures. Subsequently, Res-MLP (Touvron et al., 2021a) is proposed, which also obtains impressive performance with residual MLP only trained on ImageNet1K. gMLP (Liu et al., 2021a) and EA (Guo et al., 2021) introduce Spatial Gating Unit (SGU) and the external attention to improve the performance of the pure MLP-based architecture, respectively. Recently, Container (Gao et al., 2021) proposes a general network that unifies convolution, transformer, and MLP-Mixer. The closest concurrent methods to our work are S2-MLP (Yu et al., 2021) and ViP (Hou et al., 2021). S2-MLP uses spatial-shift MLP for feature exchange. ViP proposes a Permute-MLP layer for spatial information encoding to capture long-range dependencies. Different from their work, we focus on capturing the local dependencies with axially shifting features in the spatial dimension, which obtains better performance and can be applied to the downstream tasks. More recently, the concurrent work, CycleMLP (Chen et al., 2021b) and S2-MLPv2 (Yu et al., 2021) are proposed. S2-MLPv2 improves S2-MLP and CycleMLP designs Cycle Fully-Connected Layer (Cycle FC) to obtain a larger receptive field than Channel FC. Our AS-MLP obtains comparable performance with these architectures.

heading: 3 THE AS-MLP ARCHITECTURE

heading: 3.1 OVERALL ARCHITECTURE
text: Figure 1 shows our Axial Shifted MLP (AS-MLP) architecture2. Given an RGB image I ∈ R3×H×W , where H and W are the height and width of the image, respectively. AS-MLP first performs the patch partition operation, which splits the original image into multiple patch tokens with the patch size of 4 × 4, thus the combination of all tokens has the size of 48 × H4 × W 4 . ASMLP has four stages in total and there are different numbers of AS-MLP blocks in different stages. Figure 1 only shows the tiny version of AS-MLP, and other variants will be discussed in Sec. 3.4. All the tokens in the previous step will go through these four stages, and the final output feature will be used for image classification. In Stage 1, a linear embedding and the AS-MLP blocks are adopted for each token. The output has the dimension of C × H4 × W 4 , where C is the number of channels. Stage 2 first performs patch merging on the features outputted from the previous stage, which groups the neighbor 2×2 patches to obtain a feature with the size of 4C× H8 × W 8 and then a linear layer is adopted to warp the feature size to 2C × H8 × W 8 , followed by the cascaded AS-MLP blocks. Stage 3 and Stage 4 have similar structures to Stage 2, and the hierarchical representations will be generated in these stages.
2The design of figure follows Swin Transformer (Liu et al., 2021b).

heading: 3.2 AS-MLP BLOCK
text: The core operation of AS-MLP architecture is the AS-MLP block, which is illustrated in Figure 2a. It mainly consists of the Norm layer, Axial Shift operation, MLP, and residual connection. In the Axial Shift operation, we utilize the channel projection, vertical shift, and horizontal shift to extract features, where the channel projection maps the feature with a linear layer. Vertical shift and horizontal shift are responsible for the feature translation along the spatial directions.
As shown in Figure 2b, we take the horizontal shift as an example. The input has the dimension of C×h×w, for convenience, we omit h and assume C = 3, w = 5 in this figure. When the shift size is 3, the input feature is split into three parts and each part is shifted by {-1, 0, 1} units along the horizontal direction, respectively. In this operation, zero padding is performed (indicated by gray blocks), and we also discuss the experimental results of using other padding methods in Sec. 4. After that, the features in the dashed box will be taken out and used for the next channel projection. The same operation is also performed in the vertical shift. In the process of horizontal shift and vertical shift, since the feature performs different shift units in both directions, when the features are recombined (the feature in the dashed box), the information from different spatial positions can be combined together. In the next channel projection operation, information from different spatial locations can fully flow and interact. The code of AS-MLP block is listed in Alg. 1.
Complexity. In the transformer-based architecture, the multi-head self-attention (MSA) is usually adopted, where the attention between tokens is computed. Swin Transformer (Liu et al., 2021b) introduces a window to partition the image and propose window multi-head self-attention (W-MSA), which only considers the computation within this window. It significantly reduces the computation complexity. However, in the AS-MLP block, without the concept of the window, we only Axially Shift (AS) the feature from the previous layer, which does not require any multiplication and addition operations. Further, the time cost of Axial Shift is very low and almost irrelevant to the shift size. Given a feature map (is usually named patches in transformer) with the dimension of C × h × w, each Axial shift operation in Figure 2a only has four channel projection operations, which has the computation complexity 4hwC2. If the window size in Swin Transformer (Liu et al., 2021b) is M , the complexities of MSA, W-MSA and AS are as follows: Ω(MSA) = 4hwC 2 + 2(hw)2C, Ω(W-MSA) = 4hwC2 + 2M2hwC, Ω(AS) = 4hwC2. (1)
Therefore, the AS-MLP architecture has slightly less complexity than Swin Transformer. The specific complexity calculation of each layer is shown in Appendix A.2.

heading: 3.3 COMPARISONS BETWEEN AS-MLP, CONVOLUTION, TRANSFORMER AND MLP-MIXER
text: In this section, we compare AS-MLP with the recent distinct building blocks used in computer vision, e.g., the standard convolution, Swin Transformer, and MLP-Mixer. Although these modules are explored in completely different routes, from the perspective of calculation, they are all based on a given output location point, and the output depends on the weighted sum of different sampling location features (multiplication and addition operation). These sampling location features include local dependencies (e.g., convolution) and long-range dependencies (e.g., MLP-Mixer). Figure 3 shows the main differences of these modules in the sampling location. Given an input feature map X ∈ RH×W×C , the outputs Yi,j with different operations in position (i, j) are as follows:
Algorithm 1 Code of AS-MLP Block in a PyTorch-like style.
!"#$"%&'("# )*(#
+,-#./",01, 23452(61,
7)5234
8. 9 :; < 9 =>
7)5234
8. 9 ?; < 9 =>
7)5234
8. 9 :; < 9 @>
Convolution. For convolution operation, a sliding kernel with the shape of k × k (receptive field regionR) is performed on X to obtain the output Y convi,j :
Y convi,j = ∑
(m,n)∈R
Xi+m,j+n,:W conv m,n,:, (2)
where W conv ∈ Rk×k×C is the learnable weight. h and w are the height and width of X , respectively. As shown in Figure 3, the convolution operation has a local receptive field, thus it is more suitable at extracting features with the local dependencies.
Swin Transformer. Swin Transformer introduces a window into the transformer-based architecture to cover the local attention. The input X from a window is embedded to obtain Q,K, V matrix, and the output Yswin is the attention combination of features within the window:
Y swini,j = Softmax(Q(Xi,j)K(X) T / √ d)V (X). (3)
The introduction of locality further improves the performance of the transformer-based architecture and reduces the computational complexity.
MLP-Mixer. MLP-Mixer abandons the attention operation. It first transposes the input X , and then a token-mixing MLP is appended to obtain the output Y mixeri,j :
Y mixeri,j = (X TWmixeriW+j) T , (4)
where Wmixer ∈ Rhw×hw is the learnable weight in token-mixing MLP. MLP-Mixer perceives the global information only with matrix transposition and MLP.
AS-MLP. AS-MLP axially shifts the feature map as shown in Figure 2b. Given the input X , shift size s and dilation rate d, X is first divided into s splits in the horizontal and vertical direction. After the axial shift in Figure 2b, the output Y asi,j is:
Y asi,j = C∑ c=0 Xi+b cdC/se c−b s 2 c·d,j,cW as-h c + C∑ c=0 Xi,j+b cdC/se c−b s 2 c·d,cW as-v c (5)
where W as-h,W as-v ∈ RC are the learnable weights of channel projection in the horizontal and vertical directions (here we omit activation function and bias). Unlike MLP-Mixer, we pay more attention to the local dependencies through axial shift of features and channel projection. Such an operation is closely related to Shift (Wu et al., 2018) and TSM (Lin et al., 2019). However, our method has the following characteristics: i) we use axial shift in the horizontal and vertical directions, which focuses more on the information exchange in two directions; ii) the proposed network is built upon Swin transformer and pure MLP-based architecture, where only linear layer is used and BN is replaced by LN; iii) as will be shown in Sec. 4, our network achieves superior performance, which shows the effectiveness of our method. Although such an operation can be implemented in the original shift (Wu et al., 2018), the feature channel needs to split into k2 groups to achieve a receptive field of size k × k. However, the axial operation makes the feature channel only need to be divided into k groups instead of k2 groups, which reduces the complexity.

heading: 3.4 VARIANTS OF AS-MLP ARCHITECTURE
text: Figure 1 only shows the tiny version of our AS-MLP architecture. Following DeiT (Touvron et al., 2021b) and Swin Transformer (Liu et al., 2021b), we also stack different number of AS-MLP blocks (the number of blocks in four stages) and expand the channel dimension (C in Figure 1) to obtain variants of the AS-MLP architecture of different model sizes, which are AS-MLP-Tiny (AS-MLPT), AS-MLP-Small (AS-MLP-S) and AS-MLP-Base (AS-MLP-B), respectively. The specific configuration is as follows:
• AS-MLP-T: C = 96, the number of blocks in four stages = {2, 2, 6, 2}; • AS-MLP-S: C = 96, the number of blocks in four stages = {2, 2, 18, 2}; • AS-MLP-B: C = 128, the number of blocks in four stages = {2, 2, 18, 2}.
The detailed configurations can be found in Appendix A.1. Table 1 in Sec. 4 shows Top-1 accuracy, model size (Params), computation complexity (FLOPs) and throughput of different variants.

heading: 4 EXPERIMENTS

heading: 4.1 IMAGE CLASSIFICATION ON THE IMAGENET-1K DATASET
text: Settings. To evaluate the effectiveness of our AS-MLP, we conduct experiments of the image classification on the ImageNet-1K benchmark, which is collected in (Deng et al., 2009). It contains 1.28M training images and 20K validation images from a total of 1000 classes. We report the experimental results with single-crop Top-1 accuracy. We use an initial learning rate of 0.001 with cosine decay and 20 epochs of linear warm-up. The AdamW (Loshchilov & Hutter, 2019) optimizer is employed to train the whole model for 300 epochs with a batch size of 1024. Following the training strategy of Swin Transformer (Liu et al., 2021b), we also use label smoothing (Szegedy et al., 2016) with a smooth ratio of 0.1 and DropPath (Huang et al., 2016) strategy.
Results. All image classification results are shown in Table 1. We divide all network architectures into CNN-based, Transformer-based and MLP-based architectures. The input resolution is 224 × 224. Our proposed AS-MLP outperforms other MLP-based architectures when keeping similar parameters and FLOPs. e.g., AS-MLP-S obtains higher top-1 accuracy (83.1%) with less parameters than Mixer-B/16 (Tolstikhin et al., 2021) (76.4%) and ViP-Medium/7 (Hou et al., 2021) (82.7%). Furthermore, it achieves competitive performance compared with transformer-based architectures, e.g., AS-MLP-B (83.3%) vs. Swin-B (Liu et al., 2021b) (83.3%), which shows the effectiveness of our AS-MLP architecture.
Results of Mobile Setting. In addition to standard experiments, we also compare the results of AS-MLP in the mobile setting, which is shown in Table 2. We build the Swin (mobile) model and AS-MLP (mobile) model with similar parameters (about
10M). The specific network details can be found in Appendix A.3. The experimental results show that our model significantly exceeds Swin Transformer (Liu et al., 2021b) in the mobile setting (76.05% vs. 75.11%).

heading: 4.2 THE CHOICE AND IMPACT OF AS-MLP BLOCK
text: The core component in the AS-MLP block is the axial shift. We perform experiments to analyze the choices of different configurations3 of the AS-MLP block and its connection types and the impact of AS-MLP block. All ablations are conducted based on the AS-MLP-T, as shown in the setting of Sec. 3.4.
Different Configurations of AS-MLP Block. In order to encourage the information flow from different channels in the spatial dimension, the features from the horizontal shift and the vertical shift are aggregated together in Figure 2. We evaluate the influence of different configurations of AS-MLP block, including shift size, padding method, and dilation rate, which are similar to the configuration of a convolution kernel. All experiments of different configurations are shown in Table 3a. We have three findings as follows: i) ‘Zero padding’ is more suitable for the design of AS-MLP block than other padding methods4; ii) increasing the dilation rate slightly reduces the performance of AS-MLP, which is consistent with CNN-based architecture. Dilation is usually used for semantic segmentation rather than image classification; iii) when expanding the shift size, the accuracy will increase first and then decrease. A possible reason is that the receptive field is enlarged (shift size = 5 or 7) such that AS-MLP pays attention to the global dependencies, but when shift size is 9, the network pays too much attention to the global dependencies, thus neglecting the extraction of local features, which leads to lower accuracy. Therefore, we use the configuration (shift size = 5, zero padding, dilation rate = 1) in all experiments, including object detection and semantic segmentation.
Connection Type. We also compare the different connection types of AS-MLP block, such as serial connection and parallel connection, and the results are shown in Table 3b. Parallel connection
3There exist amounts of configurations of the block, we cannot enumerate all configurations. Here we only list some general configurations instead.
4Since we use circular shift, thus ‘No padding’ and ‘Circular padding’ are equivalent.
consistently outperforms serial connection in terms of different shift sizes, which shows the effectiveness of the parallel connection. When the shift size is 1, the serial connection is better but it is not representative because only channel-mixing MLP is used.
The Impact of AS-MLP Block. We also evaluate the impact of AS-MLP block in Table 4. Here we design five baselines: i) Global-MLP; ii) AxialMLP; iii) Window-MLP; iv) shift size (5, 1); v) shift size (1, 5). The first three baselines are designed from the perspective of how to use MLP for feature fusion at different positions, and the latter
two are designed from the perspective of the axial shift in a single direction. The specific settings are listed in Appendix A.5. The results in Table 4 show that our AS-MLP block with shift size (5, 5) outperforms other baselines.

heading: 4.3 OBJECT DETECTION ON COCO
text: The experimental setting is listed in Appendix A.4, and the results are shown in Table 5. It is worth noting that we do not compare our method with MLP-Mixer (Tolstikhin et al., 2021) because it uses a fixed spatial dimension for token-mixing MLP, which cannot be transferred to the object detection. As far as we know, we are the first work to apply MLP-based architecture to object detection. Our AS-MLP achieves comparable performance with Swin Transformer in the similar resource limitation. To be specific, Cascade Mask R-CNN + Swin-B achieves 51.9 APb with 145M parameters and 982 GFLOPs, and Cascade Mask R-CNN + AS-MLP-B obtains 51.5 APb with 145M parameters and 961 GFLOPs. The visualizations of object detection are shown in Appendix C.

heading: 4.4 SEMANTIC SEGMENTATION ON ADE20K
text: The experimental setting is listed in Appendix A.4 and Table 6 shows the performance of our ASMLP on the ADE20K dataset. Note that we are also the first to apply the MLP-based architecture to semantic segmentation. The results are impressive. With slightly fewer FLOPs, AS-MLP-T achieves better result than Swin-T (46.5 vs. 45.8 MS mIoU). For the large model, UperNet + SwinB has 49.7 MS mIoU with 121M parameters and 1188 GFLOPs, and UperNet + AS-MLP-B has 49.5 MS mIoU with 121M parameters and 1166 GFLOPs, which also shows the effectiveness of our ASMLP architecture in processing the downstream task. The visualizations of semantic segmentation are shown in Appendix C.
4.5 VISUALIZATION
We visualize the heatmap of learned features from Swin Transformer and AS-MLP in Figure 4, where the first column shows the image from ImageNet, and the second column shows the activation heatmap of the last layer of Swin transformer (Swin-B). The third, fourth, and fifth columns respectively indicate the response after the horizontal shift (AS-MLP (h)), the vertical shift (AS-MLP (v)) and the combination of both in the last layer of AS-MLP (AS-MLPB). From Figure 4, one can see that i) AS-MLP can better focus on object regions compared to Swin transformer; ii) AS-MLP (h) can better focus on the vertical part of objects (as shown in the second row) because it shifts feature in the horizontal direction. It is more reasonable because the shift in the horizontal direction can cover the edge of the vertical part, which is more helpful for recognizing the object. Sim-
ilarly, AS-MLP (v) can better focus on the horizontal part of objects (as shown in the fourth row).

heading: 5 CONCLUSION AND FUTURE WORK
text: In this paper, we propose an axial shifted MLP architecture, named AS-MLP, for vision. Compared with MLP-Mixer, we pay more attention to the local features extraction and make full use of the channel interaction between different spatial positions through a simple feature axial shift. With the proposed AS-MLP, we further improve the performance of MLP-based architecture and the experimental results are impressive. 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 method outperforms all MLP-based architectures and achieves competitive performance compared to the transformerbased architectures even with slightly lower FLOPs. We are also the first work to apply AS-MLP to the downstream tasks (e.g., object detection and semantic segmentation). The results are also competitive or even better compared to transformer-based architectures, which shows the ability of MLP-based architectures in handling downstream tasks.
For future work, we will investigate the effectiveness of AS-MLP in natural language processing, and further explore the performance of AS-MLP on downstream tasks.

heading: A THE ARCHITECTURE DETAILS

heading: A.1 THE DETAILED CONFIGURATIONS OF DIFFERENT ARCHITECTURES
text: We show the detailed configurations of different architectures in Table 7, where we assume the size of the input image is 224 × 224. The second column shows the output size of the image after each stage. Following Swin Transformer (Liu et al., 2021b), we use “Concat n × n” to indicate a concatenation of n × n neighboring features in a patch. “shift size (5, 5)” means that the shift size in the horizontal and vertical directions is 5.

heading: A.2 THE COMPUTATIONAL COMPLEXITY OF AS-MLP ARCHITECTURE
text: In this section, we show the specific computational complexity in each layer of AS-MLP architecture. The symbol definition is first given as follows. An input image: I ∈ R3×H×W ; patch size (p, p); the number of blocks in four stages: {n1, n2, n3, n4}; Channel dimension C; MLP ratio: r. The specific computational complexity is shown in Table 8, where only convolution operation is computed.

heading: A.3 THE NETWORK DETAILS IN THE MOBILE SETTING
text: In addition to AS-MLP-T, AS-MLP-S, and AS-MLP-B, we also design AS-MLP in the mobile setting. For a fair comparison, we modify the Swin Transformer correspondingly to adopt to the mobile setting. The configurations are as follow:
• Swin (mobile): C = 64, the number of blocks in four stages = {2, 2, 2, 2}, the number of heads = {2, 4, 8, 16};
• AS-MLP (mobile): C = 64, the number of blocks in four stages = {2, 2, 2, 2};

heading: A.4 THE SETTINGS OF OBJECT DETECTION AND SEMANTIC SEGMENTATION
text: Object Detection on COCO. For the object detection and instance segmentation, we employ mmdetection (Chen et al., 2019) as the framework and COCO (Lin et al., 2014) as the evaluation dataset, which consists of 118K training data and 5K validation data. We compare the performance of our AS-MLP with other backbones on COCO. Following Swin Transformer (Liu et al., 2021b), we consider two typical object detection frameworks: Mask R-CNN (He et al., 2017) and Cascade R-CNN (Cai & Vasconcelos, 2018). The training strategies are as follows: optimizer (AdamW), learning rate (0.0001), weight decay (0.05), and batch size (2 imgs/per GPU×8 GPUs). We utilize the typical multi-scale training strategy (Carion et al., 2020; Sun et al., 2021) (the shorter side is between 480 and 800 and the longer side is at most 1333). All backbones are initialized with weights pre-trained on ImageNet-1K and all models are trained with 3x schedule (36 epochs).
Semantic Segmentation on ADE20K. Following Swin Transformer (Liu et al., 2021b), we conduct experiments of AS-MLP on the challenging semantic segmentation dataset, ADE20K, which contains 20,210 training images and 2,000 validation images. We utilize UperNet (Xiao et al., 2018) and AS-MLP backbone as our main experimental results. The framework is based on mmsegmentation (Contributors, 2020). The training strategies are as follows: optimizer (AdamW), learning rate (6 × 10−5), weight decay (0.01), and batch size (2 imgs/per GPU×8 GPUs). We utilize random horizontal flipping, random re-scaling within ratio range [0.5, 2.0] and random photometric distortion as data augmentation. The input image resolution is 512× 512, the stochastic depth ratio is set as 0.3 and all models are initialized with weights pre-trained on ImageNet-1K and are trained 160K iterations.

heading: A.5 BASELINES
text: We list the specific configurations of baselines in Sec. 4.2 as follows.
• Global-MLP: following MLP-Mixer (Tolstikhin et al., 2021), we use global MLP (tokenmixing MLP) along with full spatial size instead of AS-MLP block in our architecture configurations. For Global-MLP, the model weights trained with fixed image size cannot be adapted to downstream tasks with various input sizes.
• Axial-MLP: built upon Global-MLP, Axial-MLP employs two axial MLPs along with horizontal and vertical directions instead of global MLP. Similar to Global-MLP, the model weights trained with fixed image size cannot be adapted to downstream tasks with various input sizes.
• Window-MLP: as stated in Sec 1, we set fixed window (7× 7) in our architecture configurations and perform MLP operations within the window.
• Shift size (5, 1): horizontal shift is 5 and vertical shift is 1 in AS-MLP block.
• Shift size (1, 5): horizontal shift is 1 and vertical shift is 5 in AS-MLP block.

heading: A.6 DIFFERENCES BETWEEN AS-MLP AND TSM
text: We elaborate the differences between AS-MLP and TSM as follows.
• TSM performs a shift in the temporal dimension and, as stated in (Lin et al., 2019), they target the temporal dimension for efficient video understanding. However, we explore the
shift from a spatial perspective, for the more general tasks, such as image classification, object detection, and segmentation.
• TSM shows that shifting too many channels in a network will significantly hurt the spatial modeling ability and result in performance degradation. However, Table reftable: ablation(a) shows that as the shift size increases, the channel needs to be divided into more parts, but the performance does not decrease significantly. This suggests that the argument of TSM is not obvious in the pure MLP architecture.
• Our motivation is quite different. TSM is designed to be more effective and efficient on video. Our motivation is derived from Swin Transformer’s exploration of the local receptive field of the transformer. We use shift to explore the local receptive field of MLP. Also, ASMLP is the first MLP-based architecture for object detection and semantic segmentation with the help of such a method. Furthermore, we use axial shift to reduce the complexity of the shift split.

heading: B MORE EXPERIMENTS

heading: B.1 MORE EXPERIMENTAL RESULTS ON COCO
text: Table 5 lists the object detection and instance segmentation results of different backbones with 3x schedule (36 epochs). For a complete comparison, we also conduct experiments with 1x schedule (12 epochs). The results are shown in Table 9. Our AS-MLP-T outperforms Swin-T (Liu et al., 2021b) under the Mask R-CNN (44.0 vs. 43.7 APb) and Cascade Mask R-CNN (48.4 vs. 48.1 APb) frameworks.

heading: B.2 THE COMPLETE CLASSIFICATION ACCURACY AND THROUGHPUT COMPARISON
text: Table 10 shows the complete accuracy comparison with other state-of-the-art architectures on the ImageNet-1K dataset. In addition, Table 1 shows the throughput results of the AS-MLP architecture measured with the batch size 64 on a single V100 GPU (32GB). In order to make a fair comparison with other papers, we also conduct a thorough evaluation of throughput. The results are shown in Figure 5, where we list the throughputs when the batch size is 1, 4, 8, 16, 32, 64, 128, respectively.

heading: B.3 EVALUATION ACCURACY
text: In Figure 6, we visualize the evaluation accuracy of AS-MLP and Swin transformer on the ImageNet, COCO and ADE20K datasets during training. For image classification on ImageNet, AS-MLP-T keeps pace with Swin-T in each epoch and they finally converge to the similar accuracy (81.3 vs. 81.3). For object detection and semantic segmentation on COCO and ADE20K, we can see that AS-MLP-T achieves better performance than Swin-T in the early stage, and keeps winning during the training process.

heading: B.4 COMPARISONS TO SHIFTRESNET
text: In this section, we compare the performance with ShiftResNet (Wu et al., 2018) in Table 11. The results of ShiftResNet50 with different configurations are from Table 3 of ShiftResNet paper (Wu et al., 2018). Our AS-MLP (mobile) achieves better accuracy with less parameters.

heading: C THE VISUALIZATION OF RESULTS ON COCO AND ADE20K
text: We visualize the object detection and instance segmentation results on the COCO dataset in Figure 7, where the Cascade Mask R-CNN model with the AS-MLP-T backbone is used. We also visualize the semantic segmentation results on the ADE20K dataset in Figure 8, where we utilize the UperNet model with the AS-MLP-T backbone. The object can be detected and segmented correctly.

Review summaries:
The paper proposes a new architecture for computer vision that is inspired by (a) the Swin Transformer, (b) MLP-Mixer (and colleagues) and (c) CNN-like local context via shifts (like Shift, TSM, ViP, S2-MLP). The architecture is based on the Swin Transformer, removes the windowed-attention, and then adds "local shifts" of channels to introduce local context via MLP. The architecture is applied to ImageNet-1k classification, COCO detection, and ADE20k segmentation with good results relative to model size and performance. 

This paper proposes to use the shift operation (Wu et al. CVPR 2018) in an axial manner for MLP-mixer architectures. The proposed method performs much better than previous MLP-based methods on ImageNet-1K and on par with Swin-transformer.