Complete Mixer Block with Residuals
A complete MLP-Mixer block combines token mixing and channel mixing with residual connections:
X1=X+TokenMix(X) X2=X1+ChannelMix(X1)
Where:
- TokenMix: Wtoken⋅X (mix across N tokens)
- ChannelMix: MLP with GELU applied per-token
The residual connections ensure stable training and allow the network to learn identity mappings when beneficial.
Task: Implement a complete Mixer block with both mixing stages and residual connections.
Example:
X (2×2), W_token=zeros, W1=I, W2=I
X + 0 + GELU(X) = X + GELU(X)
Zero token weights mean no token mixing. Identity channel weights apply GELU directly. Final output includes both residuals.
Constraints:
- Input X: Shape (N,D)
- Wtoken: Shape (N,N)
- W1: Shape (D,H) for channel MLP
- W2: Shape (H,D) for channel MLP
Complete Mixer Block with Residuals
Background Knowledge
A Mixer Block combines token mixing and channel mixing with residual connections and layer normalization. This is the fundamental building block of MLP-Mixer.
MLP-Mixer Block Structure
Input X
│
├──→ LayerNorm → Token Mixing ──┐
│ │
└────────────────────────────────+ (residual)
│
├──→ LayerNorm → Channel Mixing ──┐
│ │
└──────────────────────────────────+ (residual)
│
Output
Why Residual Connections?
Residuals enable:
- Gradient flow: Direct paths for gradients during backprop
- Identity mapping: Network can learn identity if needed
- Deeper networks: Enable training of 12+ layer models
The Complete Forward Pass
def mixer_block(X, W_t, W_c1, W_c2):
# Token mixing with residual
Y = X + token_mixing(layer_norm(X), W_t)
# Channel mixing with residual
Z = Y + channel_mixing(layer_norm(Y), W_c1, W_c2)
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.