Normalization layers

Which numbers get averaged

The recipe needs a group of numbers to take the mean and variance of. The two main norm layers differ only in which group they pick.

Inside a network, activations come in a grid. Each row is one example. Each column is one feature. The grid's shape is written (batch, features). Think of a class report card: one row per student, one column per subject.

• BatchNorm (batch normalization) works down each column. It compares each student with the rest of the class, one subject at a time. "How good is Ana's math compared with everyone's math?"
• LayerNorm (layer normalization) works across each row. It compares each student's subjects with their own other subjects. "Is math Ana's best subject?"

The "eval" column is about using a trained model, which comes up in a moment. CNNs are image models. RNNs are older models for text and other sequences. You'll meet both later in this phase.

Here are both layers in PyTorch. The batch has 4 examples and 3 features, and the features have very different sizes.

import torch
import torch.nn as nn

torch.manual_seed(0)
x = torch.randn(4, 3) * torch.tensor([1.0, 10.0, 100.0]) + 5      # 4 examples, 3 very different features
print("raw feature means:", x.mean(dim=0))
bn, ln = nn.BatchNorm1d(3), nn.LayerNorm(3)
b, l = bn(x), ln(x)
print("after BatchNorm, each FEATURE (column):  mean", b.mean(dim=0).detach().numpy().round(4) + 0.0, " std", b.std(dim=0, unbiased=False).detach().numpy().round(3))
print("after LayerNorm, each EXAMPLE (row):     mean", l.mean(dim=1).detach().numpy().round(4) + 0.0, " std", l.std(dim=1, unbiased=False).detach().numpy().round(3))

(As always here, torch is the course's in-browser PyTorch subset; the same code runs on real PyTorch.)

The first line shows the raw features have very different averages. After BatchNorm, each column has a mean of 0 and a spread of 1. After LayerNorm, each row does.

import statistics as st

# math, art, music test scores for three students
scores = {"Ana": [90, 60, 72], "Ben": [70, 80, 66], "Cy": [50, 70, 60]}

def standardize(values):
    mu, sd = st.mean(values), st.pstdev(values)
    return [round((v - mu) / sd, 2) for v in values]

print("LayerNorm style: each student against their own subjects")
for name, row in scores.items():
    print(" ", name, standardize(row))

print("BatchNorm style: each subject across the class")
for i, subject in enumerate(["math", "art", "music"]):
    column = [row[i] for row in scores.values()]
    print(" ", subject, standardize(column))

A network that only trains with it

Here's the payoff. The code builds a deep network: eight layers, each followed by tanh. It trains on the handwritten digits for five passes over the data (five epochs). It uses mini-batch SGD, the training method from the mini-batches lesson. It does this three times: with no norm layers, with BatchNorm, and with LayerNorm.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
y = torch.tensor(digits.target)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=0)

def deep_net(norm=None, depth=8, width=64):
    layers, n_in = [], 64
    for _ in range(depth):
        layers += [nn.Linear(n_in, width)] + ([norm(width)] if norm else []) + [nn.Tanh()]
        n_in = width
    return nn.Sequential(*layers, nn.Linear(width, 10))

loss_fn = nn.CrossEntropyLoss()
for name, norm in [("no normalization", None), ("BatchNorm1d", nn.BatchNorm1d), ("LayerNorm", nn.LayerNorm)]:
    torch.manual_seed(0)
    model = deep_net(norm)
    opt = torch.optim.SGD(model.parameters(), lr=0.05)
    for epoch in range(5):
        model.train()
        for xb, yb in DataLoader(TensorDataset(X_train, y_train), batch_size=64, shuffle=True):
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            opt.step()
    model.eval()
    with torch.no_grad():
        acc = (model(X_val).argmax(1) == y_val).float().mean().item()
        print(f"{name:<17} training loss {loss_fn(model(X_train), y_train).item():.3f}   val accuracy {acc:.3f}")

Look at the training loss, which is the error on the training data, for each version. A model that spreads its bets evenly over the 10 digits scores a loss of about 2.3. That's the natural log of 10, written ln 10, and it means pure guessing. Without normalization, the deep network's loss is still about there. It hasn't started learning.

With either norm layer, it learns quickly. Norm layers also let you use larger learning rates, and they make training less sensitive to the starting weights you picked.

Batch norm's two modes

BatchNorm has a catch. Its mean and variance come from the current batch. That's fine during training, when the data arrives in batches anyway. But once the model is trained, you want each prediction to depend only on its own input, not on whatever else happens to share its batch.

So BatchNorm keeps a running average while it trains. Think of a café owner estimating a typical day's customers. They don't recount every day since opening. They keep a rough figure and nudge it a little after each new day. BatchNorm does the same with its mean and variance after every training batch.

Calling model.eval() switches the model into evaluation mode. In that mode, BatchNorm uses the stored averages instead of the batch's own numbers. Two things follow:

• Forgetting model.eval() at test time gives different (usually worse) predictions.
• Very small batches give noisy statistics. So BatchNorm struggles with batch sizes of 1 or 2.

Why transformers use layer norm

Text brings two problems for BatchNorm. A batch can mix sentences of very different lengths. And when a model writes text, it often handles one sequence at a time, a batch of 1.

LayerNorm has neither problem. In a language model, which is a model that predicts the next token, each token (a word or piece of a word) has its own list of features. LayerNorm standardizes each token's features on their own, without looking at the rest of the batch. Because nothing comes from the batch, LayerNorm behaves the same in training and in use, at any batch size.

Two refinements in modern LLMs

Nearly every modern large language model (LLM) adds two changes to plain LayerNorm.

Pre-norm: normalize first

A transformer block has a few sub-layers. Each one is a chunk of the network, such as an MLP or the attention step you'll meet in Phase 5. Call a sub-layer $f$. Pre-norm places the norm layer before it:

In plain words: standardize the numbers, run them through the sub-layer, then add the result back onto the original $x$. The original $x$ itself is never standardized. The sub-layer's output is simply added on top.

A small example. Take $x$ = 2, 4, 4, 10 again. Standardized, that's −1, −0.333, −0.333 and 1.667. Pretend the sub-layer triples each number: −3, −1, −1 and 5. Add those to the original 2, 4, 4, 10 and you get −1, 3, 3 and 15.

import math

def norm(v):
    mu = sum(v) / len(v)
    sd = math.sqrt(sum((a - mu) ** 2 for a in v) / len(v))
    return [(a - mu) / sd for a in v]

def f(v):                       # a stand-in for an attention or MLP sub-layer
    return [3 * a for a in v]

x = [2, 4, 4, 10]
out = [a + b for a, b in zip(x, f(norm(x)))]   # pre-norm: x + f(norm(x))
print([round(v, 3) for v in out])

Pre-norm trains more stably in deep stacks. You'll see why in the lesson on residual connections, which are the shortcuts that let a signal skip a block.

RMSNorm: skip the mean

RMSNorm drops the "subtract the mean" step. It only divides by the root mean square (RMS). Read the name from right to left: square each number, take the mean, then take the square root. Then it multiplies by a learned scale, $\gamma$, as before.

In plain words: work out the typical size of the numbers, then divide every number by it. Then stretch by the learned scale $\gamma$. Nothing gets re-centered, so the mean doesn't become 0.

With 2, 4, 4 and 10, the squares are 4, 16, 16 and 100. Their mean is 136 ÷ 4 = 34, and √34 ≈ 5.83. Dividing gives about 0.343, 0.686, 0.686 and 1.715. Compare LayerNorm's −1, −0.333, −0.333 and 1.667. The RMSNorm values are all still positive.

import math

def rms_norm(v, eps=1e-6):
    rms = math.sqrt(sum(a * a for a in v) / len(v) + eps)
    return [round(a / rms, 3) for a in v]

def layer_norm(v, eps=1e-5):
    mu = sum(v) / len(v)
    var = sum((a - mu) ** 2 for a in v) / len(v)
    return [round((a - mu) / math.sqrt(var + eps), 3) for a in v]

x = [2, 4, 4, 10]
print("RMSNorm:  ", rms_norm(x))
print("LayerNorm:", layer_norm(x))

RMSNorm is slightly cheaper to compute and works as well in practice. Llama, Mistral and many other models use it.

import math
x = [1, 7]
rms = math.sqrt((1 * 1 + 7 * 7) / 2)
print(rms, [round(v / rms, 2) for v in x])