The MLP and SwiGLU — Linear algebra for LLMs

The classic MLP

\displaystyle \operatorname{MLP}(x) = \operatorname{GELU}(xW_{\text{up}})\,W_{\text{down}}

→ →

Widen each token to $d_{\text{ff}} = 4d$ , apply a nonlinearity entry by entry, and come back down to $d$ . GPT-2 small: $768 \to 3072 \to 768$ . The nonlinearity is what stops the two matrices from collapsing into one (Part 3).

Write

\operatorname{MLP}(x) = \operatorname{act}(xW_{\text{up}})W_{\text{down}}

,

d \to d_{\text{ff}} \to d

.

Its parameters

\displaystyle 2 \times d \times d_{\text{ff}} = 2 \times d \times 4d = 8d^2

Twice the attention’s $4d^2$ . Together they make the classic $12d^2$ per block that the counting deck uses.

SwiGLU: add a gate

\displaystyle \operatorname{SwiGLU}(x) = \big(\operatorname{SiLU}(xW_{\text{gate}}) \odot xW_{\text{up}}\big)\,W_{\text{down}}

\displaystyle \operatorname{SiLU}(z) = z\,\sigma(z)

Two up-projections instead of one. One goes through SiLU and then multiplies the other entry by entry (the $\odot$ of Part 2), so each hidden unit can switch the other on or off. Llama calls the three matrices gate_proj, up_proj and down_proj.

Write

(\operatorname{SiLU}(xW_g) \odot xW_u)\,W_d

: three matrices, one elementwise gate.

Shapes through SwiGLU

→ gate, through SiLU → up⊙ → down

Parameters, and the 8/3 rule

\displaystyle 3 \times d \times d_{\text{ff}}

\displaystyle 3\,d\,d_{\text{ff}} = 8d^2 \;\Rightarrow\; d_{\text{ff}} = \tfrac{8}{3}d

Three matrices instead of two, so to keep the same parameter count as a $4d$ MLP, the hidden width drops to $\tfrac{8}{3}d$ . LLaMA’s 7B models used $d_{\text{ff}} = 11008$ for $d = 4096$ ( $\tfrac{8}{3} \cdot 4096 \approx 10923$ , rounded up). Llama 3 8B goes wider: 14,336.

Write

SwiGLU params:

3\,d\,d_{\text{ff}}

. Same as a

4d

MLP when

d_{\text{ff}} = \tfrac83 d

.

✏️ Your turn

Size Llama 3 8B's MLP

\displaystyle d = 4096, \quad d_{\text{ff}} = 14336

How many parameters in one layer’s MLP?

Work it out on paper, then go to the next card to check.

Answer

\displaystyle 3 \times 4096 \times 14336 = 176{,}160{,}768

About 176 million per layer. The attention in the same layer (with GQA) has about 42 million: the MLP holds most of the model.

In PyTorch

h = F.silu(x @ W_gate.T) * (x @ W_up.T)    # (B, T, d_ff)
y = h @ W_down.T                           # (B, T, d)

📌 Recap

The MLP and SwiGLU in one card

Classic: $\operatorname{act}(xW_{\text{up}})W_{\text{down}}$ , $d_{\text{ff}} = 4d$ , $8d^2$ params.
SwiGLU: $(\operatorname{SiLU}(xW_g) \odot xW_u)W_d$ , $3\,d\,d_{\text{ff}}$ params.
$d_{\text{ff}} = \tfrac83 d$ keeps the count equal; Llama 3 8B uses 14,336.
The MLP holds most of a transformer’s parameters.

Write

The four lines above.