Normalisation and residual connections

The two structural tricks that made deep networks trainable. Both appear in every transformer, so the questions here lead directly into LLM architecture.

Why normalise

As signals pass through layers, activation distributions drift — a phenomenon originally called internal covariate shift. Each layer then has to keep re-adapting to its input distribution, and the loss surface becomes badly conditioned.

Normalisation rescales activations to a stable distribution at each layer. The original covariate-shift explanation has since been questioned; the better-supported account is that it smooths the loss landscape, which is why it permits higher learning rates and faster convergence. Being aware that the original justification is contested is a small but real currency signal.

BatchNorm vs LayerNorm

normalized = gamma * (x - mu) / sqrt(var + eps) + beta

Both apply that formula. They differ in what they average over.

For a (batch, features) tensor: BatchNorm normalises down columns, LayerNorm across rows. That one sentence answers most questions on the topic.

Why transformers use LayerNorm: sequences have variable length and per-device batches are often small, so batch statistics would be noisy and inconsistent. LayerNorm depends only on the sample itself, so it behaves identically at batch size 1 and at batch size 1024, and identically in training and inference. Autoregressive generation happens one token at a time — BatchNorm has no meaningful batch there at all.

The BatchNorm train/eval trap: it uses batch statistics while training and accumulated running statistics at evaluation. Forgetting model.eval() makes predictions depend on whatever else is in the batch — non-deterministic and wrong. Small batches also make the running estimates unreliable, which is why GroupNorm exists as a batch-independent alternative for vision with small batches.

RMSNorm

Modern LLMs typically use a cheaper variant that skips mean-centring:

def rms_norm(x, gamma, eps=1e-6):
    return gamma * x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)

No mean subtraction, no bias term. Empirically it works as well as LayerNorm while doing less work — the re-centring turns out not to matter much. Used in Llama, Mistral and most recent architectures.

Residual connections

x = x + sublayer(x)

Two things this buys, and the second is the one people miss:

A gradient highway. The derivative of x + f(x) with respect to x is 1 + f'(x). That 1 means gradient reaches earlier layers even when f' is tiny — no exponential decay with depth.

An easier optimisation target. The block only has to learn the difference from identity. If identity is the right answer, driving the weights toward zero achieves it. Learning an identity mapping explicitly through several non-linear layers is surprisingly hard, and that difficulty was exactly the degradation problem: before residuals, deeper networks had higher training error, which is an optimisation failure rather than overfitting.

Pre-norm vs post-norm

Where normalisation sits relative to the residual, and a genuinely current interview question.

# Post-norm (original Transformer, 2017)
x = LayerNorm(x + sublayer(x))

# Pre-norm (everything modern)
x = x + sublayer(LayerNorm(x))

Pre-norm keeps the residual path completely clean — nothing is applied to the skip connection, so gradients flow from the output straight to the input unmodified. Post-norm puts a normalisation on that path, which attenuates the gradient at every layer and makes deep stacks hard to train without careful warmup.

Practical consequence: post-norm transformers need aggressive learning-rate warmup and are unstable at depth. Pre-norm trains stably, tolerates higher learning rates, and needs less warmup. Essentially all transformers since GPT-2 use pre-norm, usually with a final normalisation after the last block.

The transformer block, assembled

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.norm1(x))     # pre-norm + residual
        x = x + self.mlp(self.norm2(x))      # pre-norm + residual
        return x

That’s the whole structure, and everything in this file appears in it. Being able to write those three lines from memory and explain each part is a solid answer to “describe a transformer block”.

Dropout in this context

Dropout randomly zeroes activations during training. Interacts with the above in one notable way: dropout before BatchNorm causes a train/eval mismatch, because the running statistics are estimated on dropped-out activations but evaluation runs without dropout. Put dropout after normalisation, or use LayerNorm.

Modern large models use very little dropout — data scale provides the regularisation instead. See ../01_ml_foundations/04_overfitting_regularization.md.

Interview angle

• “BatchNorm vs LayerNorm?” — BatchNorm normalises each feature across the batch; LayerNorm normalises each sample across its features. BatchNorm depends on batch size and composition and behaves differently in training and evaluation; LayerNorm doesn’t.
• “Why do transformers use LayerNorm?” — variable sequence lengths, small per-device batches, and autoregressive generation one token at a time all make batch statistics unusable. LayerNorm is batch-independent and identical in training and inference.
• “What is RMSNorm and why use it?” — LayerNorm without mean-centring or bias. Cheaper, empirically equivalent, standard in current LLMs.
• “What do residual connections do?” — provide an identity path so gradients reach early layers, and reduce each block’s job to learning a residual rather than a full mapping. They fixed the degradation problem, where deeper networks had higher training error.
• “Pre-norm or post-norm, and why?” — pre-norm. It leaves the residual path unmodified so gradients flow cleanly, giving stable training at depth with less warmup. Post-norm was the original 2017 design and is unstable in deep stacks.
• “Why does forgetting model.eval() matter more with BatchNorm than LayerNorm?” — BatchNorm switches from batch statistics to running averages at eval. Leaving it in train mode makes each prediction depend on the rest of its batch. LayerNorm behaves identically in both modes.