Module 5 — Residual connections and layer normalization

We have attention (modules 2 and 3) and position (module 4). Stack twelve of these operations and the model does not train: gradients explode or vanish, loss oscillates or plateaus, whatever hyperparameters you pick. What the paper adds around each attention block is not decorative. This module builds the exact wrapper that turned the Transformer into something you can actually train, and we assemble our first complete layer.

What a residual connection does

A residual connection is the shortest possible piece of arithmetic: add the input back to the output.

$\text{Residual}(f, x) = x + f(x).$

For a Transformer sub-layer $f$ (attention or feed-forward), that turns a stack of $L$ layers into a stack of $L$ corrections applied to an identity: the network is a running sum, not a chain of transformations.

Two consequences make this the single most important trick in modern deep learning.

• Gradients flow. The gradient of the sum with respect to $x$ contains an identity term, so it never vanishes across depth. Removing residuals from a twelve-layer Transformer sends the useful gradient to zero on the first layers within one epoch.
• Layers can no-op. If a sub-layer is not needed, it can learn $f(x) \approx 0$ and the residual short-circuits it. Depth ceases to hurt: adding a useless layer costs runtime, not accuracy.

We wrapped one attention module in module 3. Now we wrap it with a residual:

import torch.nn as nn

class SubLayerConnection(nn.Module):
    def __init__(self, d_model, dropout=0.1):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return x + self.drop(sublayer(self.norm(x)))

That is a "pre-norm" residual. The next section says why we chose it over the paper's original.

Pre-norm versus post-norm

The original 2017 paper places the normalisation after the residual sum:

$\text{PostNorm}(f, x) = \text{LayerNorm}(x + f(x)).$

Every modern Transformer instead uses:

$\text{PreNorm}(f, x) = x + f(\text{LayerNorm}(x)).$

The difference looks cosmetic. In practice it is not.

• Post-norm produces sharper gradients and is easier to fit on a small dataset — but it requires a learning-rate warmup. Without warmup, the initial updates are too large and training diverges within tens of steps. That is the "warmup 4000 steps" line every 2017-era config file inherited.
• Pre-norm trains stably without warmup, at any depth. That is why every LLM larger than a billion parameters uses it. The trade-off is a small drop in final accuracy on some benchmarks, which the industry accepts in exchange for reliable training.

We build our red-thread Transformer with pre-norm. If you copy a 2017-flavoured tutorial that diverges as soon as you increase the depth, this is almost certainly why.

LayerNorm, not BatchNorm

The naming is unfortunate: LayerNorm and BatchNorm are cousins, not synonyms. They normalise along different axes and it matters.

• BatchNorm computes the mean and variance across the batch axis, one per feature. It works in vision because batches are large and their statistics are stable.
• LayerNorm computes the mean and variance across the feature axis, one per token. It does not depend on the batch, which is exactly what we want in NLP where batch sizes vary and, at inference, the batch is often a single sentence.

Concretely, for a hidden vector $h \in \mathbb{R}^d$:

$\text{LayerNorm}(h) = \gamma \cdot \frac{h - \mu(h)}{\sqrt{\sigma^2(h) + \epsilon}} + \beta,$

with $\gamma$ and $\beta$ learned per feature. In PyTorch, nn.LayerNorm(d_model) handles all of it.

BatchNorm sits badly in NLP

A common mistake is to swap the LayerNorm for a BatchNorm to reuse a vision recipe. It trains — until you evaluate on a single sentence. BatchNorm relies on running statistics that no longer match at inference when the batch dimension collapses. LayerNorm is agnostic to batch size, which is the whole point.

The feed-forward block

A Transformer layer is not just attention. Every attention sub-layer is followed by a position-wise feed-forward block: two linear layers with a non-linearity, applied independently to each position.

$\text{FFN}(x) = W_2 \, \text{GELU}(W_1 x + b_1) + b_2.$

The hidden dimension of this block is usually $4d$. For $d = 512$ that is 2048 units. The FFN alone therefore holds $2 \cdot 4 d^2 = 8 d^2$ parameters, versus $4 d^2$ for the multi-head attention module. The FFN carries most of the parameters of a Transformer, not attention. That surprises everyone once.

Why is it there at all? Multi-head attention mixes information across positions but is essentially a weighted average — a linear combination. The FFN adds a non-linear, per-position transformation. Together they cover both axes: attention across positions, FFN within a position.

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff=None, dropout=0.1):
        super().__init__()
        d_ff = d_ff or 4 * d_model
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        return self.net(x)

Assembling one Transformer encoder layer

A single encoder layer is now within reach. It stacks two SubLayerConnection blocks: one wrapping multi-head attention, one wrapping the feed-forward network.

from mha import MultiHeadAttention

class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super().__init__()
        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ff = FeedForward(d_model, dropout=dropout)
        self.sub1 = SubLayerConnection(d_model, dropout)
        self.sub2 = SubLayerConnection(d_model, dropout)

    def forward(self, x, mask=None):
        x = self.sub1(x, lambda y: self.mha(y, mask)[0])
        x = self.sub2(x, self.ff)
        return x

Note how the residual connection wraps a callable: whatever the sub-layer computes, it is added back to x. Module 6 stacks $N$ of these into an encoder, then adds an embedding at the entry and a task-specific head at the exit.

The ResidualBlock visualisation

Course 07's ResidualBlock SVG (src/components/viz/ResidualBlock) shows the exact wiring: input flows both through the sub-layer and around it, and the two paths meet at a plus sign. Import it in your own notes if you find diagrams help:

import ResidualBlock from '@site/src/components/viz/ResidualBlock';

Parameter accounting for one layer

Adding it up for $d = 512$, $h = 8$, $d_\text{ff} = 2048$:

Component	Parameters
Multi-head attention ($4 d^2$)	1,048,576
Feed-forward ($8 d^2$)	2,097,152
Two LayerNorms ($2 \cdot 2 d$)	2,048
Total per layer	~3.15 M

Multiply by 12 layers and you are already at 38 million parameters, before the embedding and output head. That is roughly what the original "base" Transformer contained. GPT-3-scale models push $d$ up to 12,288 and stack 96 layers, and the arithmetic scales as $d^2 \cdot L$.

Read the parameter count of any model you download

Every Hugging Face model card lists the parameter count. Divide it by $12 d^2 L$ and you recover $d$ or $L$ almost exactly — the extras (embeddings, biases, LayerNorm parameters) are less than 5 % of the total. This lets you sanity-check what a "7B" or "70B" model actually is before deploying it.

In summary

• Residual connections turn a stack of layers into a running sum, so gradients flow across depth and useless layers can no-op themselves.
• Pre-norm (x + f(LayerNorm(x))) trains stably without warmup and has replaced the paper's post-norm in every large model.
• LayerNorm normalises along the feature axis, is agnostic to batch size, and is the right choice for NLP; BatchNorm is not.
• The feed-forward block, hidden size $4d$, carries about two thirds of a layer's parameters and adds the per-position non-linearity that attention alone lacks.

Next module: stacking these layers into an encoder and meeting BERT, the first family of Transformers built on the encoder alone.