Residual connections and normalization: keeping deep models trainable

What should a deep layer do when it has no useful change to make? A residual connection makes “keep the incoming representation” easy to express. Normalization addresses a different question: what scale should the next transformation receive?

Before you begin: You understand a transformer block and can calculate an average.

This analogy describes an operation, not a memory of perfect past answers. The numbers can still grow, useful information can still be damaged, and training can still fail.

Add a correction to the current state

A residual block computes y = x + F(x). Here x is the incoming representation and F is an attention or feed-forward transformation. The two terms must have matching shapes, unless a projection explicitly aligns them.

If the transformation learns a small correction, the block stays close to the identity operation. During backpropagation, the addition also creates a direct gradient path. The derivative contains an identity term as well as the derivative through F. This helps optimization; it does not prove that gradients can never vanish or explode across a whole model.

Turn the learned branch off

Take x = [2, -1] and a branch output F(x) = [0, 0]. A residual addition returns [2, -1]. Now compare the two arrangements later in this lesson. Pre-norm still returns x when the learned branch is zero. Post-norm returns Norm(x), which generally differs from x.

Before continuing, calculate the mean of [2, -1] and describe the direction of its centered values. Why is “both blocks preserve the input exactly” an incorrect explanation?

Now inspect the normalization arithmetic itself, including the small epsilon that keeps a constant vector from causing division by zero.

What survives when the learned branch contributes nothing?

Give the next operation a manageable scale

Layer normalization usually normalizes the features of each token separately. It subtracts their mean, divides by their standard deviation with a small stabilizing epsilon, and applies learned scale and bias parameters.

normalized = (x - mean(x)) / sqrt(variance(x) + epsilon)
output = learned_scale * normalized + learned_bias

For the features [1, 2, 3], the mean is 2 and the population variance is two-thirds. Before the learned transformation, the normalized values are approximately [-1.225, 0, 1.225]. The middle feature is now zero because it was exactly at the mean, not because it became irrelevant.

# Runnable: Python 3, standard library.
import math

x = [1.0, 2.0, 3.0]
mean = sum(x) / len(x)
variance = sum((v - mean) ** 2 for v in x) / len(x)
normalized = [(v - mean) / math.sqrt(variance + 1e-5)
              for v in x]
assert abs(sum(normalized)) < 1e-12
print([round(v, 3) for v in normalized])

Unlike batch normalization, this computation does not estimate statistics across other examples in a batch. That makes the operation independent of which other sentences happen to share the batch. The learned scale and bias still matter: the final output need not have mean zero or variance one.

RMSNorm uses the root mean square of the features and does not subtract their mean. It is a distinct normalization rule, not LayerNorm with a renamed variable.

Where does normalization go?

Pre-norm keeps the main residual path outside the normalization operation and often makes deep transformer optimization easier. Post-norm was used in the original Transformer. Training behavior also depends on initialization, learning rate, depth, and other design choices; “pre-norm always wins” is too strong.

Do not explain normalization only as fixing “internal covariate shift.” That phrase is an incomplete account of why normalization helps optimization. The concrete operation and its effect on scale are safer starting points.

Investigate a failure

A model starts producing non-finite losses after an architectural change. Check activation and gradient norms around each sublayer. Verify epsilon is positive, the normalized axis is the feature axis, mixed-precision reductions are appropriate, and residual shapes agree. A graph showing a loss spike tells you when trouble began; these checks help locate where.

Exercise: add 100 to every value in the example. Which normalized result stays approximately the same: LayerNorm or RMSNorm?

Practice with feedback

import torch, torch.nn as nn

class Block(nn.Module):
    def __init__(self, d, pre=True):
        super().__init__()
        self.norm = nn.LayerNorm(d)
        self.f = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.pre = pre
    def forward(self, x):
        return x + self.f(self.norm(x)) if self.pre else self.norm(x + self.f(x))

def grad_at_first(pre):
    torch.manual_seed(0)
    blocks = nn.Sequential(*[Block(128, pre) for _ in range(40)])
    x = torch.randn(4, 16, 128, requires_grad=True)
    blocks(x).pow(2).mean().backward()
    return x.grad.norm().item()

print('pre-norm ', grad_at_first(True))
print('post-norm', grad_at_first(False))
# pre-norm gradients arrive at the input with a healthy magnitude; post-norm
# at 40 layers is orders of magnitude smaller at initialisation, which is
# exactly the regime that needs a warmup schedule to train at all.

# Identity check: zero the learned branch and confirm the block is a no-op.
b = Block(128, pre=True)
for p in b.f.parameters():
    torch.nn.init.zeros_(p)
x = torch.randn(2, 3, 128)
assert torch.allclose(b(x), x)
# In the post-norm version the same test fails: norm(x + 0) is not x. The
# identity path is the thing pre-norm protects, and it is why I would choose
# it for anything deep.

Next, move outside the block and examine the tokenizer that decides which pieces of text become these representations.

Sources

See the definitions in Layer Normalization, Root Mean Square Layer Normalization, and the optimization analysis in On Layer Normalization in the Transformer Architecture.