What roles do residual connections and layer normalisation play in transformer training?
Residual connections preserve an identity path for representations and gradients, so each transformer sub-layer can learn a correction without forcing the whole stack to rebuild the signal. Layer normalisation controls the mean and scale of each token's hidden features, improving conditioning independently of batch size; their placement matters, with Pre-LN usually easier to optimise and Post-LN more sensitive to training settings.
How to think about it
Residual connections give each transformer sub-layer an identity route, so it can learn a correction while information and gradients pass through depth. Layer normalisation keeps each token’s hidden features at a well-conditioned scale; together they make deep stacks easier to optimise, but they solve different problems and their placement matters.
Why residual connections matter
A transformer block is one repeated attention-and-feed-forward unit. A sub-layer is one of those operations, such as self-attention or the feed-forward multilayer perceptron. Instead of replacing the input with the sub-layer’s output, a residual connection adds the input back:
output = input + sub_layer(input)
The vector carried from block to block is often called the residual stream. The useful mental model is not “the next layer rewrites the representation.” It is “the next layer proposes an update.”
Imagine a 24-block decoder. Each block has an attention sub-layer and an MLP sub-layer, so there are 48 residual additions. If the attention sub-layer discovers that a token should attend more strongly to a previous quotation mark, it adds that correction to the existing representation. It does not need to reconstruct every fact that earlier blocks already encoded.
Without the addition, a deep stack repeatedly composes transformations:
x_next = F(x)
The gradient then passes through a product of many local derivatives. If those derivatives mostly shrink values, early layers receive tiny gradients. If they mostly amplify values, training becomes unstable. Either way, the problem gets worse as depth increases.
With a residual connection, a Pre-LN block has the form:
x_next = x + F(LayerNorm(x))
Its local derivative is approximately:
I + J_F J_LN
Here, I is the identity matrix, and a Jacobian is the matrix describing how a function’s outputs change when its inputs change. The identity term is the important part. It gives the gradient a route that does not require multiplying through the entire sub-layer transformation.
That does not mean gradients can never vanish or explode. The branch term can still be badly scaled, and the identity contributions can interact destructively. It means every block starts with a direct “do nothing” route. If a newly added block is not useful yet, the network can behave more like the shallower network beneath it instead of having to learn a completely new transformation from scratch.
This is why residual networks made very deep vision models practical, and why the same idea is central to transformers.
Common misconception: residual connections do not make the gradient equal to one everywhere. They add an identity contribution. The actual gradient also includes the sub-layer branch, normalisation, attention, dropout, and parameter scaling. In a Post-LN block, the distinction is even more important because the normalisation comes after the addition.
What layer normalisation does
Layer normalisation, or LayerNorm, normalises the features inside one token representation. It computes statistics across the hidden dimension, usually called d_model, rather than across the examples in a batch.
For a token vector with d features:
mu = (1 / d) * sum(x_i)
var = (1 / d) * sum((x_i - mu)^2)
LayerNorm(x_i) = gamma_i * (x_i - mu) / sqrt(var + eps) + beta_i
gamma and beta are learned scale and shift parameters. eps is a small constant that prevents division by zero.
Take the four-feature vector [2, 4, 6, 8]. Its mean is 5, and its variance is 5. With gamma = 1 and beta = 0, LayerNorm produces approximately:
[-1.342, -0.447, 0.447, 1.342]
Now shift every feature by 18, giving [20, 22, 24, 26]. The mean changes to 23, but the deviations from the mean are unchanged. LayerNorm therefore produces the same normalised values.
That removes arbitrary offset and scale from the input seen by the next sub-layer. This matters because hidden-feature scale affects several sensitive operations:
- Attention forms dot products between queries and keys. Large or inconsistent magnitudes can make the softmax extremely sharp.
- MLPs pass activations through nonlinear functions whose useful operating ranges depend on scale.
- Gradients depend on the scale of activations and weights. A layer that receives wildly different scales from different examples is harder to optimise.
LayerNorm gives each token a more predictable numerical starting point. It does not make all tokens identical. Their relative feature patterns remain, and the learned gamma and beta parameters let the model restore useful scales and offsets.
LayerNorm also does not use running averages. The statistics for a token are computed from that token’s current hidden vector during both training and inference. A token can be normalised in a batch of 64, a batch of 1, or during one-token-at-a-time decoding without changing the rule.
| Property | LayerNorm | BatchNorm |
|---|---|---|
| Statistics | Hidden features of one token | Usually examples in a batch, with layout-dependent details |
| Depends on batch size | No | Typically yes |
| Running inference statistics | No | Usually yes |
| Awkward for autoregressive decoding | No | Often |
| Handles variable sequence lengths naturally | Yes, per token | Requires careful layout, padding, and masking choices |
It is too strong to say BatchNorm mathematically fails whenever the batch size is one. Its behaviour depends on the tensor layout and how many values remain in its normalisation axes. The practical issue for language models is that BatchNorm couples examples or positions and usually has a train-versus-inference distinction. A decoder generating one token at a time is a poor fit for that dependency. LayerNorm avoids it.
How the two mechanisms fit together
A common Pre-LN transformer block looks like this:
for each block:
x = x + dropout(attention(layer_norm(x)))
x = x + dropout(mlp(layer_norm(x)))
x = final_layer_norm(x)
The residual connection preserves the stream. LayerNorm prepares the stream for each branch. Dropout is applied to the learned update, not to the identity route, so the original representation still has a path through the block.
There are two major normalisation placements.
Pre-LN
Pre-LN puts normalisation before the sub-layer:
x_next = x + F(LayerNorm(x))
This gives the cleanest identity path through the block. It is usually easier to train at substantial depth and is less dependent on a carefully tuned learning-rate warm-up.
The trade-off is that the residual addition itself is not immediately normalised. Across many blocks, the residual stream can accumulate updates and change scale. A final normalisation layer is common, but it does not make every intermediate residual state normalised.
Post-LN
Post-LN puts normalisation after the residual addition:
x_next = LayerNorm(x + F(x))
This was used in the original Transformer architecture. It normalises the updated representation at every block, which can be attractive for the final behaviour of the network. But the gradient from one block to the previous block passes through the LayerNorm Jacobian after every addition:
J_LN (I + J_F)
That is not the same as an unnormalised identity route. Deep Post-LN models can therefore be more sensitive to initialisation, learning rate, and warm-up. A training run may look fine for a few steps and then produce a loss spike when the learning rate reaches its peak.
Many current large language models use Pre-LN or a close variant. One common variant is RMSNorm, or root-mean-square normalisation, which rescales by the vector’s root-mean-square value but does not subtract the mean. RMSNorm is not identical to LayerNorm, but it serves a similar conditioning role in the block.
Trade-offs and failure modes
Neither technique is a universal cure.
| Symptom seen in training | Likely issue |
|---|---|
| Early-layer gradients are much smaller than late-layer gradients, and a deeper model learns worse than a shallower one | A residual addition is missing, misplaced, or accidentally detached |
| Activation RMS grows steadily with depth, attention logits become extreme, then the loss spikes or becomes NaN | Normalisation is missing, applied on the wrong dimension, or the residual stream is poorly scaled |
| Loss becomes unstable in the first few hundred steps, especially in a deep model | Post-LN training may need a smaller learning rate, warm-up, or different initialisation |
| A model trains but quality falls at the same parameter count after switching to Pre-LN | The architecture or final normalisation may need retuning; easier optimisation does not guarantee the best final quality |
A useful debugging habit is to log activation norms and gradient norms by layer, not only the total loss. If the first blocks have gradients near zero while the last blocks have ordinary gradients, inspect residual paths and normalisation placement before changing the optimiser. If norms grow with depth, confirm that LayerNorm is operating over the hidden dimension for each token, not accidentally over the batch or sequence axes.
The senior answer also acknowledges that residuals and LayerNorm are not interchangeable. Residuals address the depth and gradient-routing problem. LayerNorm addresses the scale and conditioning problem. Removing either can hurt, but for different reasons. Gradient clipping may hide an exploding-gradient symptom; it does not replace a sound residual and normalisation design.
What they’ll ask next
“Does LayerNorm depend on sequence length?”
No. It computes statistics over the hidden features of each token independently, so adding more tokens does not change the LayerNorm statistics for an existing token. Sequence length still affects self-attention, memory use, and the context available to each token. LayerNorm is independent of those effects, not the entire transformer.
“Why not use BatchNorm instead?”
BatchNorm couples normalisation to other examples or positions and normally relies on running statistics at inference. That is awkward for variable-length text and for autoregressive decoding with batch_size = 1. LayerNorm works on each token’s feature vector directly. BatchNorm can be made to work in some sequence architectures, but it is not the natural default for standard transformers.
“Do residual connections eliminate vanishing gradients?”
No. They create an identity contribution that makes vanishing less likely, especially in Pre-LN blocks, but the complete gradient still depends on the branch Jacobians, normalisation, attention logits, parameter scales, and optimisation settings. A model can still have exploding gradients, poorly scaled updates, or ineffective early layers.
Say this in the interview
“Residual connections let every transformer sub-layer learn an update while preserving an identity path for information and gradients; LayerNorm independently stabilises each token’s hidden-feature scale, and the Pre-LN or Post-LN placement determines how much optimisation stability the model gets.”