Concept units

Layer normalization

Intermediate

Read these first if they are new to you

The idea in one paragraph

Layer normalization takes one token's numbers, the d entries of its vector, and shifts and rescales them to mean 0 and standard deviation 1. If the numbers are all nearly equal, the standard deviation comes out below 1. Layer normalization then multiplies each entry by a learned scale and adds a learned shift. It treats every token on its own, so the result does not depend on the other tokens or on the other examples in a batch. Networks use it to keep the numbers that flow from layer to layer at a steady size.

Layer normalization works across one token's numbers

Lay out a group of tokens as a table: one row per token, one column per feature. A token is the vector a network holds for one item, such as one point of a cloud; each of its numbers is a feature. The normalizing data unit normalizes a column: one feature, across many examples. Layer normalization normalizes a row: all the features of one token.

Normalizing a column versus normalizing a row
layer norm: one token, all its features normalizing data: one feature features → token 2
What to notice. Each cell is one number. Layer normalization uses only the numbers in one row to compute that row's mean and spread, so every token gets its own.

The name comes from its first use: it normalized all the outputs of one layer for one input. In a transformer, those outputs are exactly one token's features.

The recipe: subtract the mean, divide by the spread, then scale and shift

For a token x with d numbers:

  1. Find the mean: μ = (x1 + … + xd) / d.
  2. Find the variance: σ2 = ((x1 − μ)2 + … + (xd − μ)2) / d.
  3. Normalize each number: x̂i = (xi − μ) / √(σ2 + ε).
  4. Scale and shift: yi = γi x̂i + βi.

The small number ε (epsilon), usually 0.00001, prevents a division by zero when all the numbers are equal. The scale γ (gamma) and shift β (beta) are learned vectors of d numbers, one pair per feature. At the start of training, every entry of γ is 1 and every entry of β is 0, written γ = 1 and β = 0. Training then adjusts them, like the weights of a linear layer. They let the network give each feature the size and offset that suit the next layer.

Worked example: normalize the token [2, 4, 6, 8]
  1. Mean: (2 + 4 + 6 + 8) / 4 = 5.
  2. Differences from the mean: [−3, −1, 1, 3]. Variance: (9 + 1 + 1 + 9) / 4 = 5.
  3. Spread: √5 ≈ 2.236. Divide: [−3, −1, 1, 3] / 2.236 = [−1.342, −0.447, 0.447, 1.342]. (Adding ε = 0.00001 changes only the sixth decimal place.)
  4. With γ = 1 and β = 0, these are the outputs.

The outputs have mean 0 and standard deviation 1, as promised.

Predict first. What does layer normalization give for the token [20, 40, 60, 80]? And for [102, 104, 106, 108]?

Both give [−1.342, −0.447, 0.447, 1.342] again. The first is the worked example times 10; the second is the worked example plus 100. Normalization undoes a common offset exactly, and a common scale by a positive number nearly.

The scale cancels only nearly because of ε. Scaling the token by k multiplies its differences from the mean by k and its variance by k2, but leaves ε the same. So k cancels while the variance is much larger than ε, which means a spread well above √ε, about 0.003. The worked example times 0.001 has spread 0.0022, and it gives [−0.775, −0.258, 0.258, 0.775].

Shifting a token leaves its normalized values unchanged, and scaling it by a positive number nearly does

A token before and after layer normalization
The token: k × [2, 4, 6, 8] + c After layer normalization mean 5.00, spread 2.24 mean 0.00, spread 1.00

Try it. Before you move a slider, predict what the right-hand bars will do. Then search: is there any setting of k and c that changes the right-hand bars?

Answer

Only two kinds of setting change them visibly. A negative k flips the token, so the normalized bars flip too. At k = 0 all four numbers are equal, the spread is 0, and every output is 0; this is where ε prevents a division by zero. Every other setting changes them by less than 0.0001, too little to see. Even at k = 0.25, the spread is 0.56, far above √ε.

What to notice. The left bars show the token's numbers on a fixed scale; the right bars show the same token after layer normalization, with γ = 1 and β = 0.

Networks use it to keep each layer's inputs at a steady size

As numbers pass through many layers, their size can drift: some tokens' numbers grow large, others shrink. The next layer then sees inputs of very different sizes from step to step, which makes training slow and unstable. Layer normalization resets each token to a standard size before the next layer reads it. This tends to make training faster and more stable.

In a transformer, a layer normalization sits next to each attention and MLP sublayer. Some designs put it before the sublayer, others after the residual addition; the transformer block unit compares the two. Some models compute γ and β from another input, such as the time step of a generative model, instead of learning fixed values (see time conditioning). A cheaper variant, RMS normalization, skips subtracting the mean and divides by the root of the mean square.

layer_norm.py

import numpy as np

def layer_norm(x, gamma, beta, eps=1e-5):  # 1e-5 is 0.00001
    # axis=-1 averages the numbers of each token (one row). keepdims=True keeps
    # each mean in its token's row, so x - mu subtracts the right mean from each token.
    mu = x.mean(axis=-1, keepdims=True)                  # each token's mean
    var = ((x - mu) ** 2).mean(axis=-1, keepdims=True)   # and its variance
    return gamma * (x - mu) / np.sqrt(var + eps) + beta

tokens = np.array([[2.0, 4.0, 6.0, 8.0],     # one token per row
                   [20.0, 40.0, 60.0, 80.0],
                   [5.0, 5.0, 5.0, 5.0]])
# np.ones(4) holds four 1s, np.zeros(4) four 0s, and np.full(4, 2.0) four 2s.
print(np.round(layer_norm(tokens, np.ones(4), np.zeros(4)), 3))
print(np.round(layer_norm(tokens[0], np.full(4, 2.0), np.ones(4)), 3))

[[-1.342 -0.447  0.447  1.342]
 [-1.342 -0.447  0.447  1.342]
 [ 0.     0.     0.     0.   ]]
[-1.683  0.106  1.894  3.683]

PyTorch provides the same operation as torch.nn.LayerNorm(d), with ε = 0.00001 and γ, β learned.

Two details this page skips

Layer normalization throws away each token's overall offset, and nearly all of its overall size. Designs that need that information keep an unnormalized copy, which is one job of the residual path in a transformer block. A different method, batch normalization, normalizes each feature across the examples in a batch, so its result depends on the other examples and it behaves differently in training and in use. Transformers mostly use layer normalization, which avoids both effects.

Check yourself

  1. Work a similar problem. Normalize the token [0, 3], with γ = 1 and β = 0 (ignore ε).

    Answer

    Mean 1.5. Differences [−1.5, 1.5]. Variance (2.25 + 2.25) / 2 = 2.25, so the spread is 1.5. Output [−1, 1].

  2. What does layer normalization give for [5, 5, 5, 5], and why does ε matter here?

    Answer

    The mean is 5 and every difference is 0, so the variance is 0. Without ε the formula would divide 0 by 0. With ε it divides 0 by √0.00001 and gives [0, 0, 0, 0], and then the output is just β.

  3. Two tokens sit in the same batch. Does the normalized value of token 1 depend on token 2?

    Answer

    No. The mean and spread come from token 1's own numbers only. (Batch normalization, by contrast, would mix in other examples.)

  4. Training has set γ = [2, 2, 2, 2] and β = [1, 1, 1, 1]. What does the token [2, 4, 6, 8] become?

    Answer

    Normalize first: [−1.342, −0.447, 0.447, 1.342]. Then 2 × each + 1: [−1.683, 0.106, 1.894, 3.683].

Go further

Optional background, for readers who want more than this page.

Where this shows up

Units that build on this one