Source-linked AI summary
MLP-Mixer: An all-MLP Architecture for Vision
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy
TL;DR
CNNs and attention-based models dominate vision, but whether either mechanism is necessary remains open. MLP-Mixer replaces both with MLPs that mix spatial locations and feature channels, achieving competitive accuracy–cost trade-offs, including 87.94% top-1 validation accuracy on ImageNet.
Problem
Vision research is dominated by CNNs and increasingly self-attention models, motivating alternatives that remove both architectural mechanisms.
Method
MLP-Mixer uses only MLPs to mix information across image patches and feature channels.
Results
87.94% top-1 validation accuracy on ImageNet, with competitive accuracy–compute trade-offs versus CNNs and Transformers.
Takeaways & Limitations
MLP-Mixer shows that competitive vision performance does not require convolutions or self-attention, opening research beyond these established models.
Takeaways & Limitations
At more modest pre-training scale, Mixer requires modern regularization and falls slightly short of specialized CNN architectures.
Abstract
from arXiv · showhide
Convolutional Neural Networks (CNNs) are the go-to model for computer vision. Recently, attention-based networks, such as the Vision Transformer, have also become popular. In this paper we show that while convolutions and attention are both sufficient for good performance, neither of them are necessary. We present MLP-Mixer, an architecture based exclusively on multi-layer perceptrons (MLPs). MLP-Mixer contains two types of layers: one with MLPs applied independently to image patches (i.e. "mixing" the per-location features), and one with MLPs applied across patches (i.e. "mixing" spatial information). When trained on large datasets, or with modern regularization schemes, MLP-Mixer attains competitive scores on image classification benchmarks, with pre-training and inference cost comparable to state-of-the-art models. We hope that these results spark further research beyond the realms of well established CNNs and Transformers.
1 Introduction
MLP-Mixer is a conceptually simple vision architecture that replaces convolutions and self-attention with MLPs mixing spatial locations and feature channels. Despite this simplicity, it achieves competitive accuracy/cost trade-offs, including 87.94% top-1 ImageNet validation accuracy after large-scale pre-training.
- Architecture: MLP-Mixer uses only MLPs, matrix multiplication, data-layout changes, and scalar nonlinearities, without convolutions or self-attention.The architecture repeatedly applies MLPs across spatial locations or feature channels.
- Architecture: Mixer processes linearly projected image patches as a patches × channels table using interleaved token-mixing and channel-mixing MLPs.Token-mixing MLPs communicate across spatial locations independently for each channel, while channel-mixing MLPs communicate across channels independently for each token.
- Architecture: Each Mixer layer contains token- and channel-mixing MLPs built from two fully connected layers and GELU, with skip-connections, dropout, and channel-wise layer normalization.The architecture also includes per-patch linear embeddings and a classifier head.
- Results: 87.94% top-1 validation accuracy is achieved on ILSVRC2012 ImageNet after pre-training on approximately 100M images, with a near state-of-the-art accuracy/cost trade-off.Mixer also achieves strong performance with approximately 1–10M pre-training images when combined with modern regularization techniques.
2 Mixer Architecture
Mixer transforms non-overlapping image patches into a fixed-size table and alternates token-mixing and channel-mixing MLPs. Its design uses tied per-row/per-column transformations, linear complexity, and standard residual normalization components.
- Input representation: Mixer projects S non-overlapping P×P image patches into a two-dimensional table X ∈ R^S×C using one shared projection matrix.For an image of resolution (H, W), S = HW/P^2.
- MLP mixing: Each Mixer layer applies MLP blocks independently across rows and columns, respectively mixing channel features and spatial tokens.Each MLP block contains two fully-connected layers and a nonlinearity applied independently to each input row.
- Computational complexity: Mixer’s token-mixing width is independent of patch count, yielding linear complexity in patches; its channel-mixing width is independent of patch size, yielding linear pixel complexity.The paper contrasts this with ViT’s quadratic patch complexity and notes CNN-like linear complexity in image pixels.
- Parameter sharing: Parameter tying applies the same channel-mixing MLP to every row and the same token-mixing MLP to every column, providing positional invariance for channel mixing.The architecture identifies positional invariance as a prominent feature of convolutions.
- Architectural components: Beyond its MLPs, Mixer uses skip-connections, layer normalization, no position embeddings, and global average pooling followed by a linear classification head.Token-mixing MLPs are sensitive to input-token order, so Mixer does not use position embeddings.
3 Experiments
The experiments evaluate Mixer across diverse downstream classification tasks, emphasizing accuracy, pre-training cost, and test-time throughput. Results show that Mixer improves with larger pre-training datasets and can achieve competitive accuracy–efficiency trade-offs, especially at larger scales.
- Experimental setup: Mixer is evaluated on ImageNet, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102, and the 19-dataset VTAB-1k benchmark.The experiments use downstream datasets ranging from 2k to 1.3M examples, with VTAB-1k providing 19 diverse tasks.
- Experimental setup: Mixer models are compared with similarly scaled CNNs and attention-based models, including ViTs and BiT models.The comparisons cover multiple Mixer configurations and use pink, yellow, and blue markers for MLP-, convolution-, and attention-based models, respectively.
- Accuracy–compute trade-off: 84.15% top-1 ImageNet accuracy is achieved by Mixer after ImageNet-21k pre-training with additional regularization, although it remains slightly inferior to other models.Regularization is necessary in this setting because Mixer otherwise overfits.
- Accuracy–compute trade-off: 87.94% top-1 ImageNet accuracy is achieved by Mixer-H/14, outperforming BiT-ResNet152x4 by 0.5% and trailing ViT-H/14 by 0.5%.Mixer-H/14 runs 2.5 times faster than ViT-H/14 and almost twice as fast as BiT.
- Scaling with pre-training data: As pre-training data scale increases, Mixer’s performance steadily improves, with JFT-300M Mixer-H/14 only 0.3% behind ViT-H/14 on ImageNet while running 2.2 times faster.At the high end of model scales, Mixer sits confidently on the accuracy–compute frontier.
4 Related work
MLP-Mixer differs from successful vision architectures by using neither convolutions nor self-attention, while drawing design ideas from CNNs and Transformers. It also relates to fully connected networks and recent advances in architecture and training strategies.
- Architectural context: MLP-Mixer uses neither convolutional nor self-attention layers, while its design choices trace back to CNN and Transformer literature.This distinguishes MLP-Mixer from previous successful architectures in computer vision.
- CNNs: CNNs became computer vision’s de-facto standard after AlexNet surpassed approaches based on hand-crafted image features.Subsequent work improved CNN design, including small 3×3 convolutions and skip-connections with batch normalization.
- Transformers: Vision Transformers progressed from generative modeling and locality-biased or low-resolution recognition to ViT, which scales well to large data with fewer locality biases.The passage identifies ViT as a pure transformer model achieving state-of-the-art performance on popular vision benchmarks.
- Attention-based architectures: Mixer takes an orthogonal direction from recent attention-based vision architectures by avoiding both locality bias and attention mechanisms.Related work includes replacing ResNet convolutions with self-attention and designing new attention-like mechanisms.
- Fully connected networks: Compared with prior fully connected networks, MLP-Mixer uses token- and channel-mixing MLPs, standard regularization and optimization, and effective scaling to large data.Earlier work relied on heavy augmentation, autoencoder pre-training, or custom regularization and optimization, often for small-scale tasks.
5 Conclusions
The paper presents MLP-Mixer as a simple vision architecture matching existing state-of-the-art methods in the trade-off between accuracy and training and inference costs. The authors identify practical and theoretical questions and call for research beyond convolutional and self-attention-based models.
- Conclusions: MLP-Mixer is a very simple vision architecture whose accuracy–resource trade-off matches existing state-of-the-art methods for training and inference.The conclusion frames this as the central experimental finding.
- Conclusions: Future practical work should analyze MLP-Mixer’s learned features and their differences, if any, from CNN and Transformer features.This comparison is proposed as a practical research direction.
- Conclusions: Future theoretical work should study the inductive biases hidden in these features and their eventual role in generalization.The authors present this as a theoretical question raised by the results.
- Conclusions: The results are intended to spark research beyond established convolutional and self-attention models, including possible applications in NLP and other domains.The paper specifically highlights whether this design can work in NLP or elsewhere.
A Things that did not help · A.1 Modifying the token-mixing MLPs
The authors ablated several modifications to Mixer’s token-mixing MLPs across model scales pre-trained on JFT-300M. Untying parameters, grouping channels, alternative channel-view projections, and pyramid token reduction did not yield transferable improvements.
- A.1 Modifying the token-mixing MLPs: The study ablated multiple token-mixing MLP modifications in Mixer models of various scales pre-trained on JFT-300M.The experiments targeted changes intended to improve token mixing.
- A.1 Modifying the token-mixing MLPs: Untying the token-mixing MLP parameters across feature columns produced no noticeable improvement.This change introduced C independently weighted MLPs and multiplied the parameter count by C.
- A.1 Modifying the token-mixing MLPs: Grouping G neighboring channels increased the token-mixing input dimensionality from S to G · S while reducing the number of channel groups.The transformation reshaped X ∈ R^(S×C) into a matrix of dimension (S · G) × (C/G).
- A.1 Modifying the token-mixing MLPs: An alternative channel-grouping design projected each token into G trainable views, each containing C/G features, before concatenation.The procedure used G linear functions mapping R^C to R^(C/G).
- A.1 Modifying the token-mixing MLPs: Pyramid token-mixing models reduced tokens from S inputs to S′ < S outputs and significantly reduced JFT-300M training time without much initial performance loss.These early results did not transfer to ImageNet or ImageNet-21k.
- A.1 Modifying the token-mixing MLPs: The pyramid findings could not be transferred from JFT-300M to ImageNet or ImageNet-21k.Thus, the observed training-time benefit did not establish a generally useful modification.
A.2 Fine-tuning
During fine-tuning, the authors tested mixup, Polyak averaging, and inception cropping on JFT-300M-pre-trained Mixer models, but found no consistent or any improvements and dropped these techniques.
- Fine-tuning: Mixup and Polyak averaging did not produce consistent improvements during fine-tuning, so the authors dropped them.These experiments followed ideas from BiT and ViT.
- Fine-tuning: Inception cropping also failed to improve fine-tuning results.
- Fine-tuning: The experiments covered JFT-300M-pre-trained Mixer models at all scales.
B Pre-training: hyperparameters, data augmentation and regularization
The section describes the hyperparameter settings used to pre-train Mixer models and the augmentation and regularization procedures applied on ImageNet and ImageNet-21k. These procedures sweep RandAugment magnitude and mixup strength, test dropout rates, and use linearly increasing stochastic-depth probabilities.
- Hyperparameters: Table 4 reports the optimal hyperparameter settings used for pre-training Mixer models.The passage identifies Table 4 as the source of these settings but does not provide their individual values.
- Data augmentation: For ImageNet and ImageNet-21k pre-training, RandAugment uses two augmentation layers while sweeping magnitude m over {0, 10, 15, 20}.The augmentation configuration is applied to both pre-training datasets.
- Regularization: Mixup sweeps strength p over {0.0, 0.2, 0.5, 0.8}, dropout tests rates d of 0.0 and 0.1, and stochastic depth increases layer-dropping probability linearly.For stochastic depth, the probability increases from 0.0 for the first MLP to s; the supplied passage ends there.
C Fine-tuning: hyperparameters and higher image resolution
The section specifies the fine-tuning setup and extends Mixer to higher-resolution inputs without changing patch size. The extension lengthens token sequences and adapts token-mixing MLPs by widening and structurally expanding their parameters.
- Fine-tuning setup: Models are fine-tuned at 224 resolution, with task-specific learning-rate grids and modified evaluation preprocessing.The learning-rate grid excludes 0.001, adds 0.06 for CIFAR-10, CIFAR-100, Flowers, and Pets, and uses {0.003, 0.01, 0.03} for VTAB-1k; one evaluation option is resize-crop from 256 × 256 to 224 × 224.
- Higher-resolution fine-tuning: At higher resolution, Mixer keeps patch size fixed, producing longer token sequences and requiring adjusted token-mixing MLPs.The method follows the established practice of fine-tuning at higher resolution than pre-training to improve transfer performance.
- Higher-resolution fine-tuning: For an integer resolution factor K, sequence length S and token-mixing hidden width D_S both increase by K^2.The increased sequence is partitioned into K^2 equal parts, each retaining the original length S.
- Higher-resolution fine-tuning: The expanded token-mixing input matrix is initialized as a block-diagonal matrix containing copies of the pre-trained matrix W_1.The resized token sequence is formed from K^2 spatially split subsequences, and other MLP parameters are handled analogously.
D Weight visualizations · E MLP-Mixer code
The weight visualizations show that learned projection structure depends strongly on patch resolution, with 32×32 patches producing Gabor-like low-frequency units and 16×16 patches lacking that structure. The accompanying JAX/Flax listing implements MLP blocks, Mixer blocks, and the overall MLP-Mixer module.
- D Weight visualizations: The visualizations sort hidden units by a heuristic intended to display low-frequency filters first and pair each unit with its closest inverse.The authors note that heavy augmentation in ImageNet and ImageNet-21k pre-training strongly influences learned-unit structure.
- D Weight visualizations: 32×32 patches produce Gabor-like low-frequency embedding projections, whereas 16×16 patches show no such structure across Mixer scales.Figure 7 similarly reports structured low-frequency units for Mixer-B/32 and mostly high-frequency, unstructured units for Mixer-B/16.
- D Weight visualizations: Figure 6 depicts each hidden unit in Mixer-B/16’s first two token-mixing MLPs as a 14 × 14 image containing 196 incoming-token weights.Each block contains 384 hidden units, and columns correspond to models trained on three datasets.
- E MLP-Mixer code: The code imports Flax neural-network modules and JAX NumPy, then defines an MlpBlock whose call applies a dense layer to its input.The listing identifies the implementation as MLP-Mixer code written in JAX/Flax.
- E MLP-Mixer code: MixerBlock includes a channel-mixing MlpBlock parameterized by channels_mlp_dim and adds its output to the block input.The relevant method and class declarations appear in the code listing.
- E MLP-Mixer code: The MlpMixer module rearranges features from n h w c into n (h w) c before applying Mixer blocks.The code declares MlpMixer and its channels_mlp_dim and call method before performing this rearrangement.
- E MLP-Mixer code: The overall implementation iterates over num_blocks, applies MixerBlock with token and channel MLP dimensions, and uses zero kernel initialization in the listing.The final line is part of Listing 1, identified as MLP-Mixer code written in JAX/Flax.