Channel Mixing with GELU
Channel Mixing in MLP-Mixer applies an MLP independently to each token's features:
Y=W2⋅GELU(W1⋅XT)T
Or equivalently per-token: yi=W2⋅GELU(W1⋅xi)
Where:
- X∈RN×D: Input tokens
- W1∈RD×H: Expand to hidden dim H
- W2∈RH×D: Project back to D
- GELU: Gaussian Error Linear Unit activation
GELU approximation: GELU(x)≈0.5x(1+tanh(2/π(x+0.044715x3)))
Task: Implement channel mixing with GELU activation.
Example:
X (2×2), W1 = I (2×4), W2 = I (4×2)
GELU applied element-wise to X
With identity-like projections, the MLP reduces to applying GELU to each element. GELU(1) ≈ 0.8412, GELU(0) = 0.
Constraints:
- N: Number of tokens
- D: Input/output dimension
- H: Hidden dimension (typically 4×D)
Channel Mixing with GELU
Background Knowledge
Channel Mixing is the second core operation in MLP-Mixer. While token mixing shares information across spatial locations, channel mixing transforms features at each location independently.
Channel Mixing in MLP-Mixer
For each patch independently:
hidden = GELU(x @ W1) # Expand: C → 4C typically
output = hidden @ W2 # Contract: 4C → C
This is identical to the Feed-Forward Network (FFN) in Transformers!
GELU Activation Function
Gaussian Error Linear Unit is defined as:
GELU(x) = x × Φ(x)
Where Φ(x) is the CDF of standard Gaussian. Approximate formula:
GELU(x) ≈ 0.5x(1 + tanh(√(2/π)(x + 0.044715x³)))
GELU is smoother than ReLU and has become standard in BERT, GPT, and ViT.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.