Normalization layers and residual connections

Before this lesson: lesson 03, backpropagation (adjoints, Jacobians) and lesson 04, activations and initialization (variance propagation). You also need the mean and variance of a list of numbers, and the projection onto a vector (linear algebra lesson 03). Time: about 75 minutes, plus the notebook.

Why this matters #

• ResNets (He et al., 2016) stack conv → BatchNorm → ReLU inside residual blocks and train with more than 100 layers. Before them, deeper plain networks had a higher training error.
• Transformers put a LayerNorm (or RMSNorm in LLaMA) in every block, around a residual stream. The position of that norm (pre-LN or post-LN) decides whether training needs a learning-rate warmup.
• Diffusion models and science models (weather, climate, physics) use GroupNorm or LayerNorm, because their batches per GPU are tiny and BatchNorm statistics would be noise.
• Interviewers ask you to derive the BatchNorm backward pass and to explain why transformers do not use BatchNorm. This lesson does both.

1. Predict first #

2. The picture #

Normalize, then rescale #

Take one feature and a batch of $N$ numbers. BatchNorm does three things:

1. subtract the batch mean $\mu$;
2. divide by the batch standard deviation $\sigma$; now the numbers have mean 0 and variance 1;
3. multiply by a learned scale $\gamma$ and add a learned shift $\beta$.

Add a constant to every input, or multiply every input by a positive constant: steps 1 and 2 undo it, so the output does not change. The layer forgets the location and the scale of the batch, and $\gamma, \beta$ put back a location and a scale that the network learns.

Which numbers share one mean and variance #

The normalization layers differ only in the group of numbers that share one $(\mu, \sigma)$.

The key difference: BatchNorm mixes the examples of a batch; LayerNorm, RMSNorm and GroupNorm never do. The output for one example depends on the other examples only with BatchNorm.

Residual connections, simplest case #

Take 50 scalar layers. In a plain stack $x_{l+1} = f_l(x_l)$, the chain rule multiplies 50 derivatives $f_l'$. If each has size $0.1$, the product has size $10^{-50}$: the first layer receives no signal.

In a residual stack $x_{l+1} = x_l + F_l(x_l)$, the derivative of one block is $1 + F_l'$. When $F_l'$ is small, this factor is close to 1, so the product of 50 such factors does not collapse to 0. Section 4 computes $0.78$ instead of $10^{-50}$ for the same layer derivatives.

3. Definitions and main results #

BatchNorm, LayerNorm and RMSNorm (forward) #

Take one group of $n$ numbers $x_1, \dots, x_n$: one channel over the batch for BatchNorm, the features of one token for LayerNorm. Then

• $\hat x$ is the normalized input; $\gamma$ (scale, or gain) and $\beta$ (shift, or bias) are learned.
• $\varepsilon$ (PyTorch default $10^{-5}$) keeps the division safe when $\sigma \approx 0$.
• The variance used to normalize is the biased one: divide by $n$, not by $n - 1$.
• In LayerNorm, $\gamma$ and $\beta$ are vectors over the features and act entry by entry.

RMSNorm (Zhang and Sennrich, 2019) drops the mean and $\beta$:

It needs one reduction (a sum over the features) instead of two, and gives similar quality. LLaMA uses it; T5 uses the same layer under the name "layer norm without bias and mean".

Invariances. Let $\varepsilon = 0$, $c > 0$, $d \in \mathbb{R}$, and $\mathbf 1$ the all-ones vector. BatchNorm and LayerNorm give the same output for $x$ and for $cx + d\mathbf 1$. RMSNorm gives the same output for $x$ and $cx$, but not for a shift.

BatchNorm in training and in eval mode #

In training, BatchNorm normalizes with the batch statistics. It also updates running estimates $\hat\mu$, $\hat\sigma^2$ of the mean and variance with momentum $m$ (PyTorch default 0.1, the weight of the new value):

The running variance uses the unbiased estimate $\tfrac{n}{n-1}\sigma^2$; the normalization itself uses the biased one. In eval mode BatchNorm uses $\hat\mu$ and $\hat\sigma^2$. Then it is a fixed affine map per channel,

which can be folded into the linear or convolution layer before it. LayerNorm and RMSNorm do not use the batch, so they compute the same function in training and in eval mode.

The backward pass removes two components #

Theorem (backward pass of BatchNorm and LayerNorm). Take one group in training mode. Let $\bar y_i = \partial\mathcal{L}/\partial y_i$ be the upstream gradient and $g_i = \gamma\,\bar y_i$. Then

Here $\operatorname{mean}(v) = \frac1n\sum_i v_i$ and $\odot$ is the entry-wise product.

Why this is true, in one picture. Moving $x$ along $\mathbf 1$ (a shift) or along $\hat x$ (a scale change) does not change the output. So the gradient with respect to $x$ cannot have a component in these two directions. The formula takes $g$ and removes exactly these two components.

Proof.

1. From $y_i = \gamma \hat x_i + \beta$ we get $\bar\beta$ and $\bar\gamma$ directly, and $\partial\mathcal{L}/\partial \hat x_i = \gamma\,\bar y_i = g_i$.
2. The mean: $\partial \mu/\partial x_j = 1/n$.
3. The variance: $\partial \sigma^2/\partial x_j = \frac2n \sum_i (x_i - \mu)(\delta_{ij} - \frac1n) = \frac2n (x_j - \mu)$, because $\sum_i (x_i - \mu) = 0$. So $\partial\tilde\sigma/\partial x_j = (x_j - \mu)/(n\tilde\sigma)$.
4. The quotient rule for $\hat x_i = (x_i - \mu)/\tilde\sigma$:

5. The chain rule: $\bar x_j = \sum_i g_i\, \partial\hat x_i/\partial x_j = \frac{1}{\tilde\sigma}\big(g_j - \operatorname{mean}(g) - \hat x_j \operatorname{mean}(g\odot\hat x)\big)$. $\blacksquare$

Geometry. Write $\bar x = \frac1{\tilde\sigma} P g$ with $P = I - \frac1n \mathbf 1\mathbf 1^\top - \frac1n \hat x \hat x^\top$. We have $\mathbf 1 \perp \hat x$ (because $\sum_i \hat x_i = 0$), $\|\mathbf 1\|^2 = n$, and $\|\hat x\|^2 = n\sigma^2/(\sigma^2 + \varepsilon)$, which equals $n$ when $\varepsilon = 0$. In that case $P$ is the orthogonal projector onto the complement of $\operatorname{span}\{\mathbf 1, \hat x\}$ (projections, linear algebra lesson 03). Hence

• $\sum_j \bar x_j = 0$ exactly (for every $\varepsilon$), and
• $\sum_j \bar x_j \hat x_j = 0$ (exactly for $\varepsilon = 0$).

LayerNorm has the same formula over the feature axis of each token, with $g = \gamma \odot \bar y$; $\bar\gamma$ and $\bar\beta$ are then summed over all tokens.

RMSNorm. The same steps without the mean give $\bar x = \frac1r\big(g - \hat x\,\operatorname{mean}(g \odot \hat x)\big)$ with $g = \gamma \odot \bar y$. Only one component is removed, the one along $\hat x$, which matches the single invariance $x \mapsto cx$.

Why transformers use LayerNorm or RMSNorm #

BatchNorm makes the output of one example depend on the other examples of the batch. In transformers this causes concrete problems:

• Small or variable batches. Large models fit only a few sequences on one device, so the batch statistics are noisy. Synchronizing BatchNorm across devices costs communication in every layer.
• Train/eval mismatch. Training uses batch statistics; inference uses running ones. When a language model generates text, it processes one new token of one sequence at a time: there is no batch.
• Padding and variable lengths. Pad tokens enter the statistics unless they are masked.
• Causal leakage. Statistics over batch and time include future tokens, which a causal language model must not see.

LayerNorm normalizes each token over its own features. It does not depend on other examples or positions, it computes the same function in training and inference, and it works one token at a time with a KV cache (lesson 11).

The identity path of residual networks #

A residual block computes $x_{l+1} = x_l + F_l(x_l)$. $F_l$ is the branch (for example conv–BN–ReLU–conv, or attention, or an MLP); the term $x_l$ is the skip connection. The sequence $x_0, x_1, \dots$ is often called the residual stream.

Theorem (the identity path). Let $x_{l+1} = x_l + F_l(x_l)$ and let $J_k = \partial F_k/\partial x_k$ be the Jacobian of branch $k$. Then for $l < L$:

Proof.

1. Differentiating the update gives $\partial x_{k+1}/\partial x_k = I + J_k$.
2. The chain rule multiplies these factors for $k = l, \dots, L-1$, later layers on the left.
3. Expanding the product gives one term for every subset of layers; the empty subset gives $I$. $\blacksquare$

Why it helps. For the gradient (a row vector), $\partial\mathcal{L}/\partial x_l = \partial\mathcal{L}/\partial x_L\,(I + \cdots)$. The term $\partial\mathcal{L}/\partial x_L \cdot I$ reaches every layer unchanged, however small the $J_k$ are: an addition node copies its upstream gradient to both inputs (lesson 03). A plain stack gives $\partial x_L/\partial x_l = J_{L-1}\cdots J_l$, which shrinks or grows exponentially unless every $J_k$ is close to an isometry. He et al. (2016) observed this as the degradation problem: deeper plain networks had a higher training error than shallower ones. A residual network with small branches starts close to the identity map, which is easy to optimize.

The expansion has $2^{L-l}$ terms, one per path through the network. Veit et al. (2016) found that most of the gradient flows through the short paths, so a ResNet behaves a bit like an ensemble of shallower networks.

The forward pass can still explode. At initialization the branch output is uncorrelated with $x_l$ (the weights have mean zero), so

A He-initialized branch keeps the variance of its input, so the variance doubles in every block. Standard fixes:

• a normalization inside the branch (pre-LN, below);
• zero initialization of the last layer, or of the last $\gamma$, in each branch, so that each block starts as the identity;
• scaling the branch by about $1/\sqrt{L}$. GPT-2 scales the initial weights of the layers that write into the residual stream by $1/\sqrt{N}$, where $N$ is the number of residual layers.

Post-LN and pre-LN transformer blocks #

The original Transformer and BERT use post-LN. GPT-2, LLaMA (with RMSNorm) and most modern models use pre-LN.

• Identity path. In pre-LN, $x_L = x_l + \sum_{k = l}^{L-1} F_k(\operatorname{LN}(x_k))$, so $\partial x_L/\partial x_l = I + \cdots$ with no LayerNorm on the direct path. In post-LN, the gradient passes through $L - l$ LayerNorm Jacobians. Pre-activation ResNets (He et al., 2016, "Identity Mappings in Deep Residual Networks") made the same change for CNNs: BN and ReLU move into the branch, and the skip path is a pure identity.
• Gradient size at initialization (Xiong et al., 2020, Theorem 1). For the weight $W^{2,L}$ of the last FFN layer and width $d$, the gradient norm is bounded by $O(d\sqrt{\ln d})$ for post-LN, independent of the depth $L$, and by $O(d\sqrt{\ln d / L})$ for pre-LN. The reason: the Jacobian of a LayerNorm has norm $O(\sqrt d/\|x\|)$, and in pre-LN the squared norm of the residual stream grows linearly with depth, so the final LayerNorm divides the gradient by about $\sqrt L$.
• Warmup. Large gradients near the output make a large initial learning rate unstable in post-LN. A learning-rate warmup (start with a tiny learning rate and increase it over the first steps) keeps the first updates small. Xiong et al. trained pre-LN transformers without warmup to similar results in less time.
• Final norm. In pre-LN the residual stream itself is never normalized, so the stack needs a final LayerNorm before the output head.
• Growth of the residual stream. In pre-LN each branch adds an $O(1)$ amount to a stream whose squared norm grows about linearly in $l$. So block $l$ changes the stream by a relative amount of about $1/\sqrt l$: late blocks change it less. Post-LN, when it trains, can reach slightly better quality; DeepNorm (Wang et al., 2022) up-weights the skip path in post-LN and trains transformers with 1000 layers.

Answers to section 1. (1) Zero. BatchNorm gives the same output for $x$ and $x + c\mathbf 1$, so $\mathcal{L}(x + c\mathbf 1) = \mathcal{L}(x)$ for every $c$. Differentiating in $c$ gives $\sum_i \partial\mathcal{L}/\partial x_i = 0$. (2) The variance doubles in every block: $2^{24} \approx 1.7 \cdot 10^7$. Residual connections protect the identity path of the gradient, not the forward scale; normalization or branch scaling is still needed.

4. Worked example #

(a) BatchNorm forward. One channel, training mode, batch $x = (0, 2, 4, 6)$, $\gamma = 1$, $\beta = 0$, $\varepsilon = 0$, momentum 0.1, running mean 0 and running variance 1 before the step.

• Mean: $\mu = (0 + 2 + 4 + 6)/4 = 3$. Deviations: $(-3, -1, 1, 3)$.
• Biased variance: $\sigma^2 = (9 + 1 + 1 + 9)/4 = 5$, so $\sigma = \sqrt5 \approx 2.2361$.
• Output: $y = \hat x = (-3, -1, 1, 3)/\sqrt5 = (-1.3416, -0.4472, 0.4472, 1.3416)$.
• Running mean: $0.9 \cdot 0 + 0.1 \cdot 3 = 0.3$.
• Running variance: the unbiased estimate is $\tfrac43 \cdot 5 = 20/3 \approx 6.667$, so $0.9 \cdot 1 + 0.1 \cdot 6.667 = 1.5667$.
• Eval mode now maps $x = 3$ to $(3 - 0.3)/\sqrt{1.5667} \approx 2.157$. Training mode mapped it to 0. After one step the running statistics are still far from the data.

(b) BatchNorm backward. Take $\mathcal{L} = y_1$, so $\bar y = g = (1, 0, 0, 0)$ (with $\gamma = 1$).

• $\bar\beta = 1$ and $\bar\gamma = \hat x_1 = -1.3416$.
• $\operatorname{mean}(g) = 0.25$ and $\operatorname{mean}(g \odot \hat x) = -1.3416/4 = -0.3354$.
• The third term is $-\hat x_j \cdot (-0.3354) = 0.3354\,(x_j - \mu)/\sqrt5 = 0.15\,(x_j - \mu)$.
• So the bracket is $g_j - 0.25 + 0.15(x_j - \mu)$:
  • $j = 1$: $1 - 0.25 + 0.15 \cdot (-3) = 0.30$;
  • $j = 2$: $0 - 0.25 + 0.15 \cdot (-1) = -0.40$;
  • $j = 3$: $0 - 0.25 + 0.15 \cdot 1 = -0.10$;
  • $j = 4$: $0 - 0.25 + 0.15 \cdot 3 = 0.20$.
• $\bar x = (0.30, -0.40, -0.10, 0.20)/\sqrt5 = (0.1342, -0.1789, -0.0447, 0.0894)$.
• Checks: the entries sum to 0, and $\bar x \cdot \hat x \propto 0.30 \cdot (-3) - 0.40 \cdot (-1) - 0.10 \cdot 1 + 0.20 \cdot 3 = -0.9 + 0.4 - 0.1 + 0.6 = 0$.
• Meaning: raising $x_2$ lowers $y_1$. It raises the mean and shrinks the spread, and both push $y_1 = (x_1 - \mu)/\sigma$ further below zero.

(c) The same vector as one token. LayerNorm over its $D = 4$ features gives the same $\hat x$. RMSNorm: $r = \sqrt{(0 + 4 + 16 + 36)/4} = \sqrt{14} \approx 3.742$, so $y = (0, 2, 4, 6)/3.742 = (0, 0.5345, 1.0690, 1.6036)$. It is not centered.

(d) Residual vs plain, 50 scalar layers. Let the layer derivatives ($f_l'$ for the plain stack, $F_l'$ for the residual branches) be $+0.1$ for 25 layers and $-0.1$ for the other 25.

• Plain: $\big|\prod_l f_l'\big| = 0.1^{50} = 10^{-50}$.
• Residual: $\prod_l (1 + F_l') = 1.1^{25} \cdot 0.9^{25} = (1.1 \cdot 0.9)^{25} = 0.99^{25} \approx 0.778$.

5. In code #

PyTorch and autograd check the worked example. PyTorch refuses $\varepsilon = 0$ in training, so we use $\varepsilon = 10^{-12}$.

import torch
import torch.nn.functional as F

x = torch.tensor([[0.], [2.], [4.], [6.]], dtype=torch.float64, requires_grad=True)
bn = torch.nn.BatchNorm1d(1, eps=1e-12, momentum=0.1, dtype=torch.float64)
y = bn(x)                                    # training mode: batch statistics
print(y.detach().ravel().numpy().round(4))   # [-1.3416 -0.4472  0.4472  1.3416]
print(round(bn.running_mean.item(), 4), round(bn.running_var.item(), 4))   # 0.3 1.5667

y[0, 0].backward()                           # upstream gradient (1, 0, 0, 0)
dx = x.grad.ravel()
print((dx * 5 ** 0.5).numpy().round(4))      # [ 0.3 -0.4 -0.1  0.2]
xhat = y.detach().ravel()
print(abs(dx.sum().item()) < 1e-12, abs((dx @ xhat).item()) < 1e-12)   # True True

bn.eval()                                    # now uses the running statistics
print(round(bn(torch.tensor([[3.]], dtype=torch.float64)).item(), 3))  # 2.157

t = torch.tensor([0., 2., 4., 6.], dtype=torch.float64)
print(F.rms_norm(t, (4,), eps=1e-12).numpy().round(4))   # [0.     0.5345 1.069  1.6036]

The theorem of Xiong et al. with PyTorch's own blocks (norm_first=True is pre-LN):

import torch
import torch.nn as nn

def last_ffn_grad(norm_first, L, d=64, seed=0):
    """At init: gradient norm of the last block's second FFN weight, and mean |x_L|^2 / d."""
    torch.manual_seed(seed)
    blocks = nn.ModuleList(
        nn.TransformerEncoderLayer(d, 4, 4 * d, dropout=0.0, batch_first=True, norm_first=norm_first)
        for _ in range(L))
    for p in blocks.parameters():
        if p.dim() == 2:
            nn.init.xavier_normal_(p)
    final = nn.LayerNorm(d) if norm_first else nn.Identity()   # pre-LN needs a final norm
    head = nn.Linear(d, 100)
    x, target = torch.randn(8, 16, d), torch.randint(0, 100, (8, 16))
    for blk in blocks:
        x = blk(x)
    loss = nn.functional.cross_entropy(head(final(x)).flatten(0, 1), target.flatten())
    loss.backward()
    return blocks[-1].linear2.weight.grad.norm().item(), (x.detach() ** 2).sum(-1).mean().item() / d

for L in (6, 12, 24, 48):
    g_post, _ = last_ffn_grad(False, L)
    g_pre, stream = last_ffn_grad(True, L)
    print(f"L={L:2d}  post-LN {g_post:.3f}  pre-LN {g_pre:.3f}  pre-LN |x_L|^2/d {stream:.1f}")
# L= 6  post-LN 0.363  pre-LN 0.213  pre-LN |x_L|^2/d 3.8
# L=12  post-LN 0.371  pre-LN 0.162  pre-LN |x_L|^2/d 6.8
# L=24  post-LN 0.373  pre-LN 0.096  pre-LN |x_L|^2/d 13.6
# L=48  post-LN 0.398  pre-LN 0.074  pre-LN |x_L|^2/d 34.1

Post-LN stays between 0.36 and 0.40 at every depth. Pre-LN falls by a factor $2.9 \approx \sqrt 8$ from $L = 6$ to $L = 48$, and its residual stream grows about linearly with depth. These are single random draws, so the last digit is noise (the figure in section 3 averages 5 seeds). The exact constants of Xiong et al. hold in their simplified setting, not exactly here.

6. Where it shows up #

• CNNs. Conv → BN → ReLU is the standard unit. For inference, BN is folded into the conv weights, also before post-training quantization. Fine-tuning with small batches often freezes BN (keeps it in eval mode).
• Small batches. GroupNorm (Wu and He, 2018) computes statistics over groups of channels within one example. It replaces BN when the batch per device is tiny: detection, segmentation, diffusion U-Nets.
• Transformers. Pre-LN with RMSNorm (LLaMA), a final norm, and extra norms for stability such as QK-norm (lesson 12).
• ML for science. Global weather and climate fields are large, so the batch per GPU is often 1 or 2 and BN statistics are useless. ViT-style Earth-system models use LayerNorm; diffusion U-Nets use GroupNorm. DiT's adaptive LayerNorm (adaLN-Zero) computes $\gamma$ and $\beta$ from the condition (the diffusion time step, a class label, or a physical input) and initializes the branch scale at zero, so each block starts as the identity.
• Distributed training. With data parallelism, BN computes its statistics on each GPU's local batch. SyncBatchNorm gives global statistics at the cost of one all-reduce per BN layer; LayerNorm needs no communication.

7. Common mistakes #

• Forgetting model.eval(). BN then normalizes test data with the test batch statistics. With one value per channel, $\hat x = 0$ and the output is $\beta$ (PyTorch's BatchNorm1d raises an error in that case).
• Wrong variance. Normalize with the biased variance (divide by $n$); the running estimate uses $n/(n-1)$.
• Momentum convention. PyTorch's momentum=0.1 is the weight of the new batch value; Keras's momentum=0.99 is the weight of the old running value.
• Wrong LayerNorm axis. LayerNorm in a transformer normalizes over the features of each token, not over the batch and not over the sequence.
• Pre-LN without a final norm. The output head then sees a stream whose scale grows with depth.
• "Residual connections prevent all instability." The forward variance still grows (doubles per block for unscaled branches) unless the branches are normalized or scaled.

8. Check yourself #

1. BatchNorm2d on a tensor of shape $(8, 16, 32, 32)$: how many values enter each mean, how many learned parameters are there, and how many running values?
   Answer

   Each channel's statistics use $8 \cdot 32 \cdot 32 = 8192$ values. There are $16 + 16 = 32$ learned parameters ($\gamma, \beta$) and $16 + 16 = 32$ running values (mean and variance).
2. Show in one line, without the backward formula, that $\sum_i \bar x_i = 0$ for LayerNorm.
   Answer

   LayerNorm gives the same output for $x$ and $x + c\mathbf 1$, so $0 = \frac{d}{dc} \mathcal{L}(x + c\mathbf 1)\big|_{c=0} = \sum_i \partial \mathcal{L}/\partial x_i$.
3. Which component of $g$ does the RMSNorm backward pass remove, and why only one?
   Answer

   The component along $\hat x$ (the radial direction): $\bar x = \frac1r\big(g - \hat x\,\operatorname{mean}(g \odot \hat x)\big)$. RMSNorm is invariant only to $x \mapsto cx$, not to shifts, so the mean direction stays.
4. An eval-mode BN follows a linear layer $z = xW + b$ (row vector $x$). Fold the two into one affine layer $xW' + b'$.
   Answer

   With $s = \gamma/\sqrt{\hat\sigma^2 + \varepsilon}$ (one entry per output channel): $W' = W\operatorname{diag}(s)$ (column $j$ of $W$ times $s_j$) and $b' = s \odot (b - \hat\mu) + \beta$.
5. 24 residual blocks without normalization. Each branch output has variance $\operatorname{Var}(x_l)/24$ and is uncorrelated with $x_l$. What is $\operatorname{Var}(x_{24})/\operatorname{Var}(x_0)$?
   Answer

   Each block multiplies the variance by $1 + 1/24$, so $(25/24)^{24} \approx 2.66$ (close to $e$), instead of $2^{24}$ for unscaled branches.
6. In a pre-LN block, set the output projection of attention and the second FFN matrix (weights and biases) to zero. What does the block compute? And a post-LN block?
   Answer

   Pre-LN: both branches output 0, so the block is the identity and its Jacobian is $I$. Post-LN: $\operatorname{LN}_2(\operatorname{LN}_1(x))$, which is not the identity; the gradient still passes through two LayerNorm Jacobians.

9. Interview angle #

• "BatchNorm at train and test time?" — Batch statistics in training, running averages in eval (unbiased variance in the update). In eval it is a fixed affine map that can be folded into the previous layer.
• "Derive the BatchNorm backward pass." — The five steps of section 3. The result projects out $\mathbf 1$ and $\hat x$; check it with $\sum_i \bar x_i = 0$.
• "Why LayerNorm in transformers?" — Batch dependence: small batches, batch-1 decoding, padding, causal leakage. LayerNorm is per token and the same in training and inference.
• "Why do ResNets train with 1000 layers?" — The identity term in $\prod(I + J_k)$, plus normalization or branch scaling to control the forward variance.
• "Why do post-LN transformers need warmup?" — Their gradients near the output are large and do not shrink with depth at initialization (Xiong et al.); pre-LN avoids this but needs a final norm.

10. Practice #

The notebook for this lesson runs in Colab. Each exercise has a check cell and a folded solution.

1. BatchNorm forward with running statistics, in training and eval mode, checked against nn.BatchNorm1d.
2. LayerNorm and RMSNorm forward, checked against PyTorch.
3. The closed-form BatchNorm backward pass, checked against autograd.
4. Gradient norm per layer of a 30-layer plain MLP and of the same MLP with residual connections, at initialization.
5. Pre-LN and post-LN stacks: gradient norm per block at initialization.

11. Watch and read #

• d2l.ai, §8.5 Batch normalization, including layer normalization and prediction mode.
• Simon Prince, Understanding Deep Learning, free PDF: chapter 11 (§11.2 residual blocks, §11.3 exploding gradients in residual networks, §11.4 batch normalization).
• Chris Bishop and Hugh Bishop, Deep Learning: Foundations and Concepts, §7.4 Normalization.
• Ioffe and Szegedy, Batch Normalization (2015); Ba, Kiros and Hinton, Layer Normalization (2016); Zhang and Sennrich, Root Mean Square Layer Normalization (2019); Wu and He, Group Normalization (2018).
• He et al., Deep Residual Learning for Image Recognition and Identity Mappings in Deep Residual Networks (2016).
• Xiong et al., On Layer Normalization in the Transformer Architecture (2020), §3.3–3.4.