L5.4 · LESSON 4 OF 6

Pre-norm and post-norm

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

THE IDEA

Where the normalisation goes

The original transformer applied normalisation after adding the sublayer's output: y = norm(x + F(x)) (post-norm). Many modern decoders normalise before the sublayer and leave the residual path untouched: y = x + F(norm(x)) (pre-norm). The sublayer F is the same; the arrangement differs, and with it training behaviour.

blocks4

after block	stream	length
1	[5, -1]	5.1
2	[7, -3]	7.62
3	[9, -5]	10.3
4	[11, -7]	13.04

Here F doubles its input and norm is LayerNorm (start x = [3, 1]). In pre-norm the raw stream is only ever added to, so it can grow; in post-norm every block ends with a normalised vector. Which one a real model uses changes how it trains, so trace the specific architecture.

Same sublayer, same start. Compare the two arrangements over several blocks.

Predict first. Are the original transformer equations automatically identical to a modern pre-norm decoder?

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

Trace both arrangements by hand, then write both blocks and watch the stream sizes diverge.

blocks4

after block	stream	length
1	[5, -1]	5.1
2	[7, -3]	7.62
3	[9, -5]	10.3
4	[11, -7]	13.04

pre-norm y =

post-norm y =

which keeps an un-normalised stream (pre or post)

To do: x = [3, 1], norm = LayerNorm (mean 2, std 1, so norm(x) = [1, −1]) and the sublayer doubles its input. PRE-norm: y = x + sublayer(norm(x)). Type y
To do: POST-norm: y = norm(x + sublayer(x)). x + 2x = [9, 3] (mean 6, std 3). Type y
To do: Which design leaves the raw residual stream un-normalised (its size can keep growing across layers)? Type pre or post

Part 2. Do it in code

layer_norm(x) is provided (no learned scale, eps = 0).
prenorm_block(x, sublayer): y = x + sublayer(layer_norm(x)).
postnorm_block(x, sublayer): y = layer_norm(x + sublayer(x)).
stack(x, sublayers, kind): apply the blocks in order with kind "pre" or "post" (ValueError for anything else) and return the list of streams after every block (not including the input).
stream_norm(v): the Euclidean length of the vector.

import numpy as np

def layer_norm(x):
    x = np.asarray(x, dtype=float)
    return (x - x.mean()) / x.std()

def prenorm_block(x, sublayer):
    pass

def postnorm_block(x, sublayer):
    pass

def stack(x, sublayers, kind):
    pass

def stream_norm(v):
    pass

Ctrl/⌘ + Enter runs it

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

CHECK YOUR UNDERSTANDING

Are the original transformer equations automatically identical to a modern pre-norm decoder?

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.