MLP-Mixer - All-MLP architecture for vision.

Replaces attention and convolutions entirely with MLPs. Uses two types of MLP layers applied alternately: token-mixing MLPs that operate across spatial locations (patches), and channel-mixing MLPs that operate within each location independently.

Architecture

Image [batch, channels, height, width]
      |
+-----v--------------------+
| Patch Embedding           |  Split into P x P patches, linear project
+---------------------------+
      |
      v
[batch, num_patches, hidden_size]
      |
+-----v--------------------+
| Mixer Layer x N           |
|                           |
| Token Mixing:             |
|   LN -> Transpose         |
|   -> Dense(token_mlp_dim) |
|   -> GELU                 |
|   -> Dense(num_patches)   |
|   -> Transpose            |
|   + Residual              |
|                           |
| Channel Mixing:           |
|   LN -> Dense(ch_mlp_dim) |
|   -> GELU                 |
|   -> Dense(hidden_size)   |
|   + Residual              |
+---------------------------+
      |
      v
+-----v--------------------+
| LayerNorm                 |
+---------------------------+
      |
+-----v--------------------+
| Global Average Pool       |  Mean over patches
+---------------------------+
      |
      v
[batch, hidden_size]
      |
+-----v--------------------+
| Optional Classifier       |
+---------------------------+

Key Insight

Token-mixing MLPs allow communication between different spatial locations, while channel-mixing MLPs process features within each location. This separation is analogous to depthwise separable convolutions but uses fully-connected layers, achieving competitive results without attention.

Usage

# MLP-Mixer-B/16
model = MLPMixer.build(
  image_size: 224,
  patch_size: 16,
  hidden_size: 768,
  num_layers: 12,
  token_mlp_dim: 384,
  channel_mlp_dim: 3072,
  num_classes: 1000
)

# Small Mixer for CIFAR-10
model = MLPMixer.build(
  image_size: 32,
  patch_size: 4,
  hidden_size: 256,
  num_layers: 8,
  token_mlp_dim: 128,
  channel_mlp_dim: 1024,
  num_classes: 10
)

References

• "MLP-Mixer: An all-MLP Architecture for Vision" (Tolstikhin et al., NeurIPS 2021)