Channel Mixing with GELU

mediumPROMLP-MixerNeural Networks, Activation FunctionsAdvanced Architectures

Channel Mixing in MLP-Mixer applies an MLP independently to each token's features:

$Y = W_2 \cdot \text{GELU}(W_1 \cdot X^T)^T$

Or equivalently per-token: $y_i = W_2 \cdot \text{GELU}(W_1 \cdot x_i)$

Where:

$X \in \mathbb{R}^{N \times D}$ : Input tokens
$W_1 \in \mathbb{R}^{D \times H}$ : Expand to hidden dim H
$W_2 \in \mathbb{R}^{H \times D}$ : Project back to D
GELU: Gaussian Error Linear Unit activation

GELU approximation: $\text{GELU}(x) \approx 0.5x(1 + \tanh(\sqrt{2/\pi}(x + 0.044715x^3)))$

Task: Implement channel mixing with GELU activation.

Example

Input

X (2×2), W1 = I (2×4), W2 = I (4×2)

Output

GELU applied element-wise to X

Reasoning

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)

More from Advanced Architectures

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

Concept primerPREVIEW