Complete Mixer Block with Residuals

hardPROMLP-MixerNeural Networks, ArchitectureAdvanced Architectures

A complete MLP-Mixer block combines token mixing and channel mixing with residual connections:

$X_1 = X + \text{TokenMix}(X)$ $X_2 = X_1 + \text{ChannelMix}(X_1)$

Where:

TokenMix: $W_{\text{token}} \cdot X$ (mix across N tokens)
ChannelMix: $\text{GELU}(X_1 W_1)\, W_2$ , a two-layer MLP with no biases applied per-token (use the tanh-approximate GELU defined in the starter code)

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

Input

X (2×2), W_token=zeros, W1=I, W2=I

Output

X + 0 + GELU(X) = X + GELU(X)

Reasoning

Inputs: $X$ is a $2 \times 2$ matrix of 2 tokens (rows) by 2 channels (columns), here $\begin{bmatrix}1 & 2\\ 3 & 4\end{bmatrix}$ . $W_{\text{token}}$ is the $2 \times 2$ zero matrix, and both channel weights $W_1, W_2$ are the $2 \times 2$ identity, so the hidden width is $H = D = 2$ .
Token mixing: each new token row is a weighted blend of all token rows, $\sum_j w_{ij}\, x_j$ . With every weight zero the blend is the zero matrix, so $X_1 = X + 0 = X$ .
Channel mixing: each token row is multiplied by $W_1$ , passed through GELU entry by entry, then multiplied by $W_2$ . With identity weights, ChannelMix $(X_1)$ is just GELU of every entry: $\text{GELU}(1) \approx 0.8412$ , $\text{GELU}(2) \approx 1.9546$ , $\text{GELU}(3) \approx 2.9964$ , $\text{GELU}(4) \approx 3.9999$ .
Output: $X_2 = X_1 + \text{GELU}(X_1)$ , a $2 \times 2$ matrix like the input: [[1.8412, 3.9546], [5.9964, 7.9999]] (rounded to 4 decimal places).
Notice: for large positive $x$ , $\text{GELU}(x) \approx x$ , so entries approach $2x$ . The residuals keep $X$ even when a stage contributes nothing.

Constraints

Input $X$ : Shape $(N, D)$
$W_{\text{token}}$ : Shape $(N, N)$
$W_1$ : Shape $(D, H)$ for channel MLP
$W_2$ : Shape $(H, D)$ for channel MLP

Learn this in

Transformer ArchitectureLLM Study Plan

More from Advanced Architectures

Channel Mixing with GELUmedium Token Mixing Layermedium Bidirectional Scan for Visionmedium Discretize Continuous State Space Modelhard State Space Model vs Attention Complexitymedium

Concept primerPREVIEW

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)

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

Example

Constraints

Learn this in

More from Advanced Architectures

Complete Mixer Block with Residuals

Background Knowledge

MLP-Mixer Block Structure

Why Residual Connections?

The Complete Forward Pass

Continue the full explanation

Editor locked