Token Mixing Layer

mediumFREEMLP-MixerNeural Networks, Matrix MathAdvanced Architectures

MLP-Mixer replaces self-attention with simple MLPs that mix information in two stages:

Token Mixing: Mix information across spatial tokens (patches)
Channel Mixing: Mix information across feature channels

Token Mixing operates on the transposed representation: $Y = W_{\text{token}} \cdot X$

Where:

$X \in \mathbb{R}^{N \times D}$ : N tokens with D channels each
$W_{\text{token}} \in \mathbb{R}^{N \times N}$ : Learnable token mixing weights

This mixes information across the N spatial locations while keeping channels independent.

Task: Implement the token mixing forward pass.

Example

Input

X (3×2), W_token = identity (3×3)

Output

X unchanged (identity mixing)

Reasoning

With identity weights, each token keeps only its own information. Non-identity weights blend information across tokens.

Constraints

$N$ (Number of tokens): $1 \leq N \leq 256$
$D$ (Channels per token): $1 \leq D \leq 512$

More from Advanced Architectures

Complete Mixer Block with Residualshard Channel Mixing with GELUmedium Bidirectional Scan for Visionmedium Discretize Continuous State Space Modelhard State Space Model vs Attention Complexitymedium

Concept primer

Token Mixing Layer (MLP-Mixer)

Background Knowledge

MLP-Mixer (Google Brain, 2021) demonstrated that neither convolutions nor attention are necessary for competitive image classification. It uses only MLPs to mix information across tokens and channels.

Architecture Overview

MLP-Mixer alternates between two mixing operations:

Token Mixing: Mix information across spatial locations (patches)
Channel Mixing: Mix information across feature channels

Token Mixing Intuition

Token mixing answers: "How should different spatial locations share information?"

For an input $X \in \mathbb{R}^{N \times C}$ (N patches, C channels):

X_transposed = X.T # Shape: (C, N)
X_mixed = W @ X_transposed # Shape: (C, N), W is (N, N)
output = X_mixed.T # Shape: (N, C)

This is equivalent to: the same MLP applied to each channel's spatial arrangement.

Algorithm/Approach

def token_mixing(X, W):
 # X: (num_patches, channels)
 # W: (num_patches, num_patches)
 return (W @ X.T).T # Or equivalently: X @ W.T

Key insight: The weight matrix $W$ is shared across all channels.

Step-by-Step Strategy

Transpose input to shape (C, N)
Apply weight matrix W (N × N) via matrix multiplication
Transpose back to (N, C)

Common Pitfalls

Confusing token mixing (spatial) with channel mixing (feature)
Incorrect transpose order
Wrong matrix multiplication order (W @ X vs X @ W)

Time & Space Complexity

Time: $O(N^2 \cdot C)$ - N×N matrix times N×C input
Space: $O(N^2)$ for the weight matrix