PIXELBANKv8.2.1
Menu
Sign In

Complete Mixer Block with Residuals

HardNeural Networks, Architecture

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

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

Where:

  • TokenMix: Wtoken⋅XW_{\text{token}} \cdot 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:

Input:
X (2×2), W_token=zeros, W1=I, W2=I
Output:
X + 0 + GELU(X) = X + GELU(X)
Reasoning:

Zero token weights mean no token mixing. Identity channel weights apply GELU directly. Final output includes both residuals.

Constraints:

  • Input XX: Shape (N,D)(N, D)
  • WtokenW_{\text{token}}: Shape (N,N)(N, N)
  • W1W_1: Shape (D,H)(D, H) for channel MLP
  • W2W_2: Shape (H,D)(H, D) for channel MLP
Editor

Test Results

0/0
Run code to see test results.