Module 13: Residual Connections and Layer Normalization

Why this matters

Attention (modules 04-11) and the FFN (module 12) are the two computational sub-layers of a transformer block — but stacking many of them (module 15 stacks dozens) doesn't work reliably without two unglamorous but load-bearing additions: a residual (skip) connection around each sub-layer, and layer normalization. This module verifies both are doing real, measurable work, not boilerplate: a 30-layer deep network's gradient at the input collapses by 8 orders of magnitude without residual connections, and layer normalization is confirmed to force every token's representation to a fixed mean/standard-deviation regardless of its raw input scale.

Concepts

Residual connections: output = x + sublayer(x), not output = sublayer(x)

Every attention and FFN sub-layer in a real transformer is wrapped like this:

 output = x + Sublayer(x)          NOT:   output = Sublayer(x)

The + x — adding the sub-layer's input back to its output — is the residual (or "skip") connection. It looks almost too simple to matter. Verified directly below, it's the difference between a deep stack being trainable at all and a deep stack whose gradient vanishes to nothing at its earliest layers.

Verified: residual connections prevent gradient collapse at depth

Two 30-layer stacks of identical blocks — one with a residual connection, one without — measuring the gradient magnitude that actually reaches the very first layer's input:

import torch
import torch.nn as nn

d, depth = 32, 30

class BlockNoResidual(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(d, d)
    def forward(self, x):
        return torch.tanh(self.lin(x))

class BlockResidual(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(d, d)
    def forward(self, x):
        return x + torch.tanh(self.lin(x))

def measure_grad(block_cls):
    blocks = nn.ModuleList([block_cls() for _ in range(depth)])
    x = torch.randn(1, d, requires_grad=True)
    h = x
    for b in blocks:
        h = b(h)
    h.sum().backward()
    return x.grad.abs().mean().item()

print(measure_grad(BlockNoResidual))
print(measure_grad(BlockResidual))

Verified output:

No residual, depth 30:   1.84e-08
With residual, depth 30: 4.73

Roughly 8 orders of magnitude difference. Without residual connections, by the time a gradient signal has backpropagated through 30 stacked tanh-activated layers, it has shrunk to almost nothing — the earliest layers would receive essentially no useful training signal at all. With residual connections, the gradient at the same depth stays a normal, usable magnitude. This is the same underlying vanishing-gradient phenomenon track 00 module 04 covered for RNNs, verified here for depth (many stacked layers) rather than sequence length — residual connections are the transformer's specific fix for this problem at depth.

Why residual connections work: gradient has a direct path

The mechanism is simple once stated: d(x + f(x))/dx = 1 + df(x)/dx. The +1 term means a gradient can always flow straight through the + x path, completely bypassing whatever f (attention or the FFN) does — even if f's own gradient shrinks toward zero (as tanh-activated layers' gradients tend to, especially stacked many times), the +1 guarantees an unobstructed route back to earlier layers. This is exactly why the gap above is so large: without the +1 path, every layer's gradient must pass entirely through f's own gradient, multiplying together and shrinking at every step.

Layer normalization: forcing every token to a fixed scale

Layer normalization rescales each token's vector (across its own feature dimensions — module 01, module 02) to have mean 0 and standard deviation 1, regardless of the input's original scale:

ln = nn.LayerNorm(8)
x = torch.randn(3, 8) * 5 + 10   # arbitrary scale and shift per row
out = ln(x)
print(x.mean(dim=-1), x.std(dim=-1))
print(out.mean(dim=-1), out.std(dim=-1, unbiased=False))

Verified output:

input means:  [7.59, 11.75, 10.81]     input stds:  [4.91, 2.83, 6.78]
output means: [~0.0, ~0.0, ~0.0]       output stds: [1.0, 1.0, 1.0]

No matter how differently-scaled the three input rows started (means ranging from 7.6 to 11.8, standard deviations from 2.8 to 6.8), every output row lands at (effectively) exactly mean 0, standard deviation 1. This matters directly for the residual-connection story above: as values pass through many stacked blocks, each adding its sub-layer's output back onto an accumulating stream, that stream's scale could otherwise drift arbitrarily large or small across dozens of layers — layer normalization resets each token's scale at every step, keeping values in a well-behaved range the way module 06 verified scaling was needed to keep softmax well-behaved.

Pre-norm vs. post-norm: where normalization sits relative to the residual

Two real, both-used conventions differ in exactly where layer normalization is placed:

 Post-norm (original 2017 paper):    output = LayerNorm(x + Sublayer(x))
 Pre-norm (most modern models):      output = x + Sublayer(LayerNorm(x))

Pre-norm puts the raw, un-normalized x on the direct residual path (the same +1-gradient path verified above), while post-norm applies normalization after the addition, meaning that direct path is no longer purely linear. This difference has a real, documented consequence for training stability at large depth — pre-norm is the more common choice in modern large models specifically because it preserves cleaner gradient flow through the residual path, though it isn't covered numerically in this module's exercises (a good candidate for exercise 3 below).

Reference

 Concept                    What it does                     Verified
 ────────────────────       ───────────────────────────       ────────────────────
 Residual connection         output = x + Sublayer(x)           gradient magnitude
                                                                 at depth 30:
                                                                 1.84e-08 (no
                                                                 residual) vs. 4.73
                                                                 (with residual)
 Layer normalization         normalize each token's own          arbitrary input
                             features to mean 0, std 1            means/stds all
                                                                  mapped to
                                                                  ~mean 0, std 1

 Post-norm:  LayerNorm(x + Sublayer(x))    — original paper
 Pre-norm:   x + Sublayer(LayerNorm(x))     — most modern large models

Hands-on exercises

Exercise 1 — reproduce the gradient-collapse measurement

Run the exact measure_grad comparison above at a few different depths (try 5, 15, 30, 60) and confirm the gap between residual and no-residual widens as depth increases — connecting the specific mechanism (+1 gradient term compounding across layers) to why the effect gets worse, not just different, with more layers.

Exercise 2 — reproduce the layer-normalization verification

Run the exact nn.LayerNorm example above with your own random input values at a different scale and shift, and confirm the output always lands at mean ~0, std ~1 per row regardless of the input's original scale.

Exercise 3 — measure the pre-norm vs. post-norm gradient difference

Build two versions of a deep stack (5-10 layers is enough to see an effect) — one applying LayerNorm after the residual addition (post-norm) and one applying it before the sub-layer, inside the residual path (pre-norm) — and measure the gradient magnitude reaching the input in each, similar to exercise 1's method. Connect what you find to why pre-norm has become the more common choice for very deep modern models.

Independent challenge

A researcher proposes removing layer normalization entirely from a transformer, arguing "residual connections already solve the gradient-flow problem, so normalization is redundant." Using this module's verified findings about what each component specifically does, write two or three sentences on why this reasoning conflates two different problems.

Common mistakes & troubleshooting

• Assuming residual connections and layer normalization solve the same problem. Verified above: residual connections fix gradient magnitude at depth; layer normalization fixes value scale consistency per token — removing either one leaves the other problem unaddressed.
• Underestimating how severe gradient collapse gets without residual connections. Verified above: an 8-order-of-magnitude difference at just 30 layers — real modern models are often far deeper, making this effect (without a fix) considerably worse still.
• Confusing pre-norm and post-norm as interchangeable stylistic choices. They place normalization differently relative to the residual path, with a real, documented difference in gradient behavior at large depth — not a cosmetic preference.
• Assuming layer normalization normalizes across the batch or across positions. It normalizes each token's own feature vector independently (verified: each row gets its own mean/std computed across its features) — a different axis than batch normalization (common in other architectures, not used the same way here).

Checkpoint quiz

1. What is the residual-connection formula, and what did the verified gradient experiment show about depth-30 stacks with and without it?
2. What mathematical fact about d(x + f(x))/dx explains why residual connections prevent gradient collapse?
3. What did the verified LayerNorm experiment show about outputs from inputs with very different means and standard deviations?
4. What is the difference between pre-norm and post-norm placement, and which is more common in modern large models?
5. Why is "residual connections make layer normalization redundant" an incorrect claim?

Further reading & sources

• Deep Residual Learning for Image Recognition (He et al., 2015) - the original residual-connection paper (ResNet); the output = x + f(x) pattern this module verifies numerically, first demonstrated at scale for very deep networks.
• Layer Normalization (Ba et al., 2016) - the original layer normalization paper, defining the per-token mean-0/std-1 normalization verified in this module.
• On Layer Normalization in the Transformer Architecture (Xiong et al., 2020) - directly analyzes the pre-norm vs. post-norm gradient-stability difference this module's exercise 3 asks you to measure.
• Attention Is All You Need (Vaswani et al., 2017) - Section 3.1 specifies the original post-norm placement (LayerNorm(x + Sublayer(x))).
• Track 00, Module 04: RNNs/LSTMs and the Sequence Problem - covers the same underlying vanishing-gradient phenomenon in a different context (sequence length rather than network depth).

Next

Module 14: The Complete Transformer Block