L5.3 · LESSON 3 OF 6

The feed-forward sublayer

Trace a transformer block, distinguish normalization variants and count a simplified block's parameters.

THE IDEA

The feed-forward sublayer

After attention has mixed information between positions, each token goes through a small MLP on its own: expand from the model width d to a larger hidden size f, apply a nonlinearity, and project back to d. This positionwise network holds most of a block's weights and applies the same matrices to every token.

A simple MLP has two matrices (d×f and f×d). Gated variants use three, multiplying a nonlinear branch by a linear one, so they have 1.5× the weights at the same f (models often shrink f to compensate).

width d4hidden size f8

gated variant (three matrices)

output = W2 · relu(W1·x) matrices: W1 d×f, W2 f×d weights = 2·4·8 = 64

The MLP is applied to each token separately: token t's output depends only on token t's vector. All mixing between positions happens in attention.

Token walk-through 1 of 4

x = [1, -2]

Step through one token. Then change d and f and switch to the gated variant.

Predict first. Does a positionwise MLP directly average neighbouring tokens?

Chapter context

A transformer combines attention with per-position feed-forward computation and residual paths. Exact ordering and normalization differ between architectures; learn a stated variant instead of assuming every transformer has an identical block.

TRY IT YOURSELF

Part 1. Do it hands-on

Count MLP weights and run one token by hand, then implement both MLP variants and a test for positionwise behaviour.

width d4hidden size f8

gated variant (three matrices)

output = W2 · relu(W1·x) matrices: W1 d×f, W2 f×d weights = 2·4·8 = 64

The MLP is applied to each token separately: token t's output depends only on token t's vector. All mixing between positions happens in attention.

Token walk-through 1 of 4

x = [1, -2]

weights, d = 4, f = 8, two matrices gated (three matrices) output for x = [1, −2]

To do: Width d = 4, hidden size f = 8. A simple MLP has two matrices (d×f and f×d), no biases. How many weights?
To do: A gated MLP has three matrices of that size (d×f twice, f×d once). How many weights?
To do: One token x = [1, −2] through the MLP: h = relu(x @ W1) with W1 = [[1, 1], [0, −1]], then h @ W2 with W2 = [[1], [1]]. Type the result

Part 2. Do it in code

ffn(X, W1, W2): the positionwise MLP relu(X @ W1) @ W2 for a matrix X of shape (tokens, d).
gated_ffn(X, Wg, Wu, Wd): (silu(X @ Wg) * (X @ Wu)) @ Wd, where silu(z) = z / (1 + exp(−z)).
ffn_params(d, f, gated=False): 2·d·f for the simple MLP, 3·d·f for the gated one (no biases).
is_positionwise(fn, X, trials=3, seed=0): True if each output row depends only on the same input row. Test it by changing one other row and checking that this row's output is unchanged.

import numpy as np

def ffn(X, W1, W2):
    pass

def gated_ffn(X, Wg, Wu, Wd):
    pass

def ffn_params(d, f, gated=False):
    pass

def is_positionwise(fn, X, trials=3, seed=0):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Does a positionwise MLP directly average neighboring tokens?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

One possible pre-norm block is

h=x+\operatorname{Attention}(\operatorname{Norm}(x)),\quad y=h+\operatorname{MLP}(\operatorname{Norm}(h)).

For standard full-width Q/K/V/O projections and an ungated d→f→d MLP, matrix weights total $4d^2+2df$ , excluding bias and normalization. GQA changes K/V projection dimensions.

For one feature vector x of width d, LayerNorm uses $\mu=\frac1d\sum_i x_i$ and $v=\frac1d\sum_i(x_i-\mu)^2$ , then normalizes $(x_i-\mu)/\sqrt{v+\epsilon}$ before its learned affine transform. RMSNorm instead scales using $\sqrt{\frac1d\sum_i x_i^2+\epsilon}$ without the same mean subtraction. Neither definition normalizes across arbitrary axes unless those axes are part of the specified layer.

Worked example

Let d=4 and f=8. Four attention matrices contain 4×4²=64 weights. The two MLP matrices contain 2×4×8=64. The simplified block has 128 matrix weights. Embedding tables, learned normalization scales, biases and any positional parameters are outside this count.

For x=[3,4] without affine transforms or epsilon, LayerNorm produces [-1,1], while RMSNorm divides by sqrt(12.5). The outputs differ because RMSNorm does not center the vector.

Reference code

The following example uses the Python standard library and does not call a remote service.

width,hidden=4,8
attention=4*width*width
mlp=2*width*hidden
assert attention == 64 and mlp == 64
assert attention+mlp == 128
query_heads,kv_heads=4,2
assert query_heads//kv_heads == 2
from math import sqrt,isclose
x=[3.,4.]
mean=sum(x)/len(x)
var=sum((v-mean)**2 for v in x)/len(x)
assert [(v-mean)/sqrt(var) for v in x] == [-1.,1.]
rms=sqrt(sum(v*v for v in x)/len(x))
assert isclose(sum((v/rms)**2 for v in x)/len(x),1)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Parameter-count formulas depend on the architecture. RMSNorm and LayerNorm do not perform identical transformations. Smaller KV caches do not mean the model has no attention or no queries.