Backpropagation

A fresh network has random weights and produces nonsense. This section is about the loop that changes that: run a batch forwards, measure how wrong the answer was, push that error backwards to work out each weight's share of the blame, and nudge every one of them. The optimiser and the learning rate decide how big those nudges are, and reading a loss curve is how you tell whether any of it is working.

The loop#

for epoch in range(epochs):
    for batch in loader:
        optimizer.zero_grad()          # clear last batch's gradients
        output = model(batch.x)        # forward
        loss = criterion(output, batch.y)
        loss.backward()                # backward: compute gradients
        optimizer.step()               # update the weights

Here is the whole thing as something you can run, on data small enough to finish instantly. Two hundred steps of full-batch gradient descent on a two-layer network fitting a linear function with noise added.

import torch
import torch.nn as nn

torch.manual_seed(0)
X = torch.randn(256, 10)
true_w = torch.randn(10, 1)
y = X @ true_w + 0.1 * torch.randn(256, 1)      # a linear rule plus noise

def train(lr, steps=200, zero_grad=True, seed=0):
    torch.manual_seed(seed)
    model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    losses = []
    for step in range(steps):
        if zero_grad:
            opt.zero_grad()
        loss = ((model(X) - y) ** 2).mean()
        loss.backward()
        opt.step()
        losses.append(loss.item())
    return losses

losses = train(0.03)
print([round(losses[i], 4) for i in (0, 5, 20, 50, 199)])
# [5.8346, 4.2446, 0.4697, 0.0706, 0.0402]

Watch the accumulation happen, because the description is vaguer than the numbers. Printing the norm of all gradients before each step, with zeroing removed, gives 3.44, 6.73, 9.81, 12.71, 15.44 — climbing by roughly the size of one step's gradient every step, because that is literally what is being added.

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.SGD(model.parameters(), lr=0.03)

for step in range(1, 6):
    loss = ((model(X) - y) ** 2).mean()
    loss.backward()                                   # no zero_grad anywhere
    norm = torch.cat([p.grad.flatten() for p in model.parameters()]).norm()
    print(step, round(loss.item(), 4), round(norm.item(), 4))
    opt.step()

# 1 5.8346 3.4435
# 2 5.4877 6.7322
# 3 4.8473 9.8099
# 4 3.9718 12.7096
# 5 2.9005 15.4357

print([round(v, 4) for v in train(0.03, zero_grad=False)[::50]])
# [5.8346, 1.2349, 1.0015, 0.7041]   against [5.8346, 0.0706, 0.0523, 0.0447]

Backpropagation#

After the forward pass you have a loss. Backpropagation works backwards through the network applying the chain rule, computing for each parameter how much it contributed to that loss. The framework does this automatically by recording the operations performed.

The smallest network that shows the whole mechanism has one input, one hidden unit and one output, which is three weights and no matrices. Set the input to 2, the two weights to 0.5 and -1.5, the bias on the hidden unit to 0.1, and the target to 1. Forwards: the hidden pre-activation is 0.5 times 2 plus 0.1, which is 1.1; ReLU leaves it at 1.1; the output is -1.5 times 1.1, which is -1.65; the squared error against a target of 1 is 7.0225.

Backwards, one local derivative at a time. The loss changes with the output at a rate of 2 times the error, which is 2 times -2.65, or -5.3. The output weight's own gradient is that times the hidden value: -5.3 times 1.1 is -5.83. To reach the first weight the signal passes back through the output weight, giving -5.3 times -1.5, which is 7.95; through the ReLU, which multiplies by 1 because the unit is active; and then out to the weight by multiplying by the input, 2, giving 15.9. Every one of those numbers is a multiplication of things you already have, which is the entire reason backpropagation is cheap.

x, w1, b1, w2, target = 2.0, 0.5, 0.1, -1.5, 1.0

z = w1 * x + b1              # 1.1
h = max(0.0, z)              # 1.1   ReLU is active
out = w2 * h                 # -1.6500000000000001
loss = (out - target) ** 2   # 7.022500000000002

d_out = 2 * (out - target)           # -5.300000000000001
d_w2 = d_out * h                     # -5.830000000000001
d_h = d_out * w2                     #  7.950000000000001
d_z = d_h * (1.0 if z > 0 else 0.0)  #  7.950000000000001
d_w1 = d_z * x                       # 15.900000000000002
d_b1 = d_z * 1.0                     #  7.950000000000001

def loss_of(w):                       # numerical check on w1
    return (max(0.0, w * x + b1) * w2 - target) ** 2

print(d_w1, (loss_of(w1 + 1e-6) - loss_of(w1 - 1e-6)) / 2e-6)
# 15.900000000000002 15.900000001067838

Two details in that trace are worth keeping. The ReLU passes the gradient through untouched because the unit was active; had z been negative, that one multiplication by zero would have cut the signal to w1 and b1 dead, which is what a dead unit means. And the gradient for w1, 15.9, is three times the gradient for w2, purely because of the input value 2 and the weight -1.5 it passed through. Nothing is deciding which weights matter; the arithmetic is.

Optimisers#

• SGD: step against the gradient by the learning rate. Simple, and sensitive to the learning rate.
• SGD with momentum: accumulate a running direction, which smooths noisy gradients and moves faster along consistent slopes.
• Adam: adapts the step size per parameter. The default for most work, and usually the right first choice.
• AdamW: Adam with weight decay handled correctly. The standard for transformer training.

Learning rate#

Of all the hyperparameters, this one dominates. Too high and the loss oscillates or diverges. Too low and training crawls or settles somewhere poor. It is worth tuning before anything else.

Loss diverges or becomes NaN         learning rate far too high
Loss oscillates without improving    too high
Loss falls very slowly               too low
Loss falls then plateaus early       try a schedule

Those four rows are not a metaphor. Run the same function at four rates and every one of them appears. Note the order in the last row: the loss goes to infinity first and only then to NaN, because once a weight is infinite the next subtraction is infinity minus infinity, which is undefined. A NaN loss is therefore usually the second symptom, and the overflow that caused it happened a step earlier.

for lr in (0.001, 0.03, 0.3, 1.0):
    losses = train(lr)
    print(lr, [round(losses[i], 4) for i in (0, 5, 20, 50, 199)])

# 0.001  [5.8346, 5.7758, 5.6045, 5.2774, 3.704]    barely moving after 200 steps
# 0.03   [5.8346, 4.2446, 0.4697, 0.0706, 0.0402]   healthy
# 0.3    [5.8346, 17.385, 6.1761, 1.1906, 1.0414]   jumped to 17 first, settles 25x worse
# 1.0    [5.8346, inf,    nan,    nan,    nan]      gone

The 0.3 row deserves a second look. It did not diverge, and after 200 steps it reports a loss of 1.04, which someone reading a single number would accept. It is twenty-five times worse than the same model trained at 0.03. A too-high learning rate does not always announce itself; often it just costs you most of the model's quality quietly.

A schedule that warms up and then decays is standard for large models. Warmup avoids enormous early updates when the weights are random; decay lets it settle rather than bouncing around a minimum.

Reading a loss curve#

Training falls, validation falls, gap small     healthy
Training falls, validation rises               overfitting, stop earlier
Neither falls                                  learning rate, bug, or
                                               a model too small
Loss is NaN                                    exploding gradients or
                                               bad input data

What passing that check looks like is worth seeing, because the target is not a low loss but an absurdly low one. Eight examples, a model with several hundred parameters: it should reach a loss that is numerically zero, and if it stalls at 0.3 the problem is in your code rather than in your data.

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.Adam(model.parameters(), lr=0.01)
Xb, yb = X[:8], y[:8]                      # eight examples, nothing else

for step in range(1, 301):
    opt.zero_grad()
    loss = ((model(Xb) - yb) ** 2).mean()
    loss.backward()
    opt.step()
    if step in (1, 50, 100, 200, 300):
        print(step, round(loss.item(), 8))

# 1   12.29048252
# 50  0.06314296
# 100 0.00020716
# 200 1e-08
# 300 0.0

What to memorise, and what to look up#

Memorise the loop in the order zero, forward, loss, backward, step, because saying it out of order in an interview is noticed. Memorise that gradients accumulate and must be zeroed, that the learning rate is tuned before anything else, and the overfit-one-batch check. Everything else — the exact defaults of Adam, which scheduler a paper used, the parameters of a particular data loader — is documentation, and reaching for it in an interview is the correct behaviour rather than a gap.

You know how the gradients get computed. This is about the handful of settings that decide whether a run works at all. Batch size and learning rate move together, so changing one alone destabilises training. Schedules exist because the step size that helps at the start is wrong by the end. Weight decay and early stopping are how you stop a model memorising its training set. These are the knobs you will actually turn, and the usual reason a run that ought to work does not.

Batch size and learning rate move together#

A larger batch gives a less noisy gradient estimate, which allows a larger learning rate. A common heuristic scales the learning rate with batch size. Change one without the other and training can become unstable or unexpectedly slow.

The reason is worth stating precisely, because the heuristic is repeated far more often than it is explained. A gradient computed on a batch is an estimate of the true gradient, and averaging more samples reduces the noise in that estimate — by the square root of the batch size, which is where the square-root version of the scaling rule comes from. A noisier estimate has to be trusted less, so the step taken along it must be shorter. Batch size therefore does not set the step size by itself; it sets how far you can afford to trust one.

Measure it rather than assuming it. The same two-layer network and the same 200 training examples, at three batch sizes and five learning rates, reporting validation loss after sixty epochs.

                lr 0.0125   0.025   0.05    0.1     0.2
batch 8           0.616      0.474   1.060   nan     nan
batch 32          0.606      0.571   0.487   1.135   nan
batch 128         0.581      0.590   0.485   1.625   nan

So use the rule as a direction rather than a formula. Doubling the batch means you should try a larger rate, and expecting the doubled rate to work is optimistic. Note also the row for batch 8, where 0.05 gives 1.060 while 0.025 gives 0.474: the failure at a slightly too large rate is not a crash but a result twice as bad, arrived at without any warning in the training loss.

Schedules#

• Warmup: start small and ramp up, avoiding large destabilising updates while weights are random.
• Cosine decay: reduce the rate smoothly towards zero, letting the model settle into a good region.
• Step decay: drop the rate at fixed points. Simpler, and the timing needs tuning.

A schedule is four lines, and writing them out removes the mystery from the names. Warmup is a straight line from near zero to the target rate over the first few hundred or few thousand steps. Cosine decay then follows half a cosine wave from the target down to zero over the remaining steps. Every large training run you read about is some version of exactly this.

import math

def lr_at(step, base=3e-4, warmup=2000, total=100_000):
    if step < warmup:
        return base * step / warmup                       # straight line up
    progress = (step - warmup) / (total - warmup)
    return base * 0.5 * (1 + math.cos(math.pi * progress))  # cosine down to zero

for step in (0, 1000, 2000, 25_000, 50_000, 100_000):
    print(step, round(lr_at(step), 8))

# 0        0.0
# 1000     0.00015
# 2000     0.0003
# 25000    0.00026104
# 50000    0.00015481
# 100000   0.0

Weight decay#

A penalty that pulls weights towards zero, discouraging reliance on any single large weight. With Adam, the correct implementation is decoupled weight decay, as in AdamW; applying it naively through the loss interacts badly with Adam's per-parameter scaling.

What it buys is measurable and it is not free. On a network with 200 training examples, trained to memorise them completely, the best validation loss reached during training was 0.611 without weight decay, 0.583 at a decay of 0.01, and 0.549 at 0.05. But at 0.05 the training loss only reached 0.406 rather than zero, and the final validation loss was 1.02 — the model had been pulled so hard towards zero that it was still improving, or getting worse, in ways that had nothing to do with the data. Weight decay is a dial with a wrong setting on both sides of the right one.

Early stopping#

Track validation loss and keep the checkpoint where it was lowest. It costs nothing and prevents the model from continuing to overfit after it has stopped improving on unseen data.

The shape of the validation curve in that same experiment is what makes the case. It reached its lowest point, 0.611, at epoch 26; it was still being trained at epoch 200, by which point the training loss was 0.0000 — every one of the 200 examples memorised exactly — and the validation loss had risen to 0.692. The last 174 epochs made the model thirteen per cent worse on data it had not seen, and consumed most of the compute of the run doing it.

Two practical points. Keep the best checkpoint rather than the last one, because those are different models and only one of them is the one you measured. And do not stop at the first epoch where validation rises: early validation curves are noisy, and the standard rule is patience, meaning stop only after some number of epochs with no improvement on the best value seen so far.

best, best_epoch, patience, waited = float("inf"), 0, 10, 0

for epoch in range(max_epochs):
    train_one_epoch(model, optimiser, train_loader)      # your own loop
    score = validate(model, val_loader)                  # lower is better
    if score < best:
        best, best_epoch, waited = score, epoch, 0
        torch.save(model.state_dict(), "best.pt")        # keep the model, not just the number
    else:
        waited += 1
        if waited >= patience:
            break

model.load_state_dict(torch.load("best.pt"))             # go back to the best one

The order to set these in#

1 batch size      largest that fits and keeps the hardware busy
2 learning rate   sweep at that batch size; this dominates everything else
3 schedule        warmup if the model is large or the batch is; then decay
4 weight decay    only once 1 to 3 are settled; try 0, 0.01, 0.1
5 early stopping  always on, patience of a few epochs, keep the best checkpoint

Every deep learning framework rests on one mechanism: it records the operations you perform, then computes the derivative of the result with respect to any number that went into it. This section builds that mechanism in about a hundred lines and trains a neuron with it. It is the shortest route to understanding what loss.backward() is doing, why gradients have to be zeroed between steps, and where the memory in a training run actually goes.

The idea#

Every number produced during the forward pass remembers which numbers it came from and how. Backpropagation walks that record in reverse, multiplying local derivatives along the way, which is the chain rule. PyTorch does exactly this with tensors; here it is with single numbers, so every step is visible. The design follows the well-known micrograd project, which has a video walkthrough.

One design decision in the code below carries more weight than the rest, so it is worth naming before you read it: the reverse pass has to visit nodes in a particular order. A node's gradient is only final once every node that used it has contributed, so it cannot be processed until all of its consumers have been. Sorting the graph so that every node appears after its parents, then walking that list backwards, guarantees exactly that. This is a topological sort, and it is the only algorithm in the whole engine — everything else is arithmetic.

import math

class Value:
    def __init__(self, data, parents=()):
        self.data, self.grad = data, 0.0
        self.parents = parents
        self._backward = lambda: None

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other))
        def backward():
            self.grad += out.grad                  # d(a+b)/da = 1
            other.grad += out.grad
        out._backward = backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other))
        def backward():
            self.grad += other.data * out.grad     # d(ab)/da = b
            other.grad += self.data * out.grad
        out._backward = backward
        return out

    __radd__, __rmul__ = __add__, __mul__

    def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,))
        def backward():
            self.grad += (1 - t * t) * out.grad    # d tanh(x)/dx = 1 - tanh(x)^2
        out._backward = backward
        return out

    def backward(self):
        order, seen = [], set()
        def visit(v):                              # every node after its parents
            if v not in seen:
                seen.add(v)
                for p in v.parents:
                    visit(p)
                order.append(v)
        visit(self)
        self.grad = 1.0
        for v in reversed(order):
            v._backward()

Why += and not =#

A value can feed into several later operations. Its total gradient is the sum of the gradients flowing back from each use. Overwriting instead of adding is the classic bug, and the same reason PyTorch requires zeroing gradients between steps: they accumulate by design.

The smallest case that proves it: let y be x plus x. The derivative is obviously 2. With the plus sign, the add node sends a gradient of 1 back to its first parent and 1 to its second, and both parents are the same object, so x ends with 2. Replace the two plus-equals with plain assignment and x ends with 1, because the second write overwrote the first. Half the answer, no error, and a network that trains slowly for reasons nobody can find.

The other thing the graph explains is memory. Every intermediate value stays alive until the backward pass is done, because its local derivative needs the numbers from the forward pass. That is why a training step needs several times the memory of an inference step at the same batch size, why activation checkpointing works by throwing those intermediates away and recomputing them, and why keeping a tensor in a Python list across steps quietly holds the whole graph that produced it.

Check it#

x, w, b = Value(2.0), Value(-3.0), Value(1.0)
y = (x * w + b).tanh()
y.backward()
print(y.data, w.grad)        # -0.9999092042625951 0.0003631664618877206

import torch
xt, wt, bt = (torch.tensor(v, requires_grad=True) for v in (2.0, -3.0, 1.0))
torch.tanh(xt * wt + bt).backward()
print(wt.grad.item())        # 0.0003630950232036412, the same to four digits

The same idea on arrays#

The Value class is one number per node, which is clear and hopelessly slow. Real frameworks run the identical algorithm on whole arrays, and writing the backward pass of a two-layer network by hand in NumPy takes ten lines. It is worth doing once, because every line is one of the local derivatives you already know, and because the shapes are the part that actually confuses people.

import numpy as np

rng = np.random.default_rng(0)
X = rng.standard_normal((16, 4))          # 16 examples, 4 features
y = rng.standard_normal((16, 1))

params = {
    "W1": rng.standard_normal((4, 8)) * 0.5,
    "b1": np.zeros(8),
    "W2": rng.standard_normal((8, 1)) * 0.5,
    "b2": np.zeros(1),
}

def forward(p):
    z1 = X @ p["W1"] + p["b1"]            # [16, 8]
    h = np.maximum(0, z1)                 # [16, 8]  ReLU
    out = h @ p["W2"] + p["b2"]           # [16, 1]
    loss = ((out - y) ** 2).mean()
    return loss, (z1, h, out)

def backward(p, cache):
    z1, h, out = cache
    n = len(y)
    d_out = 2 * (out - y) / n             # [16, 1]  derivative of the mean squared error
    grads = {}
    grads["W2"] = h.T @ d_out             # [8, 16] @ [16, 1] -> [8, 1], matches W2
    grads["b2"] = d_out.sum(axis=0)       # sum over the batch
    d_h = d_out @ p["W2"].T               # [16, 1] @ [1, 8] -> [16, 8]
    d_z1 = d_h * (z1 > 0)                 # ReLU: pass through where the unit was active
    grads["W1"] = X.T @ d_z1              # [4, 16] @ [16, 8] -> [4, 8], matches W1
    grads["b1"] = d_z1.sum(axis=0)
    return grads

loss, cache = forward(params)
grads = backward(params, cache)
print(round(loss, 6))                     # 1.5287

Now the part that makes this trustworthy rather than plausible. A numerical gradient nudges one parameter, recomputes the loss, and divides — no calculus, no chance of repeating your own mistake. Comparing it against the analytic gradient for every parameter is the gradient check, and it is how anyone writing a new layer knows the backward pass is right.

def numerical(p, name, i, h=1e-5):
    flat = p[name].reshape(-1)            # a view, so writing to it edits the parameter
    original = flat[i]
    flat[i] = original + h
    up = forward(p)[0]
    flat[i] = original - h
    down = forward(p)[0]
    flat[i] = original                    # always put it back
    return (up - down) / (2 * h)

def check(p, grads):
    worst = 0.0
    for name in p:
        analytic = grads[name].reshape(-1)
        for i in range(len(analytic)):
            num = numerical(p, name, i)
            scale = max(1e-8, abs(num) + abs(analytic[i]))
            worst = max(worst, abs(num - analytic[i]) / scale)
    return worst

print(check(params, grads))               # 4.5294809909942426e-10

Then break it deliberately, because a check you have never seen fail is not a check you trust. Delete the ReLU mask from the backward pass — one of the most common real mistakes, since the forward pass still looks right — and run the same comparison.

def backward_buggy(p, cache):
    z1, h, out = cache
    n = len(y)
    d_out = 2 * (out - y) / n
    g = {}
    g["W2"] = h.T @ d_out
    g["b2"] = d_out.sum(axis=0)
    d_h = d_out @ p["W2"].T
    d_z1 = d_h                            # BUG: the ReLU mask is missing
    g["W1"] = X.T @ d_z1
    g["b1"] = d_z1.sum(axis=0)
    return g

print(check(params, backward_buggy(params, cache)))     # 1.0

Finally, use the gradients for what they are for. Fifty lines of NumPy, no framework, and the loss falls from 1.5287 to 0.0739 in 500 steps.

for step in range(500):
    loss, cache = forward(params)
    grads = backward(params, cache)
    for name in params:
        params[name] -= 0.05 * grads[name]

print(round(loss, 6))                     # 0.073931

When to reach for a gradient check#

• You wrote a custom loss or a custom layer. Check it before training anything, on a tiny input.
• A model trains but plateaus higher than it should. A wrong gradient often still reduces the loss, just not to where it could go.
• You changed a reduction, a mask or an axis argument. These are the edits that break a backward pass silently.
• Never in a training loop. It costs two forward passes per parameter, which is why it is run once on a four by eight network rather than on the real one.
• Use double precision while checking. In 32-bit arithmetic the numerical estimate itself is only good to about five digits, so a correct gradient can look like a failed check.

Deep networks were hard to train for years for one reason: the signal telling the early layers how to change either shrank to nothing or blew up on the way back. This section measures that directly, layer by layer, instead of describing it. Then it shows what each of the standard fixes does to those numbers — the activation you choose, how the weights are initialised, normalisation, and residual connections — which is why modern architectures look the way they do.

Measure it#

import torch, torch.nn as nn

def deep_mlp(act, depth=20, width=256):
    layers = []
    for _ in range(depth):
        layers += [nn.Linear(width, width), act()]
    return nn.Sequential(*layers, nn.Linear(width, 1))

for act in (nn.Sigmoid, nn.Tanh, nn.ReLU):
    torch.manual_seed(0)
    model = deep_mlp(act)
    model(torch.randn(64, 256)).pow(2).mean().backward()
    norms = [m.weight.grad.norm().item() for m in model if isinstance(m, nn.Linear)]
    print(act.__name__, f"first layer {norms[0]:.3e}", f"last layer {norms[-1]:.3e}")

# Sigmoid  first layer 4.237e-18  last layer 1.324e+00
# Tanh     first layer 3.593e-06  last layer 7.948e-02
# ReLU     first layer 1.330e-09  last layer 1.901e-02

With sigmoid, gradients in the first layers are typically many orders of magnitude smaller than in the last, because each layer multiplies the gradient by a derivative of at most 0.25. The early layers barely learn. ReLU passes gradients through unchanged where active, and the gap shrinks dramatically.

Notice that ReLU shrinks the gap by nine orders of magnitude and still leaves seven. The activation is not the whole story, which is the honest version of an argument usually stated as though it were. What is left is the size of the weights, and that is set by the initialisation.

Initialisation#

If weights are too large, activations and gradients grow at every layer; too small, they shrink. He initialisation, with standard deviation equal to the square root of 2 divided by the number of inputs, keeps signal size roughly steady through ReLU layers. PyTorch's defaults are chosen with the same goal in mind.

Set the initialisation deliberately and run the same measurement. With He initialisation the first layer's gradient is 4.84 and the last is 24.9 — a factor of five apart rather than ten million, on the identical architecture. Then set every weight from a normal distribution with standard deviation 0.01, which looks harmlessly small, and every gradient in the network reads exactly zero: the signal has underflowed to nothing on the way forward, and there is nothing to send back.

def measure(init_fn):
    torch.manual_seed(0)
    model = deep_mlp(nn.ReLU)
    for m in model:
        if isinstance(m, nn.Linear):
            init_fn(m.weight)
            nn.init.zeros_(m.bias)
    model(torch.randn(64, 256)).pow(2).mean().backward()
    norms = [m.weight.grad.norm().item() for m in model if isinstance(m, nn.Linear)]
    return norms[0], norms[-1]

he = lambda w: nn.init.kaiming_normal_(w, nonlinearity="relu")
tiny = lambda w: nn.init.normal_(w, 0.0, 0.01)

print("%.3e %.3e" % measure(he))      # 4.843e+00 2.486e+01
print("%.3e %.3e" % measure(tiny))    # 0.000e+00 0.000e+00

Residuals and normalisation#

Adding each block's input to its output gives gradients a direct path back through the network, and normalisation keeps each layer's inputs in a steady range. Add both to the experiment above and the first-layer gradient becomes comparable to the last even at depth 100. This is why every modern deep architecture uses them.

class Block(nn.Module):
    def __init__(self, width):
        super().__init__()
        self.norm = nn.LayerNorm(width)
        self.lin = nn.Linear(width, width)
        self.act = nn.ReLU()

    def forward(self, x):
        return x + self.act(self.lin(self.norm(x)))     # the addition is the whole trick

for depth in (20, 100):
    torch.manual_seed(0)
    model = nn.Sequential(*[Block(256) for _ in range(depth)], nn.Linear(256, 1))
    model(torch.randn(64, 256)).pow(2).mean().backward()
    norms = [m.lin.weight.grad.norm().item() for m in model if isinstance(m, Block)]
    print(depth, "%.3e %.3e" % (norms[0], norms[-1]))

# 20   3.456e+00 1.754e+00
# 100  1.388e+02 2.635e+01

The reason is visible in the forward expression. Because each block computes x plus something, the derivative of its output with respect to its input is 1 plus something. The chain rule multiplies those together, and a product of terms near 1 does not vanish the way a product of terms near 0.25 does. The gradient has a road straight back to the first layer that does not pass through any weight matrix.

How to use this on your own model#

• Log the gradient norm per layer, not just the total. A single total hides the shape entirely; the per-layer profile is what tells you where learning has stopped.
• A healthy profile is flat to within an order of magnitude or two across layers. A ratio of 1e6 between first and last means the early layers are effectively frozen.
• Exactly zero anywhere means underflow or a dead path, not a small gradient. Look at initialisation, at a ReLU whose inputs are all negative, and at anything you detached.
• Do this once on a few batches at the start of a run, not continuously. It costs a backward pass you were doing anyway plus one norm per layer.