Source-linked AI summary
Pay Attention to MLPs
Hanxiao Liu, Zihang Dai, David R. So, Quoc V. Le
TL;DR
Transformers leave open whether self-attention is essential to their effectiveness across language and vision. The paper proposes gMLP, an MLP architecture with static spatial projections and multiplicative gating, and finds comparable performance and scaling in key applications, with finetuning limitations on cross-sentence tasks.
Problem
The paper asks whether multi-head self-attention is essential to the effectiveness of Transformers in language and vision.
Method
The paper proposes gMLP, an MLP-based architecture using channel projections, static spatial projections, and multiplicative gating instead of self-attention.
Results
gMLPs are competitive with Transformers on BERT pretraining perplexity and ViT accuracy, and show comparable scalability with increased data and compute.
Takeaways & Limitations
Self-attention is not necessary for strong performance and scaling in the evaluated applications, though small attention additions can help on downstream alignment tasks.
Takeaways & Limitations
Transformers can be more advantageous on finetuning tasks requiring cross-sentence alignment, although larger gMLPs can close the gap.
Abstract
from arXiv · showhide
Transformers have become one of the most important architectural innovations in deep learning and have enabled many breakthroughs over the past few years. Here we propose a simple network architecture, gMLP, based on MLPs with gating, and show that it can perform as well as Transformers in key language and vision applications. Our comparisons show that self-attention is not critical for Vision Transformers, as gMLP can achieve the same accuracy. For BERT, our model achieves parity with Transformers on pretraining perplexity and is better on some downstream NLP tasks. On finetuning tasks where gMLP performs worse, making the gMLP model substantially larger can close the gap with Transformers. In general, our experiments show that gMLP can scale as well as Transformers over increased data and compute.
1 Introduction
The paper asks whether self-attention is essential to Transformer effectiveness and proposes gMLP, a gated MLP alternative, for language and vision. Experiments find comparable results in vision, parity in BERT pretraining, and similar scaling, while some finetuning gaps remain.
- Motivation: Self-attention’s role is questioned because Transformers use dynamic cross-token aggregation, whereas statically parameterized MLPs can represent arbitrary functions.The paper frames the necessity of self-attention’s inductive bias as an open question.
- Approach: gMLP replaces self-attention with channel and spatial projections using static parameterization and multiplicative gating.The architecture is built from basic MLP layers with gating.
- Vision results: With 66% less parameters, gMLP is 3% more accurate than MLP-Mixer and achieves comparable ImageNet performance with DeiT.These comparisons question the necessity of self-attention in Vision Transformers.
- Language results: gMLP matches Transformers on BERT pretraining perplexity, and both pretraining and finetuning metrics improve as quickly with increasing capacity.The authors report that gMLP scales as well as Transformers despite lacking self-attention.
- Language limitations: For finetuning tasks requiring cross-sentence alignment, Transformers can outperform gMLP, but larger gMLPs or a small amount of self-attention can close or reverse the gap.A single-head module of size up to 128 is reported as sufficient to outperform Transformers on evaluated NLP tasks.
- Conclusion: The results suggest self-attention is not necessary for scaling models, although simpler spatial interaction mechanisms can be supplemented by reduced self-attention when useful.The conclusion covers both language and vision experiments.
2 Model
gMLP stacks identical blocks that combine channel projections with a spatial interaction layer and multiplicative gating. Its static spatial projections enable cross-token communication without self-attention, while careful initialization and normalization support stable training.
- Block structure: gMLP consists of L identical blocks operating on token representations X ∈ R^n×d.Each block has the same size and structure.
- Channel projections: The channel projections use linear maps U and V analogous to Transformer feed-forward projections, with an activation such as GeLU between them.Shortcuts, normalizations, and biases are omitted from the abbreviated formulation.
- Spatial interaction: The spatial layer s(·) supplies cross-token communication; if it is an identity, the block reduces to an independently processed regular feed-forward network.Designing s(·) to capture spatial interactions is a central focus.
- Spatial projection: The spatial projection uses W ∈ R^n×n, whose size matches sequence length, and unlike self-attention, W is independent of the input representations.For a 128-token padded sequence, W has shape 128×128.
- Spatial Gating Unit: Linear gating combines spatial projections through element-wise multiplication, with near-zero W and unit b making s(Z) approximately Z at initialization.This makes early blocks behave like regular FFNs before gradually injecting spatial information.
- Spatial Gating Unit: Splitting Z into Z1 and Z2, normalizing the gating input, and using the resulting Spatial Gating Unit improves stability and narrows the gap with self-attention.The split supports separate gating and multiplicative bypass paths.
- Relation to existing layers: SGU computes gating over the spatial cross-token dimension, distinguishing it from channel-based gating in GLUs and related layers.The paper relates SGU to GLUs, Highway Networks, LSTMs, and Squeeze-and-Excite blocks while identifying this architectural distinction.
3 Image Classification
The ImageNet experiments compare gMLP with Transformer and other MLP-like models under ViT/B16-style protocols and capacity-matched configurations. gMLPs are competitive with DeiT and show learned spatial projections with locality and spatial invariance.
- Experimental setup: gMLP models use 16×16 image patches at the stem, with depth and width selected to match ViT/DeiT capacity.The models use a DeiT-like regularization recipe, varying only stochastic-depth strength across sizes.
- Results: Properly regularized gMLPs achieve ImageNet accuracy comparable with DeiT, suggesting self-attention is not required for data-efficient image classification.Their accuracy appears more correlated with model capacity than with the presence of self-attention.
- Results: gMLPs surpass concurrently proposed MLP-like architectures on the accuracy-parameter/FLOPs tradeoff, which the authors attribute to the Spatial Gating Unit.They remain behind the best existing ConvNet and hybrid-model results.
- Learned spatial projections: After learning, gMLP spatial projection weights exhibit locality and spatial invariance across selected tokens.Each projection matrix effectively behaves like a convolution with a data-driven, irregular non-square kernel.
4 Masked Language Modeling with BERT
gMLP matches Transformers on masked-language-modeling perplexity and scales comparably as capacity increases, while downstream transfer depends on task and can benefit from tiny attention or larger models.
- Ablation: The Importance of Gating: Learned gMLP spatial weights become Toeplitz-like, capturing shift invariance from the MLM task without positional encodings.The spatial projection acts like a one-dimensional convolution whose kernel spans the sequence length.
- Ablation: The Importance of Gating: gMLP with the Spatial Gating Unit achieves perplexity comparable to Transformer baselines in MLM pretraining.The strongest baseline has perplexity=4.26, while gMLP has perplexity=4.35; the difference is described as insignificant relative to scaling effects.
- Case Study: Model Size: As model capacity grows, gMLP matches or outperforms Transformer perplexity, and both architecture families approximately follow a power law.The results suggest Transformer language-model scaling laws may apply across model families.
- Case Study: Model Size: Pretraining perplexity does not determine finetuning equally: gMLPs outperform Transformers on SST-2 but underperform on MNLI, despite comparable scalability.The authors relate this discrepancy to architectural inductive bias and possible cross-sentence alignment requirements.
- Tiny Attention in Finetuning: A tiny single-head attention module can close gMLP’s MNLI gap and produce consistent finetuning gains over Transformers.The hybrid aMLP adds 64-dimensional single-head attention, and the combined scaling results show it outperforming Transformers on both finetuning tasks.
- Main Results for MLM: Increasing gMLP capacity narrows finetuning deficits, with the SQuAD-v2.0 gap shrinking from 8.5% at base scale to 2.7% at larger scale.The larger gMLP also achieves 89.5% F1 on SQuAD-v1.1 without self-attention and can outperform BERTlarge on SQuAD-v2.0 at a larger model size.
5 Conclusion
The paper presents gMLPs as a simpler alternative to Transformer self-attention, competitive in language and vision and scalable with increased data and compute. For BERT finetuning, gMLPs can underperform on cross-sentence alignment tasks, but larger models or limited self-attention can close the gap.
- gMLPs provide a simpler alternative to Transformer multi-head self-attention while remaining competitive in BERT pretraining perplexity and ViT accuracy.
- gMLPs achieve appealing results on challenging tasks such as SQuAD without self-attention and can significantly outperform Transformers in certain cases.
- Transformer self-attention is useful on downstream tasks requiring cross-sentence alignment, where gMLPs may otherwise lag.
- Making gMLPs substantially larger closes the finetuning gap, while blending a small single-head self-attention module provides a more parameter-efficient alternative.
A.1 Image Classification
The ImageNet experiments use TPUv2 hardware and a specified hyperparameter configuration for image classification.
- ImageNet models are trained using TPUv2 with 128 cores, and each run takes 1–4 hours.
- Table 7 lists the hyperparameters used for ImageNet-1K image classification.
A.2 Masked Language Modeling
The masked-language-modeling experiments specify separate hardware and hyperparameter configurations for ablations, pretraining, and finetuning.
- MLM ablation studies use TPUv3 with 32 cores, while full BERT models use TPUv2 with 128 cores.
- Full BERT training runs take 1–5 days depending on model size, and the vocabulary contains 32K cased SentencePieces.
- Table 8 specifies hyperparameters for MLM pretraining on C4.
- Table 9 specifies hyperparameters for MLM finetuning on GLUE and SQuAD.
B Deep-and-Thin Transformers
At fixed capacity, increasing Transformer depth does not improve perplexity, and performance worsens beyond 48 layers under the reported setup.
- At fixed capacity, perplexity is insensitive to Transformer depth, while performance worsens beyond 48 layers.The experiments adjust model width as depth increases to maintain comparable capacity, using a similar but different training setup.
C Shift Invariance in MLM
The MLM experiments use Toeplitz spatial projection matrices, and the implementation constructs these matrices from a learnable vector of 2n−1 weights.
- gMLP learns Toeplitz-like spatial projection matrices during MLM pretraining, even without imposing a shift-invariance prior.
- The MLM experiments use Toeplitz matrices as their spatial projection structure.
- The implementation defines a function that creates an n × n Toeplitz matrix.
- The matrix is generated by padding, tiling, slicing, reshaping, and selecting entries from the learnable weight vector.
D Visualizing Tiny Attention
The paper visualizes token-pair attention in tiny attention modules after MNLI-m finetuning, with heatmaps summarizing the strongest attention received during the first half of the network.
- The visualization concerns tiny attention modules in aMLP after finetuning on MNLI-m.
- Each heatmap element records the maximum attention weight ever received by a token pair during the network’s first half.