Pay Attention to MLPs

1 Introduction . Transformers [ 1 ] have enabled many breakthroughs in natural language processing ( e.g. , [ 2 , 3 , 4 , 5 , 6 ] ) and have been shown to work well for computer vision ( e.g. , [ 7 , 8 , 9 , 10 ] ) . Thanks to this success , Transformers have largely replaced LSTM-RNN [ 11 ] as the default architecture in NLP , and have become an appealing alternative to ConvNets [ 12 , 13 , 14 , 15 , 16 , 17 ] in computer vision . The Transformer architecture combines two important concepts : ( 1 ) a recurrent-free architecture which computes the representations for each individual token in parallel , and ( 2 ) multi-head selfattention blocks which aggregate spatial information across tokens . On one hand , the attention mechanism [ 18 ] introduces the inductive bias that the spatial interactions should be dynamically parameterized based on the input representations . On the other hand , it is known that MLPs with static parameterization can represent arbitrary functions [ 19 ] . It therefore remains an open question whether the inductive bias in self-attention is essential to the remarkable effectiveness of Transformers . Here we study the necessity of self-attention modules in key language and vision applications of Transformers . Specifically , we propose an MLP-based alternative to Transformers without self-attention , which simply consists of channel projections and spatial projections with static parameterization . We experiment with several design choices for this architecture and find spatial projections work well when they are linear and paired with multiplicative gating ( Figure 1 ) . We name the model gMLP because it is built out of basic MLP layers with gating . We apply gMLP to image classification and obtain strong results on ImageNet . gMLP achieves comparable performance with DeiT [ 8 ] , namely Vision Transformer ( ViT ) [ 7 ] with improved regularization , in a similar training setup . With 66 % less parameters , a gMLP model is 3 % more accurate than MLP-Mixer [ 20 ] . Together with Tolstikhin et al . [ 20 ] , Melas-Kyriazi [ 21 ] , Touvron et al . [ 22 ] and Ding et . al . [ 23 ] , our results question the necessity of self-attention layers in Vision Transformers . We apply gMLP to masked language modeling ( MLM ) in the BERT [ 2 ] setup , one of the most wellestablished applications of Transformers , and find that it is as good as Transformers at minimizing perplexity during pretraining . Our experiments indicate that perplexity is only correlated with model capacity and is insensitive to the presence of self-attention . As capacity increases , we observe that 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . both pretraining and finetuning metrics for gMLPs improve as quickly as for Transformers . This is remarkable because it indicates gMLPs scale just as well as Transformers despite the absence of self-attention , and any performance gap can always be offset by training a larger model with increased data and compute . With a standard 256-batch size × 1M-step training setup as in original BERT , a large gMLP model achieves 87.7 % accuracy on MNLI and 82.1 % F1 on SQuAD v2.0 . Note , these are better than the BERTlarge results reported in Devlin et al . [ 2 ] obtained using Transformers . For BERT ’ s finetuning , Transformers can be more practically advantageous over gMLPs on tasks that require cross-sentence alignment ( e.g. , by 0.8 % on MNLI-m in the 300M-param regime ) , even with similar pretraining perplexity . This problem can be addressed by making gMLPs substantially larger—3× as large as Transformers . A more practical solution is to blend in only a tiny bit of selfattention—a single-head self-attention with size up to 128 is sufficient to make gMLPs outperform Transformers on all NLP tasks we evaluated with even better parameter efficiency . The improvement is sometimes very significant ( e.g. , +4.4 % on SQuAD v2.0 over BERTlarge ) . Overall , the surprising effectiveness of gMLPs in both vision and NLP domains suggests that selfattention is not a necessary ingredient for scaling up machine learning models , although it can be a useful addition depending on the task . With increased data and compute , models with simpler spatial interaction mechanisms such as gMLP can be as powerful as Transformers and the capacity allocated to self-attention can be either removed or substantially reduced . 2 Model . Our model , gMLP , consists of a stack of L blocks with identical size and structure . Let X ∈ Rn×d be the token representations with sequence length n and dimension d. Each block is defined as : Z = σ ( XU ) , Z̃ = s ( Z ) , Y = Z̃V ( 1 ) where σ is an activation function such as GeLU [ 24 ] . U and V define linear projections along the channel dimension—the same as those in the FFNs of Transformers ( e.g. , their shapes are 768× 3072 and 3072× 768 for BERTbase ) . Shortcuts , normalizations and biases are omitted for brevity . A key ingredient in the aforementioned formulation is s ( · ) , a layer which captures spatial interactions ( see below ) . When s is an identity mapping , the above transformation degenerates to a regular FFN , where individual tokens are processed independently without any cross-token communication . One of our major focuses is therefore to design a good s capable of capturing complex spatial interactions across tokens . The overall block layout is inspired by inverted bottlenecks [ 25 ] which define s ( · ) as a spatial depthwise convolution . Note , unlike Transformers , our model does not require position embeddings because such information will be captured in s ( · ) . Our model uses exactly the same input and output protocols as BERT ( for NLP ) and ViT ( for vision ) . For example , when finetuning on language tasks , we concatenate together multiple text segments followed by paddings , and the predictions are deduced from the last-layer representation of a reserved < cls > symbol . Although many of these protocols were introduced for Transformers and hence can be suboptimal for gMLPs , strictly following them helps avoid confounding factors in our experiments and makes our layers more compatible with existing Transformer implementations . 2.1 Spatial Gating Unit . To enable cross-token interactions , it is necessary for the layer s ( · ) to contain a contraction operation over the spatial dimension . The simplistic option would be a linear projection : fW , b ( Z ) =WZ + b ( 2 ) where W ∈ Rn×n is a matrix for which the size is the same as the sequence length , n , and b refers token-specific biases . For example , if the padded input sequence has 128 tokens , the shape for W will be 128×128 . Unlike self-attention where W ( Z ) is dynamically generated from Z , the spatial projection matrix W here in Equation ( 2 ) is independent from the input representations . In this work , we formulate layer s ( · ) as the output of linear gating : s ( Z ) = Z fW , b ( Z ) ( 3 ) where denotes element-wise multiplication . For training stability , we find it critical to initialize W as near-zero values and b as ones , meaning that fW , b ( Z ) ≈ 1 and therefore s ( Z ) ≈ Z at the beginning of training . This initialization ensures each gMLP block behaves like a regular FFN at the early stage of training , where each token is processed independently , and only gradually injects spatial information across tokens during the course of learning . We further find it effective to split Z into two independent parts ( Z1 , Z2 ) along the channel dimension for the gating function and for the multiplicative bypass : s ( Z ) = Z1 fW , b ( Z2 ) ( 4 ) We also normalize the input to fW , b which empirically improves stability of large NLP models . This gives us the unit illustrated in Figure 1 , which we refer to as the Spatial Gating Unit ( SGU ) in the rest of the paper . In Table 3 , we provide ablation studies to compare SGU with several other variants of s ( · ) , showing that it works better and narrows the performance gap with self-attention . Connections to Existing Layers . The overall formulation of SGU resembles Gated Linear Units ( GLUs ) [ 26 , 27 , 28 ] as well as earlier works including Highway Networks [ 29 ] and LSTM-RNNs [ 11 ] . A key distinction is that our gating is computed based on a projection over the spatial ( cross-token ) dimension rather than the channel ( hidden ) dimension . SGU is also related to Squeeze-and-Excite ( SE ) blocks [ 30 ] in terms of element-wise multiplication . However , different from SE blocks , SGU does not contain cross-channel projections at all , nor does it enforce permutation invariance ( a key feature for content-based attentive modules ) due to its static parameterization for the spatial transformation . The spatial projection in SGU could in theory learn to express superficial depthwise convolutions—unlike typical depthwise convolutions with channel-specific filters , SGU learns only a single transformation shared across channels . Finally , we note SGUs offer an alternative mechanism to capture high-order relationships other than self-attention . Specifically , the output for Equation ( 3 ) contains up to 2nd-order interactions ( e.g. , zizj ) whereas output for self-attention ( assuming no nonlinearity ) contains up to 3rd-order interactions ( e.g. , qikjvk ) . In terms of computation cost , SGU has n2e/2 multiply-adds which is comparable to the 2n2d of dot-product self-attention.1 Both are linear over the input channel size and quadratic over the sequence length n . 3 Image Classification . Here we examine gMLP in the vision domain by applying it to the image classification task on ImageNet [ 31 ] without using extra data . We compare our MLP-like models with recent attentive 1The input channel size e for SGU is typically larger than the input channel size d for self-attention , because the former is applied in the middle of the block after a channel expansion . models based on vanilla Transformers , including Vision Transformer ( ViT ) [ 7 ] , DeiT [ 8 ] ( ViT with improved regularization ) , and several other representative convolutional networks . Table 1 summarizes the configurations of our gMLP image classification models . The input and output protocols follow ViT/B16 where the raw image is converted into 16×16 patches at the stem . The depth and width are chosen so that the models are comparable with ViT/DeiT in capacity . Like Transformers , we find gMLPs tend to drastically overfit the training data . We therefore apply a similar regularization recipe as the one used in DeiT.2 To avoid extensive tuning , we adjust only the strengths of stochastic depth [ 32 ] as we move from smaller to larger models in Table 1 . All the other hyperparameters remain shared across our three models . See Appendix A.1 for details . Our ImageNet results are summarized in Table 1 and Figure 2.3 It is interesting to see that gMLPs are comparable with DeiT [ 8 ] , namely ViT [ 7 ] trained using improved regularization . The results suggest that models without self-attention can be as data-efficient as Transformers for image classification . In fact , when the models are properly regularized , their accuracies seem better correlated with capacity instead of the presence of self-attention . Moreover , the accuracy-parameter/FLOPs tradeoff of gMLPs surpasses all concurrently proposed MLP-like architectures [ 20 , 21 , 22 ] , which we attribute to the effectiveness of our Spatial Gating Unit ( see Table 3 in the next section for an ablation ) . We also note while gMLPs are competitive with vanilla Transformers , their performance is behind the best existing ConvNet models ( e.g. , [ 33 , 34 ] ) or hybrid models ( e.g. , [ 35 , 36 , 37 , 38 , 10 ] ) . Figure 3 visualizes the spatial projection matrices in gMLP-B . Remarkably , the spatial weights after learning exhibit both locality and spatial invariance . In other words , each spatial projection matrix effectively learns to perform convolution with a data-driven , irregular ( non-square ) kernel shape . 2Unlike DeiT , we do not use repeated augmentation or random erasing . 3Additional results on ImageNet-21K and COCO are available in Appendix E .

Summary: This paper presents a simple architecture based on MLPs, which only adopts channel projections and spatial projections with the spatial gating units. On image classification, the presented gMLP could achieve on par performance with DeiT and ViT, and better than previous MLP-based approaches, such as MLP Mixer. On masked language modeling tasks, gMLP could also achieve on par performance with Transformer on Perplexity, and gMLP could perform slightly worse than Transformer on MNLI.