🌀 MLP-Mixer: How image patches communicate without attention

Convolutional networks mix information with filters that slide across an image. Vision Transformers use attention to decide how image patches should interact. MLP-Mixer asks a simpler question: can ordinary multilayer perceptrons do both jobs?

The answer in Ilya Tolstikhin and colleagues' 2021 paper is qualified but important. With enough data and regularization, an architecture made from matrix multiplications, transpositions, and element-wise nonlinearities can be competitive with strong CNNs and Vision Transformers on image classification. The paper does not claim that MLP-Mixer is universally better, or that attention and convolution have become unnecessary for every vision task.

This article follows the NeurIPS 2021 paper, first submitted to arXiv on May 4, 2021.

Figure 1. One Mixer layer. The first MLP communicates across image patches; the second communicates across feature channels. Both sub-blocks use layer normalization and residual connections. Adapted from the paper's figure 1.

The idea in one minute🔗

MLP-Mixer turns an image into a table. Each row represents one image patch, and each column represents one learned feature channel.

A Mixer layer alternates between two operations:

1. Token mixing: for one channel at a time, combine information from all patch locations.
2. Channel mixing: for one patch at a time, combine information from all channels.

Repeating these operations lets information travel along both axes of the table. A final average over patches and a linear classifier produce the class prediction.

The useful mental model is not "one large MLP over every pixel." It is two smaller, shared MLPs that take turns mixing the spatial and feature dimensions.

From pixels to a patch-by-channel table🔗

Suppose an image has height \(H\), width \(W\), and is divided into non-overlapping square patches of width \(P\). The number of patches is:

For a \(224\times224\) image with \(16\times16\) patches, \(S=196\): a \(14\times14\) grid. Each patch is flattened and passed through the same learned linear projection, producing a vector with \(C\) channels. Stacking the vectors creates the input table:

Here \(S\) is the number of patch tokens and \(C\) is the hidden channel width. Unlike a standard CNN, the network keeps this table shape throughout its stack of Mixer layers; it does not progressively shrink the spatial grid while increasing the channel count.

Figure 2. The full classification pipeline: patch projection, a stack of equally sized Mixer layers, global average pooling, and a linear classification head.

Follow one table through a Mixer layer🔗

The table orientation matters. In \(X\in\mathbb{R}^{S\times C}\), rows are patches and columns are channels.

Step 1: token mixing communicates across space🔗

Token mixing takes one channel, reads its values at all \(S\) patch locations, and passes that length-\(S\) vector through an MLP. The same token-mixing MLP is reused for every one of the \(C\) channels.

This operation has a global receptive field: a value at one patch can influence the output at every patch in a single block. Unlike self-attention, however, the mixing weights do not change according to the content of the current image. They are learned parameters shared across examples.

Step 2: channel mixing communicates within each patch🔗

Channel mixing takes one patch, reads its \(C\) feature values, and passes that length-\(C\) vector through a second MLP. The same channel-mixing MLP is reused at every patch location. This is closely related to applying a stack of \(1\times1\) convolutions: it changes which features interact at a location without directly moving information between locations.

Layer normalization precedes each MLP, and a residual connection adds the sub-block's input back to its output. Each MLP contains two fully connected layers separated by a GELU nonlinearity.

Figure 3. The two views of the same table. Transposing exposes the patch dimension to the token-mixing MLP; transposing back exposes the channel dimension to the channel-mixing MLP.

The sharing pattern is the key:

• token-mixing parameters are shared across channels;
• channel-mixing parameters are shared across patch locations.

The first choice keeps the model from learning a separate spatial mixer for every feature. The second applies the same feature transformation wherever a patch occurs.

Why no positional embedding?🔗

Vision Transformers add positional embeddings because self-attention alone does not know the order of its input tokens. MLP-Mixer does not add them because the token-mixing matrices already assign different learned weights to different positions in their length-\(S\) input.

That does not make Mixer permutation-invariant at inference time. If the patches of a trained model are arbitrarily reordered, the token-mixing weights no longer correspond to the positions on which they were learned. The paper's permutation experiment instead trains a new model using the same fixed permutation for every image; under that controlled setup, Mixer can learn the rearranged coordinate system.

The same position-specific weights also tie a token-mixing MLP to a particular number of patches. Fine-tuning at a higher resolution changes \(S\) and requires resizing the token-mixing weights. The paper uses a structured initialization for that resizing rather than accepting arbitrary image sizes without modification.

Is MLP-Mixer a CNN in disguise?🔗

The paper says Mixer can be viewed as a very special CNN:

• channel mixing resembles \(1\times1\) convolution;
• token mixing resembles a single-channel depth-wise convolution with a full-image receptive field, with the same spatial kernel shared across channels.

This analogy explains some mechanics, but not an equivalence between the model families. A conventional convolution uses a local kernel that slides across locations. Mixer's token-mixing weights connect particular input positions to particular output positions across the entire patch grid. Typical CNNs can also change resolution and build spatial hierarchies; the original Mixer remains isotropic.

Nor are Mixer's weights "chosen by attention." Attention computes input-dependent mixing weights from each example. Mixer applies the same learned token-mixing transformation to every example. Both allow long-range interaction, but by different mechanisms.

Complexity: what "linear in the number of patches" means🔗

For one layer, token mixing costs roughly \(O(CSD_S)\) and channel mixing costs roughly \(O(SCD_C)\). If \(C\), \(D_S\), and \(D_C\) are held fixed while \(S\) grows, both costs grow linearly with the number of patches. Standard global self-attention contains an \(S\times S\) attention matrix and therefore has a quadratic term in \(S\).

This is the paper's asymptotic comparison, not a promise that Mixer is always faster. Actual throughput depends on model width, input resolution, hardware, batch size, and implementation. The authors therefore report measured TPUv3 throughput alongside accuracy and pretraining cost.

What the experiments show🔗

The stated goal was not to set the highest score in every column. It was to test whether a simple all-MLP model could reach a competitive accuracy-cost trade-off.

The experiments first pretrain models on ImageNet, ImageNet-21k, or the proprietary JFT-300M dataset. They then evaluate transfer to ImageNet with its original and cleaned ReaL labels, CIFAR-10, CIFAR-100, Oxford-IIIT Pets, Oxford Flowers-102, and the 19-task VTAB-1k benchmark.

Most reported accuracies come from fine-tuning, which updates the pretrained model for the downstream dataset. Two scaling experiments instead freeze the representation and fit an \(\ell_2\)-regularized linear classifier from only five labeled examples per class. Top-1 accuracy is the percentage of examples for which the model's highest-scoring class matches the target. Keeping these protocols separate is essential: five-shot linear accuracy and full fine-tuning answer different questions.

Model sizes🔗

The paper evaluates Small, Base, Large, and Huge variants. The letters describe the whole model scale; the number after the slash is the patch width. For example, Mixer-B/16 is a Base model using \(16\times16\) patches.

The parameter counts exclude the classification head. The Small model is not 8 layers, Base 16, Large 24, and Huge 32; the paper's Base variants use 12 layers.

Accuracy, throughput, and training cost🔗

Figure 4. Selected rows from the paper's table 2. Throughput is measured in images per second per TPUv3 core; TPUv3 core-days estimate total pretraining cost. The highlighted Mixer rows should be compared within the same pretraining-data group.

With public ImageNet-21k pretraining, Mixer-L/16 reaches 84.15% ImageNet top-1 accuracy at 105 images per second per core. ViT-L/16 reaches 85.30% at 32 images per second per core. Mixer is faster in this comparison but less accurate, and it uses more reported pretraining compute: 0.41k versus 0.18k TPUv3 core-days.

With proprietary JFT-300M pretraining, Mixer-H/14 reaches 87.94%, compared with 88.55% for ViT-H/14 and 87.54% for BiT-R152x4. Its reported throughput is 40 images per second per core, versus 15 for ViT and 26 for BiT. These results support competitiveness at large scale, not dominance on every metric.

When trained from scratch on ImageNet, Mixer-B/16 reaches 76.44%, about three percentage points behind the matched ViT-B/16 result. The training losses are similar, suggesting that Mixer overfits more. Modern augmentation and regularization are therefore important, especially with smaller pretraining datasets.

More data narrows the gap🔗

The authors pretrain Mixer, ViT, and BiT variants on progressively larger subsets of JFT-300M while keeping the number of training steps fixed. They then fit an \(\ell_2\)-regularized linear classifier to frozen representations using five labeled examples per class.

Figure 5. Linear five-shot ImageNet accuracy versus the size of the JFT pretraining subset, reproduced from the paper's figure 2 (right). This is a frozen-feature evaluation, not full fine-tuning.

On the smallest subsets, Mixer models overfit strongly. As the dataset grows, Mixer-L/16 continues to improve and the gap to ViT-L/16 shrinks. BiT plateaus earlier in this experiment. The evidence supports the narrower conclusion that Mixer benefits substantially from data scale under the tested setup; it does not establish that Mixer will necessarily overtake other architectures on still larger datasets.

What the permutation experiment actually tests🔗

To probe inductive bias, the authors train Mixer-B/16 and ResNet50x1 on three versions of JFT-300M:

1. ordinary images;
2. images with a fixed patch-order shuffle and a fixed within-patch pixel shuffle;
3. images with one fixed global pixel permutation.

The same permutation is applied to every image throughout training and evaluation. Mixer performs almost identically under the first two pipelines, whereas ResNet degrades substantially. Under global pixel permutation, both degrade, but Mixer loses less five-shot accuracy than ResNet.

This does not show that spatial structure is irrelevant. It shows that Mixer depends less than a CNN on the conventional arrangement of local neighborhoods when it can learn a consistent alternative arrangement from a very large dataset. Its token-mixing layer can directly connect distant patch positions; a CNN's local filters are built around neighborhood structure.

What the learned weights look like🔗

Figure 6. Hidden units from the first, second, and third token-mixing MLPs of Mixer-B/16 trained on JFT-300M. Each small map contains 196 weights, one for each location in the \(14\times14\) patch grid. Adapted from the paper's figure 5 and supplementary visualizations.

Some early hidden units operate over broad regions, while others are localized. The paper also finds pairs of filters with approximately opposite phases, a pattern observed in CNNs. Deeper token-mixing layers do not simply progress from general to increasingly detailed features; their weights have less clearly identifiable spatial structure.

Where MLP-Mixer fits🔗

MLP-Mixer isolates two operations that many vision systems combine: mixing information within a location and between locations. Its value is partly experimental. By removing convolution and attention while retaining competitive classification results, it helps separate what those mechanisms make convenient from what is strictly necessary for this benchmark.

The trade-offs are equally important:

• Mixer has weaker built-in assumptions about local image structure and is comparatively data-hungry.
• The original experiments focus on image classification and transfer benchmarks; they do not establish performance on detection, segmentation, language modeling, or arbitrary modalities.
• Token-mixing weights depend on the patch-grid size, so changing resolution is less direct than it is for many convolutional models.
• Global mixing is fixed after training, whereas attention adapts its mixing pattern to each input.

Richard Sutton's The Bitter Lesson offers useful context for the paper's interest in simpler methods that can exploit scale. But Mixer does not prove that architectural inductive biases are counterproductive. Its own parameter sharing, patch construction, fixed token order, normalization, and global pooling are inductive biases. The experiment asks how far a different set of biases can go.

The takeaway🔗

MLP-Mixer's central move is easy to state precisely: represent an image as a patch-by-channel table, then alternate an MLP across patches with an MLP across channels. That factorization gives every layer spatial communication and feature communication without convolution or self-attention.

The results are strongest at large scale. Mixer trades a small amount of accuracy against high measured throughput in several comparisons and becomes more competitive as the pretraining dataset grows. The careful conclusion is therefore not "MLP is all you need," but a well-structured MLP is enough to be a serious image-classification baseline.

Related reading🔗

The Vision Transformer article explains the attention-based architecture that provides the closest comparison in this paper.

For a different use of efficient matrix structure, the QLoRA article explains how quantization and low-rank adapters reduce the memory needed to fine-tune language models. MLP-Mixer is an image-classification architecture; QLoRA is a parameter-efficient training method. They solve different problems.

Original sources🔗

• Tolstikhin et al., MLP-Mixer: An all-MLP Architecture for Vision, NeurIPS 2021. The equations and results above follow the published paper.
• Google Research's official Vision Transformer and MLP-Mixer repository, including pretrained Mixer checkpoints and fine-tuning examples.

Figures are reproduced or adapted from Tolstikhin et al.'s paper and supplementary material.