Inside the model · MLP

MLP

feed-forward · block 1

shown
toy scale
tokens
5
dmodel
8 GPT-2 768
dff
32 GPT-2 3,072
params
552 GPT-2 4.7M

All steps

  1. Up-projection

    Attention moved information between tokens; the MLP works on each token alone, and is thought to hold much of what the model has learned. X (after ln2) times Wfc plus bfc widens every token from 8 numbers to 32 (768 → 3,072 in GPT-2): each column is one neuron.

    GPT-2 [N×768]·[768×3072] · 4.7 MFLOPs / token

  2. GELU

    GELU bends every cell on its own: large positive inputs pass almost unchanged, small ones are damped, negative ones are squeezed toward 0 (never below −0.17). Without this bend, Wfc then Wproj would collapse into one 768 × 768 matrix.

    elementwise · 5 × 32 cells

  3. Down-projection

    The 32 activations are projected back down to 8 (3,072 → 768), and the bias bproj is added. Wproj is drawn transposed, so each output dimension is a row sitting directly above the activations it multiplies.

    GPT-2 [N×3072]·[3072×768] · 4.7 MFLOPs / token

  4. Residual add

    The MLP output is added to the residual stream. Positions never exchanged information in this sub-layer; each token was processed on its own. Next is block 2.

    h″ = h′ + MLP(ln2(h′))

Code

def forward(self, x):          # x = ln_2(h): (B, T, 768)
    x = self.c_fc(x)           # (B, T, 3072)
    x = self.gelu(x)           # GPT-2 used the tanh approximation
    x = self.c_proj(x)         # (B, T, 768)
    x = self.dropout(x)
    return x
# in the block: x = x + self.mlp(self.ln_2(x))

Go deeper