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 proposes a MLP-based model that achieves comparable performance with transformers for vision tasks and language modeling. While this MLP-based model performs slighly worse on natural language understanding tasks (e.g. MNLI), introducing a small single-head self-attention module can compenate for most of the performance loss. Compared to other MLP-based models introduced in concurrent works, the proposed model uses a gating module before the spatial projection. Performance wise, the proposed method is stronger than the concurrent MLP variant MLP-mixer on vision and language modeling tasks, showing the advantage of the proposed gating methods  Overall I find the proposed method to be interesting and empirically strong. I am leaning toward accept but also wish the authors had done more work in understanding the proposed models. I offer some suggestions and have some questions below.