MLP-Mixer: An all-MLP Architecture for Vision

1 Introduction . As the history of computer vision demonstrates , the availability of larger datasets coupled with increased computational capacity often leads to a paradigm shift . While Convolutional Neural Networks ( CNNs ) have been the de-facto standard for computer vision , recently Vision Transformers [ 14 ] ( ViT ) , an alternative based on self-attention layers , attained state-of-the-art performance . ViT continues the long-lasting trend of removing hand-crafted visual features and inductive biases from models and relies further on learning from raw data . We propose the MLP-Mixer architecture ( or “ Mixer ” for short ) , a competitive but conceptually and technically simple alternative , that does not use convolutions or self-attention . Instead , Mixer ’ s architecture is based entirely on multi-layer perceptrons ( MLPs ) that are repeatedly applied across either spatial locations or feature channels . Mixer relies only on basic matrix multiplication routines , changes to data layout ( reshapes and transpositions ) , and scalar nonlinearities . Patches Mixer Layer MLP Fully-connected GELU Fully-connected Patches C ha nn el s N x ( Mixer Layer ) Global Average Pooling Per-patch Fully-connected Channels Class L ay er N or m Fully-connected MLP 1 MLP 1 MLP 1 MLP 1 MLP 2 MLP 2 MLP 2 MLP 2 MLP 2 MLP 2 Skip-connections Skip-connections T T L ay er N or mFigure 1 : MLP-Mixer consists of per-patch linear embeddings , Mixer layers , and a classifier head . Mixer layers contain one token-mixing MLP and one channel-mixing MLP , each consisting of two fully-connected layers and a GELU nonlinearity . Other components include : skip-connections , dropout , and layer norm on the channels . they operate on each token independently and take individual rows of the table as inputs . The token-mixing MLPs allow communication between different spatial locations ( tokens ) ; they operate on each channel independently and take individual columns of the table as inputs . These two types of layers are interleaved to enable interaction of both input dimensions . In the extreme case , our architecture can be seen as a very special CNN , which uses 1×1 convolutions for channel mixing , and single-channel depth-wise convolutions of a full receptive field and parameter sharing for token mixing . However , the converse is not true as typical CNNs are not special cases of Mixer . Furthermore , a convolution is more complex than the plain matrix multiplication in MLPs as it requires an additional costly reduction to matrix multiplication and/or specialized implementation . Despite its simplicity , Mixer attains competitive results . When pre-trained on large datasets ( i.e. , ∼100M images ) , it reaches near state-of-the-art performance , previously claimed by CNNs and Transformers , in terms of the accuracy/cost trade-off . This includes 87.94 % top-1 validation accuracy on ILSVRC2012 “ ImageNet ” [ 13 ] . When pre-trained on data of more modest scale ( i.e. , ∼1– 10M images ) , coupled with modern regularization techniques [ 49 , 54 ] , Mixer also achieves strong performance . However , similar to ViT , it falls slightly short of specialized CNN architectures . 2 Mixer Architecture . Modern deep vision architectures consist of layers that mix features ( i ) at a given spatial location , ( ii ) between different spatial locations , or both at once . In CNNs , ( ii ) is implemented with N ×N convolutions ( for N > 1 ) and pooling . Neurons in deeper layers have a larger receptive field [ 1 , 29 ] . At the same time , 1×1 convolutions also perform ( i ) , and larger kernels perform both ( i ) and ( ii ) . In Vision Transformers and other attention-based architectures , self-attention layers allow both ( i ) and ( ii ) and the MLP-blocks perform ( i ) . The idea behind the Mixer architecture is to clearly separate the per-location ( channel-mixing ) operations ( i ) and cross-location ( token-mixing ) operations ( ii ) . Both operations are implemented with MLPs . Figure 1 summarizes the architecture . Mixer takes as input a sequence of S non-overlapping image patches , each one projected to a desired hidden dimension C. This results in a two-dimensional real-valued input table , X ∈ RS×C . If the original input image has resolution ( H , W ) , and each patch has resolution ( P , P ) , then the number of patches is S = HW/P 2 . All patches are linearly projected with the same projection matrix . Mixer consists of multiple layers of identical size , and each layer consists of two MLP blocks . The first one is the token-mixing MLP : it acts on columns of X ( i.e . it is applied to a transposed input table X > ) , maps RS 7→ RS , and is shared across all columns . The second one is the channel-mixing MLP : it acts on rows of X , maps RC 7→ RC , and is shared across all rows . Each MLP block contains two fully-connected layers and a nonlinearity applied independently to each row of its input data tensor . Mixer layers can be written as follows ( omitting layer indices ) : U∗ , i = X∗ , i +W2 σ ( W1 LayerNorm ( X ) ∗ , i ) , for i = 1 . . . C , ( 1 ) Yj , ∗ = Uj , ∗ +W4 σ ( W3 LayerNorm ( U ) j , ∗ ) , for j = 1 . . . S. Here σ is an element-wise nonlinearity ( GELU [ 16 ] ) . DS and DC are tunable hidden widths in the token-mixing and channel-mixing MLPs , respectively . Note that DS is selected independently of the number of input patches . Therefore , the computational complexity of the network is linear in the number of input patches , unlike ViT whose complexity is quadratic . Since DC is independent of the patch size , the overall complexity is linear in the number of pixels in the image , as for a typical CNN . As mentioned above , the same channel-mixing MLP ( token-mixing MLP ) is applied to every row ( column ) of X. Tying the parameters of the channel-mixing MLPs ( within each layer ) is a natural choice—it provides positional invariance , a prominent feature of convolutions . However , tying parameters across channels is much less common . For example , separable convolutions [ 9 , 40 ] , used in some CNNs , apply convolutions to each channel independently of the other channels . However , in separable convolutions , a different convolutional kernel is applied to each channel unlike the token-mixing MLPs in Mixer that share the same kernel ( of full receptive field ) for all of the channels . The parameter tying prevents the architecture from growing too fast when increasing the hidden dimension C or the sequence length S and leads to significant memory savings . Surprisingly , this choice does not affect the empirical performance , see Supplementary A.1 . Each layer in Mixer ( except for the initial patch projection layer ) takes an input of the same size . This “ isotropic ” design is most similar to Transformers , or deep RNNs in other domains , that also use a fixed width . This is unlike most CNNs , which have a pyramidal structure : deeper layers have a lower resolution input , but more channels . Note that while these are the typical designs , other combinations exist , such as isotropic ResNets [ 38 ] and pyramidal ViTs [ 52 ] . Aside from the MLP layers , Mixer uses other standard architectural components : skip-connections [ 15 ] and layer normalization [ 2 ] . Unlike ViTs , Mixer does not use position embeddings because the token-mixing MLPs are sensitive to the order of the input tokens . Finally , Mixer uses a standard classification head with the global average pooling layer followed by a linear classifier . Overall , the architecture can be written compactly in JAX/Flax , the code is given in Supplementary F . 3 Experiments . We evaluate the performance of MLP-Mixer models , pre-trained with medium- to large-scale datasets , on a range of small and mid-sized downstream classification tasks . We are interested in three primary quantities : ( 1 ) Accuracy on the downstream task ; ( 2 ) Total computational cost of pre-training , which is important when training the model from scratch on the upstream dataset ; ( 3 ) Test-time throughput , which is important to the practitioner . Our goal is not to demonstrate state-of-the-art results , but to show that , remarkably , a simple MLP-based model is competitive with today ’ s best convolutional and attention-based models . Downstream tasks We use popular downstream tasks such as ILSVRC2012 “ ImageNet ” ( 1.3M training examples , 1k classes ) with the original validation labels [ 13 ] and cleaned-up ReaL labels [ 5 ] , CIFAR-10/100 ( 50k examples , 10/100 classes ) [ 23 ] , Oxford-IIIT Pets ( 3.7k examples , 36 classes ) [ 33 ] , and Oxford Flowers-102 ( 2k examples , 102 classes ) [ 32 ] . We also use the Visual Task Adaptation Benchmark ( VTAB-1k ) , which consists of 19 diverse datasets , each with 1k training examples [ 58 ] . Pre-training We follow the standard transfer learning setup : pre-training followed by fine-tuning on the downstream tasks . We pre-train our models on two public datasets : ILSVRC2021 ImageNet , and ImageNet-21k , a superset of ILSVRC2012 that contains 21k classes and 14M images [ 13 ] . To assess performance at larger scale , we also train on JFT-300M , a proprietary dataset with 300M examples and 18k classes [ 44 ] . We de-duplicate all pre-training datasets with respect to the test sets of the downstream tasks as done in Dosovitskiy et al . [ 14 ] , Kolesnikov et al . [ 22 ] . We pre-train all models at resolution 224 using Adam with β1 = 0.9 , β2 = 0.999 , linear learning rate warmup of 10k steps and linear decay , batch size 4 096 , weight decay , and gradient clipping at global norm 1 . For JFT-300M , we pre-process images by applying the cropping technique from Szegedy et al . [ 45 ] in addition to random horizontal flipping . For ImageNet and ImageNet-21k , we employ additional data augmentation and regularization techniques . In particular , we use RandAugment [ 12 ] , mixup [ 60 ] , dropout [ 43 ] , and stochastic depth [ 19 ] . This set of techniques was inspired by the timm library [ 54 ] and Touvron et al . [ 48 ] . More details on these hyperparameters are provided in Supplementary B . Fine-tuning We fine-tune using momentum SGD , batch size 512 , gradient clipping at global norm 1 , and a cosine learning rate schedule with a linear warmup . We do not use weight decay when finetuning . Following common practice [ 22 , 48 ] , we also fine-tune at higher resolutions with respect to those used during pre-training . Since we keep the patch resolution fixed , this increases the number of input patches ( say from S to S′ ) and thus requires modifying the shape of Mixer ’ s token-mixing MLP blocks . Formally , the input in Eq . ( 1 ) is left-multiplied by a weight matrix W1 ∈ RDS×S and this operation has to be adjusted when changing the input dimension S. For this , we increase the hidden layer width from DS to DS′ in proportion to the number of patches and initialize the ( now larger ) weight matrix W′2 ∈ RDS′×S ′ with a block-diagonal matrix containing copies of W2 on its diagonal . This particular scheme only allows for S′ = K2S with K ∈ N. See Supplementary C for further details . On the VTAB-1k benchmark we follow the BiT-HyperRule [ 22 ] and fine-tune Mixer models at resolution 224 and 448 on the datasets with small and large input images respectively . Metrics We evaluate the trade-off between the model ’ s computational cost and quality . For the former we compute two metrics : ( 1 ) Total pre-training time on TPU-v3 accelerators , which combines three relevant factors : the theoretical FLOPs for each training setup , the computational efficiency on the relevant training hardware , and the data efficiency . ( 2 ) Throughput in images/sec/core on TPU-v3 . Since models of different sizes may benefit from different batch sizes , we sweep the batch sizes and report the highest throughput for each model . For model quality , we focus on top-1 downstream accuracy after fine-tuning . On two occasions ( Figure 3 , right and Figure 4 ) , where fine-tuning all of the models is too costly , we report the few-shot accuracies obtained by solving the ` 2-regularized linear regression problem between the frozen learned representations of images and the labels . Models We compare various configurations of Mixer , summarized in Table 1 , to the most recent , state-of-the-art , CNNs and attention-based models . In all the figures and tables , the MLP-based Mixer models are marked with pink ( ) , convolution-based models with yellow ( ) , and attention-based models with blue ( ) . The Vision Transformers ( ViTs ) have model scales and patch resolutions similar to Mixer . HaloNets are attention-based models that use a ResNet-like structure with local selfattention layers instead of 3×3 convolutions [ 51 ] . We focus on the particularly efficient “ HaloNet-H4 ( base 128 , Conv-12 ) ” model , which is a hybrid variant of the wider HaloNet-H4 architecture with some of the self-attention layers replaced by convolutions . Note , we mark HaloNets with both attention and convolutions with blue ( ) . Big Transfer ( BiT ) [ 22 ] models are ResNets optimized for transfer learning . NFNets [ 7 ] are normalizer-free ResNets with several optimizations for ImageNet classification . We consider the NFNet-F4+ model variant . We consider MPL [ 35 ] and ALIGN [ 21 ] for EfficientNet architectures . MPL is pre-trained at very large-scale on JFT-300M images , using meta-pseudo labelling from ImageNet instead of the original labels . We compare to the EfficientNetB6-Wide model variant . ALIGN pre-train image encoder and language encoder on noisy web image text pairs in a contrastive way . We compare to their best EfficientNet-L2 image encoder .

Summary: This paper presents an all-MLP based architecture for visual recognition tasks. The presented model is to replace the multi-head self-attention block with a spatial MLP in the original Transformer model. With a large-scale pre-trained dataset, this model with few inductive biases could achieve relatively good performance on image classification tasks.