Layer Normalization

Layer normalization (LayerNorm, LN) standardizes the features of one example with that example’s own mean and standard deviation, and then applies a learned scale and shift to each feature. Ba, Kiros and Hinton introduced it in 2016. In a transformer it acts on each token separately: the $d$ features of a token are centred, divided by their standard deviation and rescaled. The statistics never involve other examples, so the layer computes the same function in training and inference, for a batch of one as for a batch of thousands. A standard transformer block for images, from the Vision Transformer to the denoiser of a Diffusion Transformer, contains two of these layers. Where they sit in the residual block matters as much as what they compute. The original “post-norm” placement normalizes after each residual sum. The “pre-norm” placement used by vision transformers normalizes each sub-layer’s input. The two placements give different gradients at initialization, and they differ in whether training needs a learning-rate warmup.

The comparison of layer normalization with batch, instance and group normalization, as different choices of the axes that share one mean and variance, belongs to normalization layers. This article covers the layer itself and how it is used. Ba, Kiros and Hinton’s paper is on arXiv, which this site’s research process does not use, so it was not read. The operation is described here from Xiong et al. (2020), from the PyTorch implementation and from the vision papers that use it.

Purpose

A deep network’s activations change scale as signals pass through many layers and as the weights change during training. A normalization layer resets that scale at fixed points of the network. Batch normalization does this with statistics pooled over the images in a mini-batch. That makes each image’s output depend on its batch, so inference needs separate running averages. Layer normalization takes the statistics from the example itself, so there is no batch dependence and no difference between modes. PyTorch’s documentation states that nn.LayerNorm uses statistics computed from the input data in both training and evaluation modes.

In transformers the layer has a second job, which Xiong et al. (2020) analysed. Each block adds its sub-layers’ outputs to a residual stream. In our reading of their analysis, the normalization layers are what keep the gradients that flow back through that stream at a controlled size, and their position decides how large those gradients are at initialization.

When to Use It

• Transformers on any modality. Every encoder and decoder block of the original Transformer, ViT, Swin and DiT uses LN or a variant of it. Xiong et al. describe LN as playing a key role in the Transformer’s success.
• Small or variable batches. The output for one example does not depend on any other example, so the layer is unaffected by the batch size, which is BN’s weakness in detection and video models that fit few images per device. Group normalization is the other common answer to that problem; normalization layers gives Wu and He’s comparison.
• Modern convolutional networks. ConvNeXt (Liu et al., 2022) uses one LN per block, applied over channels at each position. Directly substituting LN for BN in the original ResNet gives worse accuracy, as the ConvNeXt authors note when citing Wu and He.
• Conditioning a generator. Adaptive variants replace LN’s learned scale and shift with values computed from a class label, a diffusion timestep or another conditioning vector (Variants below).

LN is a weaker choice in classic convolutional classifiers trained with large batches. There, Wu and He measured LN behind BN and GN in ResNet-50 on ImageNet, as normalization layers reports.

How It Works

For one token $\mathbf{x} \in \mathbb{R}^d$, layer normalization computes

$$\mu = \frac{1}{d}\sum_{k=1}^{d} x_k, \qquad \sigma^2 = \frac{1}{d}\sum_{k=1}^{d} (x_k - \mu)^2, \qquad \operatorname{LN}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x} - \mu\mathbf{1}}{\sqrt{\sigma^2 + \epsilon}} + \boldsymbol{\beta},$$

with a learned scale $\boldsymbol{\gamma} \in \mathbb{R}^d$ initialized to ones and a learned shift $\boldsymbol{\beta} \in \mathbb{R}^d$ initialized to zeros. This is the form in Xiong et al. (§3.1) and in PyTorch’s nn.LayerNorm. PyTorch uses the biased variance and a default $\epsilon = 10^{-5}$. Each LN in a transformer of width $d$ therefore has $2d$ parameters: 1,536 for ViT-Base’s $d = 768$. With elementwise_affine=False the layer has no parameters, as in the adaptive variants below; with bias=False it keeps $\boldsymbol{\gamma}$ but drops $\boldsymbol{\beta}$.

The normalized_shape argument fixes which axes share the statistics: the trailing axes of the input. On a token sequence of shape $(N, T, d)$, nn.LayerNorm(d) computes one mean and variance per token. On a feature map of shape $(N, C, H, W)$, nn.LayerNorm([C, H, W]) computes one per image. These are two different operations that share a name; normalization layers counts the statistics of each on the same map.

Invariance and its gradient. The output does not change if all $d$ features of a token are multiplied by the same positive factor or shifted by the same constant (apart from the small effect of $\epsilon$). The same fact shapes the backward pass. Write $\mathbf{y} = \mathbf{x} - \mu\mathbf{1}$, and set $\boldsymbol{\gamma} = \mathbf{1}$, $\boldsymbol{\beta} = \mathbf{0}$ and $\epsilon = 0$. Then $\operatorname{LN}(\mathbf{x}) = \sqrt{d}\,\mathbf{y} / \lVert\mathbf{y}\rVert$, which is the form Xiong et al. use in the proof of their Lemma 3 (supplementary material, Eq. 20). Differentiating gives

$$\frac{\partial \operatorname{LN}(\mathbf{x})}{\partial \mathbf{x}} = \frac{\sqrt{d}}{\lVert \mathbf{y} \rVert}\left(I - \hat{\mathbf{y}}\hat{\mathbf{y}}^{\top}\right)\left(I - \tfrac{1}{d}\mathbf{1}\mathbf{1}^{\top}\right), \qquad \hat{\mathbf{y}} = \mathbf{y} / \lVert \mathbf{y} \rVert .$$

The two brackets project out the directions to which the output is invariant. The factor in front scales the gradient by $\sqrt{d}$ divided by the size of the input (our derivation). Xiong et al.’s Lemma 3 states the corresponding bound, $\lVert J_{\mathrm{LN}}(\mathbf{x}) \rVert_2 = O(\sqrt{d} / \lVert \mathbf{x} \rVert_2)$. A large input to an LN shrinks every gradient passing back through it. In Xiong et al.’s account, this normalization of the gradients is the main idea behind the difference between the two placements.

Post-norm and pre-norm. Xiong et al. (§3.1, Table 1) define the two placements for a block with self-attention and a feed-forward network (FFN). In the notation of the transformer article, with FFN (Xiong et al.’s term) for its MLP, where block $\ell$ maps $Z_{\ell-1}$ to $Z_\ell$:

$$\text{post-norm:}\quad Z' = \operatorname{LN}\big(Z_{\ell-1} + \operatorname{MSA}(Z_{\ell-1})\big), \qquad Z_\ell = \operatorname{LN}\big(Z' + \operatorname{FFN}(Z')\big);$$ $$\text{pre-norm:}\quad Z' = Z_{\ell-1} + \operatorname{MSA}\big(\operatorname{LN}(Z_{\ell-1})\big), \qquad Z_\ell = Z' + \operatorname{FFN}\big(\operatorname{LN}(Z')\big),$$

and pre-norm adds one final LN after the last block. Xiong et al. call the first the Post-LN Transformer: the original design of the Transformer and of BERT. They credit the Pre-LN Transformer to earlier systems (Baevski and Auli; Child et al.; Wang et al., among others). In post-norm, every block output is a normalized vector, so the residual stream has a fixed size. In pre-norm, the stream is never normalized until the end: each block adds its two sub-layer outputs to it.

Xiong et al. analysed both at initialization with mean-field theory, in a simplified setting: single-head attention with zero query and key matrices, so that attention is uniform; $d \times d$ weight matrices with Xavier initialization; and Gaussian inputs. In this setting:

• Lemma 2. In post-norm, the expected squared norm of the sum entering each block’s second LN is $\tfrac{3}{2}d$ at every layer. In pre-norm, the stream’s expected squared norm grows linearly with depth, between $(1 + \ell/2)d$ and $(1 + 3\ell/2)d$.
• Theorem 1. The gradient of the loss with respect to the last FFN matrix $W^{2,L}$ satisfies $\lVert \partial \tilde{\mathcal{L}} / \partial W^{2,L} \rVert_F \le O(d\sqrt{\ln d})$ for post-norm, independent of the depth $L$. For pre-norm, the bound is $O\big(d\sqrt{\ln d / L}\big)$. In the authors’ explanation, the input to pre-norm’s final LN grows linearly with $L$, so the LN Jacobian divides every gradient by about $\sqrt{L}$.
• Other layers. A “preliminary” analysis of the other layers, the authors’ word, found the post-norm gradient norm large near the output and likely to decay toward the input, and the pre-norm gradient norm likely to stay the same across layers.

They checked these empirically on a 6-layer-encoder, 6-layer-decoder translation Transformer (width 512, FFN width 1,024, 4 heads; supplementary §1.1). In their Figure 3, the gradient norm of $W^{2,L}$ stayed at about 1.6 for post-norm models from 6-6 to 14-14 layers and decreased with size for pre-norm. Across the layers of the 6-6 model, their text says that the scale of the expected FFN gradients grew with the layer index for post-norm and stayed almost the same for every layer in pre-norm. Their Figure 3(a) and (b) show the post-norm growth. For pre-norm, the $W^1$ bars in Figure 3(a) are nearly level, but the $W^2$ bars in Figure 3(b) fall from about 0.55 in layer 1 to about 0.3 in layer 6 (values read approximately from the figure), so for $W^2$ the figure does not match the text.

Why warmup. Xiong et al. hypothesize that the large gradients near the output are one reason the post-norm Transformer needs a learning-rate warmup. Large steps on those parameters at the start make training unstable. Their translation experiments (§3.2) support the need: on IWSLT14 German–English with Adam, a post-norm model trained without warmup reached a BLEU of only 8.45, against about 34 with a 4,000-step warmup. With a 500-step warmup it reached 31.16 BLEU at a maximum learning rate of $5 \times 10^{-4}$ and 2.77 at $10^{-3}$. Removing the warmup from pre-norm models did not hurt them (§4.2). On IWSLT14, pre-norm reached about 34 BLEU, comparable to post-norm with warmup, and its 9th checkpoint nearly matched post-norm’s 15th. For BERT pre-training they report a 40% speed-up for pre-norm. In their Figure 5(a), the pre-norm model’s validation loss reaches about 1.69 near 500k updates and the post-norm model’s near 700k (values read approximately from the figure). Their text gives the two models the other way round, post-norm at 500k and pre-norm at 700k, which does not fit its own conclusion. They also report that a larger learning rate, $3 \times 10^{-4}$, made the post-norm BERT diverge. Learning rate schedule covers warmup in general and Liu et al.’s competing explanation through Adam’s variance.

How to Apply It

In PyTorch, nn.TransformerEncoderLayer takes a norm_first flag. The default, False, gives post-norm; True gives pre-norm. An nn.TransformerEncoder built from pre-norm layers needs its own final nn.LayerNorm, which goes in its norm argument or after the stack. Vision transformers are pre-norm (the Vision Transformer article gives ViT’s equations), and Swin V1 inherited the same configuration (Liu et al., 2022, Swin V2 paper, §3.1).

Example: gradients at initialization, post-norm against pre-norm

The script below measures the quantity Xiong et al. analyse. It records the Frobenius norm of the gradient of the loss with respect to each block’s second FFN matrix $W_2$, for both placements, in untrained transformer stacks of 6, 12 and 24 layers.

• Tokens. The data are the first 256 of scikit-learn’s bundled $8 \times 8$ digit images. Pixels are standardized, and each image is cut into $2 \times 2$ patches, giving 16 tokens of 4 pixels.
• Embedding. Each token is a random linear patch embedding plus a random position embedding, scaled so that its entries have variance about 1 (the “unit” input scale).
• Blocks. The blocks are PyTorch’s nn.TransformerEncoderLayer with width $d = 64$, 4 heads, FFN width 256, ReLU and no dropout, followed by a linear classifier.
• Initialization. As in Xiong et al.’s §3.3: every weight matrix Xavier normal, $\mathcal{N}\big(0, 2/(n_\text{in} + n_\text{out})\big)$, biases zero, LN scales one. PyTorch stores the query, key and value projections as one $3d \times d$ matrix, and the code draws it as that one matrix.
• Loss. Cross-entropy of every token against its image’s digit class, averaged over all 4,096 tokens, as Xiong et al. average the loss over positions.

Nothing is trained. The code also prints:

• C1, a check of autograd against a central finite difference;
• C2, a check that RMSNorm equals LayerNorm when every token has zero mean (Variants);
• the squared size of the stream that each LN receives;
• the last-layer gradient against depth for a second input scale, $1/\sqrt{d}$, closer to Xiong et al.’s assumption that the inputs are drawn from the same Gaussian as the weights. The sources leave this scale open.

Weights are drawn for five seeds, 0 to 4, and the tables give the median with the minimum and maximum. The script runs in double precision on one thread with PyTorch 2.14.1, NumPy 2.5.4 and scikit-learn 1.9.1, and takes about 20 seconds on a laptop CPU.

import sys
import numpy as np
import torch, torch.nn as nn
from sklearn.datasets import load_digits

torch.set_num_threads(1)
torch.set_default_dtype(torch.float64)
SEEDS = [int(s) for s in sys.argv[1:]] or [0, 1, 2, 3, 4]
print("torch", torch.__version__, "| seeds", SEEDS)

# ---- data: 256 digit images, 2x2 patches -> 16 tokens of 4 pixels each ----
digits = load_digits()
X = torch.tensor(digits.images[:256])
y = torch.tensor(digits.target[:256]).repeat_interleave(16)        # every token carries its image's class
X = (X - X.mean()) / X.std()
patches = X.reshape(256, 4, 2, 4, 2).permute(0, 1, 3, 2, 4).reshape(256, 16, 4)
print("patches", tuple(patches.shape), "pixel mean %.1e std %.3f" % (X.mean(), X.std()),
      "| class counts", np.bincount(digits.target[:256]).tolist())

d, HEADS, FFN = 64, 4, 256

def build(L, pre, seed, scale=1.0):
    """L encoder layers, post-norm (pre=False) or pre-norm (pre=True), Xavier-normal weights, zero biases."""
    g = torch.Generator().manual_seed(seed)
    E = torch.randn(4, d, generator=g) / 2.0                         # patch embedding: entries ~ variance 1
    P = torch.randn(16, d, generator=g)                              # position embedding
    layers = nn.ModuleList([nn.TransformerEncoderLayer(d, HEADS, FFN, dropout=0.0, activation="relu",
                                                       batch_first=True, norm_first=pre) for _ in range(L)])
    final = nn.LayerNorm(d) if pre else nn.Identity()                # pre-norm adds a final LN (Xiong, Table 1)
    head = nn.Linear(d, 10)
    for name, p in list(layers.named_parameters()) + list(head.named_parameters()):
        with torch.no_grad():
            if p.dim() == 2:                                         # Xavier normal: N(0, 2 / (n_in + n_out))
                p.copy_(torch.randn(p.shape, generator=g) * (2.0 / sum(p.shape)) ** 0.5)
            elif "norm" in name and "weight" in name:
                p.fill_(1.0)
            else:
                p.zero_()
    x0 = scale * (patches @ E + P) / 2 ** 0.5
    return x0, layers, final, head

def loss_of(x0, layers, final, head):
    h = x0
    for layer in layers:
        h = layer(h)
    return nn.functional.cross_entropy(head(final(h)).reshape(-1, 10), y)

def grad_norms(L, pre, seed, scale=1.0):
    """Frobenius norm of dLoss/dW2 for every layer, and the squared norm / d of the stream each LN sees."""
    x0, layers, final, head = build(L, pre, seed, scale)
    seen = []
    hook_on = [lay.norm2 if not pre else lay for lay in layers] + ([final] if pre else [])
    hooks = [m.register_forward_pre_hook(lambda m, a: seen.append((a[0] ** 2).sum(-1).mean().item() / d))
             for m in hook_on]
    loss_of(x0, layers, final, head).backward()
    for hk in hooks:
        hk.remove()
    return np.array([lay.linear2.weight.grad.norm().item() for lay in layers]), np.array(seen)

# ---- C1: autograd against a central finite difference (L = 6, seed 0, largest entry of W2) ----
for pre in (False, True):
    x0, layers, final, head = build(6, pre, 0)
    loss_of(x0, layers, final, head).backward()
    for li in (0, 5):
        W = layers[li].linear2.weight
        i, j = divmod(W.grad.abs().argmax().item(), W.shape[1])
        h = 1e-6
        with torch.no_grad():
            W[i, j] += h; lp = loss_of(x0, layers, final, head).item()
            W[i, j] -= 2 * h; lm = loss_of(x0, layers, final, head).item()
            W[i, j] += h
        fd = (lp - lm) / (2 * h)
        print(f"C1 {'pre ' if pre else 'post'}-norm layer {li + 1}: autograd vs finite difference, "
              f"rel err {abs(W.grad[i, j].item() - fd) / abs(fd):.1e}")

# ---- C2: RMSNorm equals LayerNorm when every token has zero mean ----
x0 = build(1, True, 0)[0].detach()
ln, rms = nn.LayerNorm(d, eps=1e-5), nn.RMSNorm(d, eps=1e-5)
with torch.no_grad():
    cen = x0 - x0.mean(-1, keepdim=True)
    print("C2 max |RMSNorm - LayerNorm|: centred tokens %.1e, uncentred tokens %.1e"
          % ((rms(cen) - ln(cen)).abs().max(), (rms(x0) - ln(x0)).abs().max()))

# ---- M1 and M3: per-layer gradient norms and stream norms, L = 12 ----
L = 12
res = {pre: [grad_norms(L, pre, s) for s in SEEDS] for pre in (False, True)}
G = {pre: np.array([r[0] for r in res[pre]]) for pre in res}
N = {pre: np.array([r[1] for r in res[pre]]) for pre in res}
print(f"\nL = {L}: ||dLoss/dW2|| per layer (median over seeds [min, max]); |x|^2/d of the stream each LN sees")
print("layer   post-norm grad          pre-norm grad          post |x5|^2/d  pre |x_l|^2/d")
for l in range(L):
    print(f"{l + 1:5d}  " + "  ".join(f"{np.median(G[p][:, l]):.4f} [{G[p][:, l].min():.4f}, {G[p][:, l].max():.4f}]"
                                      for p in (False, True))
          + f"  {np.median(N[False][:, l]):12.2f}  {np.median(N[True][:, l]):12.2f}")
print(f"final LN input (pre-norm) |x|^2/d: {np.median(N[True][:, L]):.2f}")
for p in (False, True):
    ratio = G[p][:, -1] / G[p][:, 0]
    print(f"{'pre ' if p else 'post'}-norm top/bottom gradient ratio: median {np.median(ratio):.2f} "
          f"[{ratio.min():.2f}, {ratio.max():.2f}]")

# ---- M2: last-layer gradient against depth, two input scales ----
print("\nlast layer ||dLoss/dW2_L|| (median over seeds) against depth L")
print("input scale    L   post-norm   pre-norm   pre-norm * sqrt(L)")
for name, scale in (("unit", 1.0), ("1/sqrt(d)", d ** -0.5)):
    for L in (6, 12, 24):
        post = np.median([grad_norms(L, False, s, scale)[0][-1] for s in SEEDS])
        pre = np.median([grad_norms(L, True, s, scale)[0][-1] for s in SEEDS])
        print(f"{name:12s} {L:3d}   {post:9.4f}  {pre:9.4f}   {pre * L ** 0.5:9.4f}")

Output:

torch 2.14.1 | seeds [0, 1, 2, 3, 4]
patches (256, 16, 4) pixel mean 1.1e-16 std 1.000 | class counts [26, 26, 26, 26, 25, 26, 25, 25, 26, 25]
C1 post-norm layer 1: autograd vs finite difference, rel err 1.3e-09
C1 post-norm layer 6: autograd vs finite difference, rel err 1.7e-09
C1 pre -norm layer 1: autograd vs finite difference, rel err 6.2e-10
C1 pre -norm layer 6: autograd vs finite difference, rel err 7.1e-09
C2 max |RMSNorm - LayerNorm|: centred tokens 8.9e-16, uncentred tokens 5.0e-01

L = 12: ||dLoss/dW2|| per layer (median over seeds [min, max]); |x|^2/d of the stream each LN sees
layer   post-norm grad          pre-norm grad          post |x5|^2/d  pre |x_l|^2/d
    1  2.5547 [2.2356, 4.7709]  2.0740 [1.6157, 2.4987]          1.32          0.98
    2  2.6244 [2.1467, 5.6966]  1.9730 [1.2945, 2.2989]          1.34          1.39
    3  2.3655 [1.9953, 5.3012]  1.5216 [1.2256, 1.6837]          1.33          1.92
    4  2.5199 [2.2892, 6.0051]  1.4641 [1.0564, 1.8008]          1.34          2.33
    5  2.7639 [2.1897, 5.9066]  1.3770 [1.0096, 1.7774]          1.22          2.81
    6  2.2463 [2.0483, 4.5133]  1.0946 [0.8300, 1.5170]          1.24          3.35
    7  2.2406 [1.9323, 4.7517]  1.0727 [0.7997, 1.1480]          1.20          3.83
    8  2.3955 [2.1652, 4.1749]  0.9505 [0.8421, 1.2995]          1.28          4.13
    9  2.4251 [2.1639, 4.1113]  0.9648 [0.8275, 1.2966]          1.40          5.03
   10  2.3281 [1.5022, 4.4927]  0.7847 [0.7288, 0.9872]          1.25          5.57
   11  2.3088 [1.4843, 4.6279]  0.8610 [0.7341, 1.0374]          1.30          6.15
   12  2.8242 [1.7681, 4.3794]  0.8588 [0.6586, 0.9463]          1.33          6.58
final LN input (pre-norm) |x|^2/d: 7.03
post-norm top/bottom gradient ratio: median 0.92 [0.63, 1.71]
pre -norm top/bottom gradient ratio: median 0.40 [0.38, 0.46]

last layer ||dLoss/dW2_L|| (median over seeds) against depth L
input scale    L   post-norm   pre-norm   pre-norm * sqrt(L)
unit           6      2.0313     0.9351      2.2905
unit          12      2.8242     0.8588      2.9751
unit          24      3.8435     0.5253      2.5734
1/sqrt(d)      6      2.0551     1.1273      2.7614
1/sqrt(d)     12      2.8455     1.1349      3.9313
1/sqrt(d)     24      3.7306     0.5563      2.7251

The checks. C1 shows autograd agreeing with finite differences to about $10^{-9}$, so the gradient norms are the true derivatives of the loss. C2 shows that on tokens with zero mean, RMSNorm and LayerNorm agree to rounding error, while on the raw tokens they differ by 0.5.

The stream. The size of what each LN receives follows Xiong et al.’s Lemma 2 qualitatively. In pre-norm, the squared norm of the stream entering block $\ell$ grows steadily, from about $d$ at the input to $6.6d$ before the 12th block and $7.0d$ at the final LN, because every block adds its two branch outputs and nothing normalizes the sum. In post-norm, each block starts from a normalized vector, and the squared norm of the sum entering its second LN stays between $1.2d$ and $1.4d$. The lemma’s constants are for its simplified setting and are not expected here.

The last layer against depth. For pre-norm, the gradient of $W_2$ in the last layer falls with depth at the unit input scale. Multiplied by $\sqrt{L}$ it varies less than it does itself, as Theorem 1’s $1/\sqrt{L}$ factor predicts. The mechanism is the one above: the final LN receives a stream that grows with $L$, and its Jacobian divides by that stream’s size. At the $1/\sqrt{d}$ input scale, the median does not fall from 6 to 12 layers and then halves at 24. Post-norm does not show the independence of depth that Xiong et al. observed (about 1.6 from 6-6 to 14-14 layers). Its last-layer gradient grows from about 2.0 to 3.7–3.8 between 6 and 24 layers, at both input scales. Theorem 1 is an upper bound that is independent of $L$, which growth at this size does not contradict. However, the paper’s empirical statement that the post-norm gradient stays the same is not reproduced here.

Across layers. The per-layer profiles go against the authors’ stated preliminary analysis and what their text reports for the Figure 3 setting:

• Post-norm. The gradients do not grow with the layer index. The median ratio of the top to the bottom layer is 0.92, and across seeds the ratio ranges from 0.63 to 1.71, so even its direction depends on the seed.
• Pre-norm. The gradients are not flat. They fall from the first layer to the last, to 0.40 of the first layer’s value, consistently across seeds. The $W^2$ bars of Xiong et al.’s own Figure 3(b) also fall with the layer index, unlike their text’s description.

In our reading, the pre-norm decline comes from the construction: the gradient reaching an early block’s $W_2$ passes back both along the identity path and through the later blocks’ branches. Each of those branches starts with an LN whose Jacobian is larger where the stream is small, and the stream is small in early layers (last column). The set-up differs from Xiong et al.’s in several ways that may matter, and the example cannot separate them:

• it is an encoder only, $d = 64$ instead of 512;
• its tokens are embedded 4-pixel patches, not word embeddings;
• it uses one batch of 256 images instead of an average over mini-batches;
• it has no dropout and no label smoothing.

What the example does reproduce is the part of the theory that rests on LN’s Jacobian: pre-norm’s stream grows with depth, and pre-norm’s last-layer gradient shrinks with it at the unit input scale (at the $1/\sqrt{d}$ scale, only from 12 to 24 layers).

Variants

RMSNorm. Root mean square normalization (Zhang and Sennrich, 2019) drops the mean subtraction and the shift. PyTorch’s nn.RMSNorm computes

$$\operatorname{RMSNorm}(\mathbf{x}) = \mathbf{g} \odot \frac{\mathbf{x}}{\sqrt{\epsilon + \frac{1}{d}\sum_{k} x_k^2}},$$

with a default $\epsilon$ equal to the machine epsilon of the computation type. When a token has zero mean, $\frac{1}{d}\sum_k x_k^2$ is its variance, and RMSNorm with the same $\epsilon$ is LayerNorm without the shift, as check C2 confirms. Otherwise the two differ, and RMSNorm is not invariant to a constant shift of the features. Zhang and Sennrich’s paper was not read for this article. Normalization layers summarizes its hypothesis and results. ViT-22B (Dehghani et al., 2023) applied all its LayerNorms without bias and centering, citing Zhang and Sennrich. This improved accelerator utilization by 3% without quality degradation, the authors report.

LayerNorm in convolutional networks. ConvNeXt applies LN to the channel vector at each pixel, the per-token form of transformers applied to a feature map. The ConvNeXt modernization of ResNet-50 first cut each block to a single BN layer, which reached 81.4% ImageNet-1K top-1 accuracy. Replacing that BN by LN gave 81.5%, and the authors found no difficulty training with LN (Liu et al., 2022, §2.6). The final design also adds LNs wherever the resolution changes: after the stem, before each downsampling layer and after the final global average pooling. Separate downsampling layers without these LNs made training diverge, and with them accuracy rose to 82.0%. The official code has a LayerNorm class for both memory layouts. It normalizes the last axis with F.layer_norm for channels-last tensors $(N, H, W, C)$, and computes the same per-pixel channel statistics by hand for channels-first tensors $(N, C, H, W)$, with $\epsilon = 10^{-6}$. Its block docstring describes the two layouts as equivalent and uses the channels-last one, permuting the tensor around the LN and the two linear layers because the authors found it slightly faster in PyTorch.

Residual post-norm. Swin Transformer V2 (Liu et al., 2022) moved each LN from the beginning of the residual branch to its end: the branch output is normalized before it is added back, so that $Z_\ell = Z_{\ell-1} + \operatorname{LN}(F(Z_{\ell-1}))$. The authors observed that in pre-norm models the activation amplitude of the main branch grows with depth. In their large models, the discrepancy between the layers with the highest and lowest amplitudes reached $10^4$. A huge-size Swin V1 model failed to complete self-supervised pre-training, and they give the large amplitude discrepancy across layers as the cause of the instability. With “res-post-norm” the amplitude does not accumulate. Their largest models also add an LN on the main branch every 6 blocks.

In their ablation (Table 6, ImageNet-1K top-1), res-post-norm alone raised Swin-B from 83.6% to 83.8%, and 84.1% with their scaled cosine attention. Footnote 1 of their paper says that the original post-normalization harms training stability, citing Xiong et al. Table 7 notes that post-norm diverged at the default learning rate and was trained at a quarter of it.

Query–key normalization. ViT-22B (Dehghani et al., 2023) applies LN to the queries and keys before the dot product, computing the attention weights as $\operatorname{softmax}\big(\operatorname{LN}(XW_Q)\operatorname{LN}(XW_K)^{\top}/\sqrt{d}\big)$, following Gilmer et al. Without it, models of about 8 billion parameters diverged after a few thousand steps: attention logits grew past 50,000 in magnitude and the attention weights became almost one-hot (their §2 and Appendix B). ViT-22B also runs attention and the MLP in parallel from a single LN, $\mathbf{y} = \mathbf{x} + \operatorname{MLP}(\operatorname{LN}(\mathbf{x})) + \operatorname{Attention}(\operatorname{LN}(\mathbf{x}))$. Stable Diffusion 3 (Esser et al., 2024, Appendix C.2) uses RMSNorm with a learnable scale on queries and keys in both streams of its MM-DiT. This kept the attention logits from growing when fine-tuning at high resolution in mixed precision. The authors add that the method is not a universal recipe.

Adaptive layer norm. In a conditional generator, the scale and shift can be computed from the condition instead of learned as constants. DiT (Peebles and Xie, 2023) replaces each LN’s $\boldsymbol{\gamma}$ and $\boldsymbol{\beta}$ with values regressed from the sum of the timestep and class embeddings (“adaLN”), following adaptive normalization in GANs and diffusion U-Nets. adaLN-Zero also regresses a per-dimension gate applied to each residual branch and initializes it to zero, so that every block starts as the identity. In the authors’ comparison of DiT-XL/2 blocks, adaLN-Zero had lower FID than cross-attention and in-context conditioning and clearly outperformed plain adaLN. Unlike cross-attention, adaLN applies the same function to every token. The Diffusion Transformer article gives the block’s equations and the official code’s details, including that the LNs there have no affine parameters of their own.

Best Practices

• Use pre-norm for new vision transformers trained from scratch, with a final LN after the last block. This is the configuration of ViT and Swin V1. Xiong et al. trained pre-norm models without warmup; a short warmup remains common in vision recipes (learning rate schedule).
• If a model must be post-norm, for instance to reproduce an original Transformer or BERT, keep the warmup and the learning rates the original used. Xiong et al. found post-norm sensitive to both, and found that a larger learning rate made post-norm BERT diverge.
• Normalize over the right axes. For tokens, normalize the last axis of size $d$. For a channels-first feature map, either permute to channels-last as ConvNeXt does, or write the per-pixel reduction over axis 1. nn.LayerNorm([C, H, W]) is a different, per-image operation.
• Watch depth in pre-norm models. The stream grows with depth, as the example shows. Swin V2’s res-post-norm, ViT-22B’s normalizations and DiT’s zero-initialized gates are each a response to instability in deep or large models.

Pitfalls

• Expecting pre-norm and post-norm checkpoints to be interchangeable. The two placements compute different functions from the same weights. A norm_first flag that differs between training and loading code silently changes the model.
• Forgetting the final LN in pre-norm. Without it, the head receives the unnormalized stream, whose size grows with depth (a squared norm of $7.0d$ after 12 blocks in the example).
• Assuming LN’s output is independent of scale everywhere. It is invariant to a common scale and shift of one token’s features, but $\epsilon$ breaks the invariance for very small inputs. RMSNorm is not shift-invariant at all.
• Reading theory as measurement. Xiong et al.’s bounds are upper bounds under simplifying assumptions, and their per-layer analysis is preliminary. The example above reproduces the pre-norm depth scaling at the unit input scale, but not the post-norm per-layer growth.

Evidence and Impact

The strongest evidence on placement is Xiong et al.’s. In theory, the post-norm gradient near the output does not depend on depth, and the pre-norm gradient shrinks as $1/\sqrt{L}$. In their translation and BERT experiments, pre-norm models trained without warmup matched the post-norm baselines in less training time. They also found that the position of LN “dominates” the change of optimizer, in their word: RAdam trained post-norm well without warmup but made little difference to pre-norm.

In vision, ViT used pre-norm, and Swin V1 inherited it from language transformers and ViT without, in the Swin V2 authors’ words, extensive study. ConvNeXt carried LN into a convolutional network at no accuracy cost, 81.5% against 81.4% with BN in its modernized ResNet-50. At larger scales, the papers reviewed here report that the normalization configuration had to change to keep training stable: Swin V2 moved LN to the end of each branch, ViT-22B normalized queries and keys, and Stable Diffusion 3 did the same with RMSNorm. In generation, DiT’s ablation found adaptive LN with zero-initialized gates the best of its conditioning mechanisms.