The transformer block

Residual stream, normalization and the gated MLP

Chapter 7 of 21 · II. The transformer · about 40 minutes. Builds on: 2. Neural networks and the multilayer perceptron, 6. Attention.

Skip this chapter if

you can write the pre-norm transformer block from memory, give the formula for RMSNorm and for a SwiGLU MLP, and estimate the parameter count of a dense model from its width, depth and vocabulary.

Inspect a block’s pieces

See a block’s pieces with real numbers, starting at the output projection and first residual add of layer 0. Open in the Model Explorer.

The block and the residual stream

A transformer is a stack of identical-shaped blocks. Each block does two things to the token vectors: first it lets tokens exchange information (attention, Chapter 6), then it processes each token individually (an MLP, Chapter 2). Both are wrapped in residual connections and normalization:

$$\begin{aligned}x&\leftarrow x+\operatorname{Attention}\bigl(\operatorname{Norm}(x)\bigr)\\x&\leftarrow x+\operatorname{MLP}\bigl(\operatorname{Norm}(x)\bigr)\end{aligned}$$

Because every sublayer adds its output to $x$, there is one vector per token that runs the full height of the network, the residual stream. Layers read from it (after normalization) and write back to it. Early layers write simple features, later layers write more abstract ones, and the last layer’s stream is turned into logits. This “pre-norm” placement, with the norm inside the residual branch, trains stably at depth: show that with the norm placed after the residual (the original arrangement) gradients at initialization are large near the output and need careful warmup, while the pre-norm arrangement is well behaved.

Normalization

Left alone, activations can drift in scale from layer to layer. Layer normalization normalizes each token’s vector across its $d$ components:

$$\operatorname{LayerNorm}(x)=g\odot\frac{x-\operatorname{mean}(x)}{\sqrt{\operatorname{var}(x)+\varepsilon}}+b$$

RMSNorm drops the mean subtraction and the bias, keeping only the rescaling:

$$\operatorname{RMSNorm}(x)=g\odot\frac{x}{\sqrt{\operatorname{mean}(x^2)+\varepsilon}}$$

It is cheaper and works as well, which is why Llama, DeepSeek and gpt-oss all use it. The gain $g$ has only $d$ parameters per norm, negligible next to the matrices.

The MLP, and why it is “gated” now

Chapter 2’s MLP was $W_{down}\cdot \varphi (W_{up}\cdot x)$ with a hidden width of 4d. Current models use a gated variant from . With SiLU as the activation (called SwiGLU):

$$\operatorname{MLP}(x)=W_{\text{down}}\Bigl(\operatorname{SiLU}\bigl(W_{\text{gate}}\,x\bigr)\odot\bigl(W_{\text{up}}\,x\bigr)\Bigr)$$

There are two parallel up-projections. One passes through the activation and acts as a per-feature “gate” that multiplies the other. Three matrices instead of two means each hidden unit costs 1.5× as much, so the hidden width is reduced (about $8 / 3\cdot d$ instead of 4d) to keep parameter counts comparable: Llama 3 8B uses $d = 4096$ and hidden width 14,336 (3.5d).

What does the MLP do? argue it acts as a key–value memory: each row of the first matrix is a “key” that fires on a certain pattern in the input, and each column of the second matrix is a “value” written into the residual stream when its key fires. Much of a model’s stored factual knowledge appears to sit in these matrices. Mixture-of-experts models (Chapter 14) split the MLP rather than attention, since it holds most of the parameters and works on each token independently.

Counting parameters

With model width $d$, $L$ layers, vocabulary $V$ and full multi-head attention:

$$\begin{aligned} \text{attention per layer}&\approx4d^2\\ \text{MLP per layer}&\approx3\,d\,d_{ff}\ \text{(gated), or } 2\,d\,d_{ff}\ \text{(classic)}\\ \text{embeddings}&\approx V\,d\\ \text{total}&\approx L\,(\text{attention}+\text{MLP})+\text{embeddings} \end{aligned}$$

Attention counts $W_Q, W_K, W_V, W_O$. The embedding term doubles if the output matrix is not tied to the input one.

Use the calculator. Choose a preset and check the result against the published size; then modify the settings.

Two things to notice. First, the MLP holds most of the layer parameters: two thirds in the classic design and about 80% in Llama 3 8B. Second, grouped-query attention (Chapter 11) reduces the key/value matrices, so real models have fewer than $4d^{2}$ attention parameters per layer.

Variations you will meet

• Clamping and tweaks. The gpt-oss model card notes its SwiGLU implementation is unconventional (it clamps values and includes a residual connection) . Reading the exact formula in model code, not just the name, matters when you reproduce or analyze a model.
• Normalization of queries and keys and other stability tricks are common in recent models.
• Final norm. After the last block, one more normalization is applied before the output matrix.

Check yourself

What is the gradient path from the loss back to layer 1 in a pre-norm residual stack?

The identity path through all residual additions, plus paths through each sublayer. The identity path gives a gradient that does not shrink with depth.

Why does RMSNorm have no bias term, and what does it lose compared with LayerNorm?

It only rescales; it does not recenter. Empirically the recentering contributes little, and skipping it saves computation.

Estimate parameters: L = 32, d = 4096, d_ff = 14336 (gated), 8 KV heads of 128, 32 query heads, V = 128,256, untied.

Attention per layer: $4096\cdot4096$ ($Q$) $+\,2\cdot4096\cdot1024$ ($K$, $V$) $+\,4096\cdot4096$ ($O$) $=41.9$M. MLP: $3\cdot4096\cdot14336=176.2$M. Per layer $218.1\text{M}\times32=6.98$B. Embeddings $2\times525.3\text{M}=1.05$B. Total $\approx8.03$B.

Security link

Every token of every user flows through the same weights, and nothing in the block distinguishes trusted from untrusted input. Also, a few components of the residual stream can have extremely large values (“outlier features”) that make low-precision quantization hard (Chapter 15); a quantized copy can therefore behave measurably differently from the full-precision model you evaluated.

Papers

• Section 3 for the original block.
• Layer normalization.
• RMSNorm.
• Why pre-norm trains more easily than post-norm.
• Gated MLPs; a short paper that is worth reading in full.
• MLPs as key–value memories.
• Table 3 lists the exact dimensions of Llama 3 8B, 70B and 405B used above.