[Paper Review] MLP-Mixer: An all-MLP Architecture for Vision

Ilya Tolstikhin, Neil Houlsby+10 more|arXiv (Cornell University)|May 4, 2021

Advanced Neural Network Applications58 references1,444 citations

TL;DR

MLP-Mixer shows that an architecture built entirely from MLPs, with token-mixing and channel-mixing blocks, can achieve competitive image classification performance without convolutions or self-attention, given large-scale pre-training and regularization.

ABSTRACT

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.

Motivation & Objective

• Motivate scaling laws and architectures beyond CNNs and Transformers for vision.
• Propose an all-MLP architecture that separates token mixing (spatial communication) from channel mixing (feature communication).
• Demonstrate competitive accuracy versus compute cost on ImageNet and transfer tasks under large-scale pre-training and regularization.

Proposed method

• Represent input as a patches-by-channels table X of size S x C.
• Use two interleaved MLP blocks per Mixer layer: token-mixing MLP operating on columns (shared across rows) and channel-mixing MLP operating on rows (shared across columns).
• Apply LayerNorm, GELU nonlinearities, skip connections, and dropout as in standard deep nets.
• Avoid position embeddings; rely on order sensitivity of token-mixing MLPs and isotropic architecture with fixed width across layers.
• Pre-train on large datasets (ImageNet-21k, JFT-300M, etc.) with regularization (RandAugment, mixup, dropout, stochastic depth) and fine-tune on downstream tasks.

Experimental results

Research questions

• RQ1 Can an architecture built purely from MLPs achieve competitive accuracy in vision tasks without convolutions or self-attention?
• RQ2 How does the performance of MLP-Mixer scale with model size and pre-training data volume?
• RQ3 What are the trade-offs between accuracy, pre-training cost, and test-time throughput for Mixer compared to CNNs and ViTs?
• RQ4 How does the model’s inductive bias (per-location channel mixing vs. cross-location token mixing) affect robustness to input permutations and data efficiency?

Key findings

• MLP-Mixer achieves competitive ImageNet top-1 accuracy when pre-trained on large datasets (e.g., ~100M images) and regularized, with pre-training and inference costs comparable to state-of-the-art models.
• The architecture’s token-mixing and channel-mixing MLPs enable information exchange across spatial locations and channels, respectively, yielding performance similar to CNNs and Transformers on a range of transfer tasks.
• With larger upstream data, Mixer performance improves significantly and can approach or slightly exceed some CNN/Transformer baselines, while maintaining favorable compute/throughput characteristics.
• Mixer demonstrates invariance to patch/order permutations and robust performance under global pixel permutations, highlighting its different inductive biases compared with traditional CNNs.
• At large scales, Mixer-H/14 achieves near- ViT-H/14 performance on ImageNet while running substantially faster, indicating a favorable accuracy-versus-compute frontier.