🎓 A1Learner
🔍
0 of 53 lessons completed

PyTorch LayerNorm by Hand

LayerNorm is a normalisation step that rescales each token's vector to mean zero and variance one. Learned scale and shift numbers then adapt the normalised vector for the rest of the network.

One row, centred and scaled

One row of eight numbers stands for a single token's vector. The hand steps subtract the row mean and divide by the row standard deviation.

torch.manual_seed(1)
row = torch.randn(8) * 3 + 5
mean = row.mean()
var = row.var(unbiased=False)
normed = (row - mean) / torch.sqrt(var + 1e-5)
print("Before: mean", round(mean.item(), 4), "std", round(row.std(unbiased=False).item(), 4))
print("After: mean", round(normed.mean().item(), 4), "std", round(normed.std(unbiased=False).item(), 4))

Scale and shift come after

Gamma is a learned scale and beta is a learned shift for each position in the vector. A scale of two and a shift of one move the normalised row to mean one and std two.

scaled = normed * 2.0 + 1.0
print("Scaled: mean", round(scaled.mean().item(), 4), "std", round(scaled.std(unbiased=False).item(), 4))

The library layer matches the hand steps

The library layer runs with weight one and bias zero. The hand function receives the same scale and shift as plain tensors.

def by_hand(xx, gamma, beta, eps=1e-5):
    m = xx.mean(dim=-1, keepdim=True)
    vv = xx.var(dim=-1, keepdim=True, unbiased=False)
    return (xx - m) / torch.sqrt(vv + eps) * gamma + beta

torch.manual_seed(1)
xm = torch.randn(2, 8)
lib = nn.LayerNorm(8)
diff = (by_hand(xm, torch.ones(8), torch.zeros(8)) - lib(xm)).abs().max().item()
print(f"Largest difference: {diff:.2e}")

Example

import torch
from torch import nn


def layer_norm_by_hand(x, gamma, beta, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    normalised = (x - mean) / torch.sqrt(var + eps)
    return normalised * gamma + beta


torch.manual_seed(0)
x = torch.randn(4, 128) * 3 + 5
gamma = 1 + 0.1 * torch.randn(128)
beta = 0.1 * torch.randn(128)

library = nn.LayerNorm(128)
with torch.no_grad():
    library.weight.copy_(gamma)
    library.bias.copy_(beta)

hand = layer_norm_by_hand(x, gamma, beta)
lib = library(x)
diff = (hand - lib).abs().max().item()

print("Input row mean:", round(x[0].mean().item(), 4))
print("Input row std:", round(x[0].std(unbiased=False).item(), 4))
print("Hand output, first 4 numbers:", [round(v, 4) for v in hand[0, :4].tolist()])
print("Library output, first 4 numbers:", [round(v, 4) for v in lib[0, :4].tolist()])
print(f"Largest difference: {diff:.2e}")

Key points

  • Each token's vector is normalised on its own, with no numbers taken from other tokens
  • Biased variance: the variance computed with division by the count of numbers, not the count minus one
  • eps: a tiny constant added inside the square root, set to 1e-5 in both versions here
  • The closing example's hand and library outputs differ by at most 4.77e-07

In this example, a hand function normalises four token vectors of width 128. The function subtracts each row's mean and divides by the square root of the biased variance plus eps. Learned gamma and beta values then scale and shift each row. An nn.LayerNorm layer receives the same gamma as its weight and the same beta as its bias. The first four numbers of the first row match the hand output exactly at four decimals. The largest difference across all 512 numbers is 4.77e-07. That gap is float rounding, not a different operation.

Note

nn.LayerNorm stores the scale as weight and the shift as bias, starting at one and zero

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Over which numbers does LayerNorm compute one mean?