CycleMLP: A MLP-like Architecture for Dense Prediction

1 INTRODUCTION . Vision models in computer vision have been long dominated by convolutional neural networks ( CNNs ) ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; He et al. , 2016 ) . Recently , inspired by the successes in Natural Language Processing ( NLP ) field , Transformers ( Vaswani et al. , 2017 ) are adopted into the computer vision community . Built with multi-head self-attention layers , multi-layer perceptrons ( MLPs ) , and skip connections ( He et al. , 2016 ) , Transformers make numerous breakthroughs on visual tasks ( Dosovitskiy et al. , 2020 ; Touvron et al. , 2020 ; Liu et al. , 2021b ) . More recently , ( Tolstikhin et al. , 2021 ; Touvron et al. , 2021a ; Liu et al. , 2021a ) have validated that building models solely on MLPs and skip connections without the self-attention layers can achieve surprisingly promising results on ImageNet ( Deng et al. , 2009 ) classification . Compared with CNNs , these MLP-based models even achieve comparable recognition performance when trained on large scale datasets ( e.g. , ImageNet-21K and JFT-300M ) . Despite promising results on visual recognition tasks , these MLP-like models can not be used in dense prediction tasks ( e.g. , object detection and semantic segmentation ) due to the three challenges : ( 1 ) Current models are composed of blocks with non-hierarchical architectures , which make the model infeasible to provide pyramid and high-resolution feature representations . ( 2 ) Current models can not deal with flexible input scales due to the Spatial FC as shown in Figure 1b . The spatial FC is configured by an image-size related weight1 , which requires input image to have a fixed scale . Thus , this structure typically requires the input image with 1We omit bias here for discussion convenience . a fixed scale during both the training and inference procedure . It contradicts the requirements of dense prediction tasks , which usually adopt a multi-scale training strategy ( Carion et al. , 2020 ) and different input resolutions in training and inference stages ( Lin et al. , 2014 ; Cordts et al. , 2016 ) . ( 3 ) The computational and memory costs of the current MLP models are quadratic to input image sizes for dense prediction tasks ( e.g. , COCO benchmark ( Lin et al. , 2014 ) ) . To address the first challenge , we construct a hierarchical architecture to generate pyramid features . For the second and third issues , we propose a novel variant of fully connected layer , named as Cycle Fully-Connected Layer ( Cycle FC ) , as illustrated in Figure 1c . The Cycle FC is capable of dealing with various image scales and has linear computational complexity to image size . Our Cycle FC is inspired by Channel FC layer illustrated in Figure 1a , which is designed for channel information communication ( Lin et al. , 2013 ; Szegedy et al. , 2015 ; He et al. , 2016 ; Howard et al. , 2017 ) . The main spirit of Channel FC lies in that it can deal with flexible image sizes since it is configured by image-size agnostic weight of Cin and Cout . However , the Channel FC is infeasible to aggregate spatial context information due to its limited receptive field . Our Cycle FC is designed to enjoy Channel FC ’ s merit of taking input with arbitrary resolution and linear computational complexity while enlarging its receptive field for context aggregation . Specifically , Cycle FC samples points in a cyclical style along the channel dimension ( Figure 1c ) . In this way , Cycle FC has the same complexity ( both the number of parameters and FLOPs ) as channel FC while increasing the receptive field simultaneously . To this end , we adopt Cycle FC to replace the Spatial FC for spatial context aggregation ( i.e. , token mixing ) and build a family of MLP-like models for both recognition and dense prediction tasks . The contributions of this paper are as follows : ( 1 ) We propose a new MLP-like operator , Cycle FC , which is computational friendly to cope with flexible input resolutions . ( 2 ) We take the first attempt to build a family of hierarchical MLP-like architectures ( CycleMLP ) based on Cycle FC operator for dense prediction tasks . ( 3 ) Extensive experiments on various tasks ( e.g. , ImageNet classification , COCO object instance detection , and segmentation , and ADE20K semantic segmentation ) demonstrate that CycleMLP outperforms existing MLP-like models and is comparable to and sometimes better than CNNs and Transformers on dense predictions . Related Work . Convolution Neural Networks ( CNNs ) has dominated the visual backbones for several years ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; He et al. , 2016 ) . ( Dosovitskiy 5 10 15 20 25 Model capacity ( MAdds : G ) 72 74 76 78 80 82 Im ag eN et To p- 1 Ac c ( % ) DeiT Mixer S2MLP ResMLP gMLP ViP CycleMLP ( ours ) et al. , 2020 ) introduced the first pure Transformer-based ( Vaswani et al. , 2017 ) model into computer vision and achieved promising performance , especially pre-trained on the large scale JFT dataset . Recently , some works ( Tolstikhin et al. , 2021 ; Touvron et al. , 2021a ; Liu et al. , 2021a ) removed the attention in Transformer and proposed pure MLP-based models . Please see Appendix A for a comprehensive review of the literature on the visual backbones . 2 METHOD . In this section , we introduce CycleMLP models for vision tasks including recognition and dense predictions . To begin with , in Sec . 2.1 we formulate our proposed novel operator , Cycle FC , which serves as a basic component for building CycleMLP models . Then we compare Cycle FC with Channel FC and multi-head attention adopted in recent Transformer-based models ( Dosovitskiy et al. , 2020 ; Touvron et al. , 2020 ; Liu et al. , 2021b ) in Sec . 2.2 . Finally , we present the detailed configurations of CycleMLP models in Sec . 2.3 . 2.1 CYCLE FULLY-CONNECTED LAYER . Notation . We denote an input feature map as X ∈ RH×W×Cin , where H , W denote the height and width of the image and Cin is the number of feature channels . We use subscripts to index the feature map . For example , Xi , j , c is the value of cth channel at the spatial position ( i , j ) and Xi , j , : are values of all channels at the spatial ( i , j ) . The motivation behind Cycle FC is to enlarge receptive field of MLP-like models to cope with downstream dense prediction tasks while maintaining the computational efficiency . As illustrated in Figure 1a , Channel FC applies weighting matrix on X along the channel dimension on fixed position ( i , j ) . However , Cycle FC introduces a receptive field of ( SH , SW ) , where SH and SW are stepsize along with the height and width dimension respectively ( illustrated in Figure 1 ( d ) ) . The basic Cycle FC operator can be formulated as below : CycleFC ( X ) i , j , : = Cin∑ c=0 Xi+δi ( c ) , j+δj ( c ) , c ·W mlp c , : + b ( 1 ) where Wmlp ∈ RCin×Cout and b ∈ RCout are parameters of Cycle FC . δi ( c ) and δj ( c ) are the spatial offset of the two axis on the cth channel , which are defined as below : δi ( c ) = ( c mod SH ) − 1 , δj ( c ) = ( b c SH c mod SW ) − 1 ( 2 ) Examples . We provide several examples ( Figure 1 ( d ) - ( f ) ) to illustrate the stepsize . For the sake of visualization convenience , we set the tensor ’ s W = 1 . Thus , these three examples naturally all have SW = 1 . Figure 1 ( d ) illustrates the offsets along two axis when SH = 3 , that is δj ( c ) ≡ 0 and δi ( c ) = { −1 , 0 , 1 , −1 , 0 , 1 , · · · } when c = 0 , 1 , 2 , · · · , 8 . Figure 1a ( e ) shows that when SH = H , Cycle FC has a global receptive field . Figure 1a ( f ) shows that when SH = 1 , there will be no offset along either axis and thus Cycle FC degrades to Channel FC ( Figure 1a ) . We also provide a more general case where W 6= 1 and SH = 3 , SW = 3 in Figure 7 ( Appendix ) . The offsets δi ( c ) and δj ( c ) enlarge the receptive field of Cycle FC as compared to Channel FC ( Figure 1a ) , which applies weights solely on the same spatial position for all channels . The larger receptive field in return brings improvements on dense prediction tasks like semantic segmentation and object detection as shown in Table 1 . Meanwhile , Cycle FC still maintains computational efficiency and flexibility on input resolution . Both the FLOPs and the number of parameters are linear to the spatial scale which are exactly the same as those of Channel FC . In contrast , although Spatial FC has a global receptive field over the whole spatial space , its computational cost is quadratic to the image scale . Besides , it fails to handle inputs with different resolutions . 2.2 COMPARISON BETWEEN MULTI-HEAD SELF-ATTENTION ( MHSA ) AND CYCLE FC . Inspired by Cordonnier et al . ( 2020 ) , when re-parametried properly , a multi-head self-attention layer withNh heads can be formulated as below , which is similar to a convolution with kernel size √ Nh×√ Nh . ( Please refer to Appendix C for detailed derivation ) MHSA ( X ) i , j , : = ∑ h∈ { 1,2 , ... , Nh } Xi+∆i ( h ) , j+∆j ( h ) , :W mhsa , h + b ( 3 ) where Wmhsa , h ∈ RCin×Cout is the parameter matrix for hth head in MHSA . b ∈ RCout is the bias vector . { ∆i ( h ) , ∆j ( h ) } = { ( 0 , 0 ) , ( 1 , 0 ) , ( −1 , 0 ) , · · · } contains all possible positional shift in convolution with kernel size √ Nh × √ Nh . Further , we stack all Wmhsa , h together and reshape it into Wmhsa ∈ RK×K×Cin×Cout . Then a relationship between Wmlp and Wmhsa can be formulated as follow . Wmlpc , : = W mhsa δi ( c ) +1 , δj ( c ) +1 , c , : ( 4 ) equation 4 shows that only the weights of Wmhsa on spatial shift ( δi ( c ) +1 , δj ( c ) +1 ) are taken into account in Wmlp . This indicates that Cycle FC introduce an inductive bias that the weighting matrix in MHSA should be sparse . Thus Cycle FC inherits the large receptive field introduced in MHSA . The receptive field in Cycle FC is enlarged to ( SH , SW ) , which enables Cycle FC to tackle with downstream dense prediction tasks better . Meanwhile , with the sparsity inductive bias , Cycle FC maintains computational efficiency in MLP-based methods as compared to convolution and multihead self-attention . The parameter size in Cycle FC isCin×Cout while Wmhsa ∈ RK×K×Cin×Cout .

Summary: This paper proposes an improvement to existing MLP style models for image tasks that extends the general MLP competitive work compared to CNN models while attempting to capture some desirable CNN properties such as applicability to varying image sizes.  In addition the method has linear computational complexity as compared to previous MLP work which have quadratic complexity.  The basic idea is a variant of previous channel and spatial FC ideas by having a node's receptive field sample across spatial and channel domains rather than being fixed to either just channel or spatial aggregation.  The overall architecture uses the previous work on patch embeddings and idea of stages with stacked blocks.  Comprehensive experimental results are show improvement over previous MLP models and competitive results to other methods such as CNNs and transformers on multiple tasks including classification and segmentation.