🎓 A1Learner
🔍
0 of 87 lessons completed

PyTorch Layer Normalization

Layer normalization is a layer that rescales each row to mean 0 and std 1 using that row alone. No other row in the batch takes part in the calculation.

A row, on its own

A row's own values are enough to rescale it. BatchNorm cannot run this calculation, because it needs the whole column.

row = torch.tensor([1.0, 2.0, 3.0, 4.0])
# its mean is 2.5, whatever the other rows hold

Two rows, two scales, one result

The second row is the first row times ten. Both rows come out identical, because each row is rescaled by its own mean and std.

import torch

norm = torch.nn.LayerNorm(3)
x = torch.tensor([[1.0, 2.0, 3.0],
                  [10.0, 20.0, 30.0]])
out = norm(x)
print([[round(v, 4) for v in row] for row in out.tolist()])

One row is enough

A single row works in both modes, and the two results match exactly. No batch is needed at any point.

import torch

norm = torch.nn.LayerNorm(3)
row = torch.tensor([[2.0, 4.0, 6.0]])
norm.train()
first = norm(row)
norm.eval()
second = norm(row)
print(f"Same for one row: {torch.equal(first, second)}")

Example

import torch

row = torch.tensor([1.0, 2.0, 3.0, 4.0])
mean = row.mean()
var = ((row - mean) ** 2).mean()
by_hand = (row - mean) / torch.sqrt(var + 1e-5)

norm = torch.nn.LayerNorm(4)
by_layer_norm = norm(row)

hand_text = ", ".join(f"{v:.4f}" for v in by_hand.tolist())
norm_text = ", ".join(f"{v:.4f}" for v in by_layer_norm.tolist())
print(f"By hand: [{hand_text}]")
print(f"LayerNorm: [{norm_text}]")
print(f"Equal: {torch.allclose(by_hand, by_layer_norm)}")

Key points

  • Row: one example's full set of values at a layer
  • Use torch.nn.LayerNorm(size) to normalise each row over its last dimension
  • Variance: the mean squared distance from the mean, the square of the std
  • Eps: a tiny number added inside the square root, so a flat row never divides by zero
  • LayerNorm holds no running stats, because each row supplies its own mean and variance
  • Train mode and eval mode give the same result, because the batch never enters the maths
  • Layer normalization pins each row's own scale, which is why transformers use it

In this example, one row holds the values 1, 2, 3, and 4. The by-hand lines subtract the row mean and divide by the square root of the variance plus eps. LayerNorm runs the same calculation in one call. Both lines print the same four values, from -1.3416 to 1.3416. The final line confirms the two results are close enough to count as equal. One row, rescaled by its own numbers, needs no batch and no stored stats.

Note

BatchNorm normalises down each column, while LayerNorm normalises across each row.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does LayerNorm use to rescale one row?