Residual connections

A clear path for the gradient

In the vanishing-gradients lesson, the gradient traveled backward through the network. Each layer multiplied it by its own factor. Factors below 1, multiplied many times, made it vanish.

A residual block changes that factor. Two symbols first. $y$ is the block's output and $x$ is its input. $\frac{\partial y}{\partial x}$ is the block's factor: how much $y$ changes when you nudge $x$ a little.

In plain words: the factor for a residual block is 1, plus the branch's own factor. The 1 comes from the skip connection, which passes $x$ through unchanged.

So a stack of plain layers multiplies factors like "something". A stack of residual blocks multiplies factors like "1 + something". Even if every branch passes back a tiny gradient, the 1 carries the signal through.

Here's a small case. Each branch has a factor of 0.05. Twenty plain layers give $0.05^{20}$ (0.05 multiplied by itself 20 times), about $10^{-26}$. That's far too small to learn from. Twenty residual blocks give $1.05^{20}$, about 2.65. The gradient is still a usable size.

plain, residual = 1.0, 1.0
for layer in range(20):
    plain = plain * 0.05              # plain layer: only the branch's factor
    residual = residual * (1 + 0.05)  # residual block: 1 + the branch's factor
print(f"plain    {plain:.1e}")
print(f"residual {residual:.2f}")

This prints plain 9.5e-27 and residual 2.65.

Note: this is a simplification. In a real network, each factor is a whole table of numbers, not one number. The single-number version still shows the main effect.

print(round((1 + 0.1) ** 3, 3), round(0.1 ** 3, 3))

Twenty layers, three ways

Now try it on a real network. The code below trains three networks on the handwritten digits. Each one has 20 hidden layers. Each trains for three epochs (passes over the data) with mini-batch SGD.

• plain: a normal stack of 20 layers.
• residual: the same layers, each wrapped as a residual block.
• residual + pre-norm: residual blocks with a LayerNorm at the start of each branch. LayerNorm, from the normalization lesson, rescales a list of numbers to a standard size. This placement is called pre-norm, and modern transformers use it.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
y = torch.tensor(digits.target)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=0)

class PlainBlock(nn.Module):
    def __init__(self, width):
        super().__init__()
        self.lin = nn.Linear(width, width)
    def forward(self, x):
        return torch.relu(self.lin(x))

class ResidualBlock(nn.Module):
    def __init__(self, width, prenorm=False):
        super().__init__()
        self.norm = nn.LayerNorm(width) if prenorm else nn.Identity()
        self.lin = nn.Linear(width, width)
    def forward(self, x):
        return x + torch.relu(self.lin(self.norm(x)))      # the skip connection

loss_fn = nn.CrossEntropyLoss()
for name, make_block in [("plain", lambda: PlainBlock(32)),
                         ("residual", lambda: ResidualBlock(32)),
                         ("residual + pre-norm", lambda: ResidualBlock(32, prenorm=True))]:
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(64, 32), *[make_block() for _ in range(20)], nn.Linear(32, 10))
    opt = torch.optim.SGD(model.parameters(), lr=0.05)
    for epoch in range(3):
        for xb, yb in DataLoader(TensorDataset(X_train, y_train), batch_size=64, shuffle=True):
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            opt.step()
    with torch.no_grad():
        loss = loss_fn(model(X_train), y_train).item()
        acc = (model(X_val).argmax(1) == y_val).float().mean().item()
    print(f"{name:<20} training loss {loss:.3f}   val accuracy {acc:.3f}")

(As always here, torch is the course's in-browser PyTorch subset.)

Read the three lines one at a time.

• Plain never gets started. Its gradients vanish, and its loss stays at about 2.3. That's the loss you get from guessing evenly among the 10 digits. (It equals ln 10, the natural log of 10.)
• Residual fixes the vanishing, but goes wrong the other way. Every block adds to the signal, so the numbers grow layer after layer until training blows up. The result is nan, short for "not a number".
• Residual + pre-norm learns. The LayerNorm in each branch keeps every block's addition a sensible size. The 20-layer network reaches about 67% after only three epochs, and it's still improving.

Keeping the sum in check

You saw the growth in the worked example: each block added a bit more. Large language models use two standard ways to stop it. Both are common, and many models use both.

Normalize inside the branch

This is pre-norm. The block becomes:

In plain words: standardize a copy of the input, run the copy through the branch, then add the result back onto the original $x$. The branch always sees well-scaled numbers, however big the running total has become.

Here's a small example. Start with 1 and 3, and let the branch be "halve every number". Without the norm, each block adds half of the current total, so the total grows by half every time. After 20 blocks that's $1.5^{20}$, about 3,325 times bigger. With pre-norm, the branch always sees $-1$ and 1, so each block adds the same $-0.5$ and 0.5. After 20 blocks the numbers have moved by only 10 each.

import math

def norm(v):                    # standardize: mean 0, spread 1
    mu = sum(v) / len(v)
    sd = math.sqrt(sum((a - mu) ** 2 for a in v) / len(v))
    return [(a - mu) / sd for a in v]

def f(v):                       # a stand-in branch: halve every number
    return [0.5 * a for a in v]

no_norm = [1.0, 3.0]
pre_norm = [1.0, 3.0]
for block in range(20):
    no_norm = [a + b for a, b in zip(no_norm, f(no_norm))]          # x + f(x)
    pre_norm = [a + b for a, b in zip(pre_norm, f(norm(pre_norm)))]  # x + f(norm(x))
print("no norm: ", [round(v, 1) for v in no_norm])
print("pre-norm:", [round(v, 1) for v in pre_norm])

This prints [3325.3, 9975.8] without the norm and [-9.0, 13.0] with it.

Start the branches small

The second fix is about the starting weights. Set the last layer of each branch close to zero. At the start, every block then does almost nothing, and the deep network behaves like a shallow one. It "grows into" its depth as training goes on.

GPT-2 does this with a rule. It scales those starting weights by $1/\sqrt{N}$, where $N$ is the number of residual layers.

In plain words: the more blocks you stack, the smaller each one's starting weights.

With 16 residual layers, $\sqrt{16} = 4$, so the weights are scaled by 1/4 = 0.25. With 48, they're scaled by about 0.144.

import math
for n in [16, 48]:
    print(n, "residual layers: scale by", round(1 / math.sqrt(n), 3))

The transformer's residual stream

A transformer block is two residual blocks in a row. Here it is in two lines of Python-style code:

x = x + attention(norm1(x))     # tokens exchange information
x = x + mlp(norm2(x))           # each token is processed on its own

The first line lets each token look at the others. That's attention, which you'll meet in Phase 5. The second line runs a small MLP on each token by itself. Both lines are pre-norm residual blocks.

Now zoom out to the whole model. Each token has a list of numbers (a vector) that starts as the token's embedding. Every block reads from that vector and adds a small update to it. At the end, the final vector is turned into logits, the raw scores for each possible next token.

That running vector is called the residual stream. Researchers who study how LLMs work (interpretability researchers) largely do it by asking what each block writes into that stream.

For marketers: it's the difference between each editor rewriting a draft from scratch and each editor adding tracked changes to the same document. With tracked changes, nothing from earlier is lost, and a bad edit is easy to undo.

x = [1.0, 1.0]
for block in range(2):
    update = [0.5, -0.5]
    x = [a + u for a, u in zip(x, update)]
print(x)