L5.2 · LESSON 2 OF 6

LayerNorm and RMSNorm

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

THE IDEA

Keeping the numbers in a usable range

Residual streams can drift to very different scales from layer to layer. A normalisation layer rescales each token's vector before a sublayer sees it. There are two common versions:

LayerNorm subtracts the vector's mean over the features, then divides by their standard deviation (plus a tiny epsilon). The result has mean 0 and spread 1.
RMSNorm only divides by the root-mean-square. It skips the mean subtraction, so it is a little cheaper, and the result generally does not have mean 0.

Both then usually multiply by a learned scale; details such as a learned bias vary by implementation.

x[0]1x[1]2x[2]3x[3]4

feature	x	LayerNorm	RMSNorm
0	1	-1.342	0.365
1	2	-0.447	0.73
2	3	0.447	1.095
3	4	1.342	1.461

mean(x) = 2.5 rms(x) = 2.739 LayerNorm output: mean 0, spread 1 (centres AND scales) RMSNorm output: mean 0.913, root-mean-square 1 (scales only, no centring)

Move the features. Note LayerNorm's mean stays 0 while RMSNorm's does not.

Predict first. Does RMSNorm ordinarily subtract the feature mean?

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

Normalise a vector both ways by hand, then implement LayerNorm and RMSNorm over the last axis.

x[0]1x[1]2x[2]3x[3]4

feature	x	LayerNorm	RMSNorm
0	1	-1.342	0.365
1	2	-0.447	0.73
2	3	0.447	1.095
3	4	1.342	1.461

mean(x) = 2.5 rms(x) = 2.739 LayerNorm output: mean 0, spread 1 (centres AND scales) RMSNorm output: mean 0.913, root-mean-square 1 (scales only, no centring)

mean of [1,2,3,4] LayerNorm first output RMSNorm first output

To do: Normalise x = [1, 2, 3, 4]. First, the mean of the features
To do: LayerNorm: (x − mean) / std (std = √1.25 = 1.118, no epsilon). Type the first output (3 decimals)
To do: RMSNorm: x / rms, rms = √(mean of squares) = √7.5 = 2.7386. It does NOT subtract the mean. Type the first output (3 decimals)

Part 2. Do it in code

layer_norm(x, gamma=None, beta=None, eps=1e-5): normalise over the last axis: (x − mean) / sqrt(variance + eps), where variance is the mean squared deviation (not divided by n−1). Then multiply by gamma and add beta if they are given.
rms_norm(x, gamma=None, eps=1e-6): x / sqrt(mean(x²) + eps) over the last axis (no mean subtraction), times gamma if given.
stats(x): return (mean, rms) of the last axis of x as two arrays.

import numpy as np

def layer_norm(x, gamma=None, beta=None, eps=1e-5):
    pass

def rms_norm(x, gamma=None, eps=1e-6):
    pass

def stats(x):
    pass

Ctrl/⌘ + Enter runs it

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

CHECK YOUR UNDERSTANDING

Does RMSNorm ordinarily subtract the feature mean?

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.