AS-MLP: An Axial Shifted MLP Architecture for Vision

1 INTRODUCTION . 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 can not 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 . 2 RELATED WORK . 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 can not 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 . 3 THE AS-MLP ARCHITECTURE . 3.1 OVERALL ARCHITECTURE . 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 ) .

Summary: 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.