Checking that gradients reach every layer

Setup#

pip install torch

Captured with torch 2.5.1 on CPU; seeded and deterministic.

The dry-soil detector#

This model contains a planted bug — a detach in the middle of the forward pass:

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.first = nn.Linear(10, 32)
        self.middle = nn.Linear(32, 32)
        self.last = nn.Linear(32, 2)

    def forward(self, x):
        h = torch.relu(self.first(x))
        h = torch.relu(self.middle(h)).detach()   # the bug: detach cuts the wire here
        return self.last(h)

torch.manual_seed(0)
model = Net()
loss = nn.CrossEntropyLoss()(model(torch.randn(16, 10)), torch.randint(0, 2, (16,)))
loss.backward()

for name, p in model.named_parameters():
    if p.grad is None:
        print(f"{name:15s} grad: MISSING")
    else:
        print(f"{name:15s} grad norm: {p.grad.norm().item():.6f}")

first.weight    grad: MISSING
first.bias      grad: MISSING
middle.weight   grad: MISSING
middle.bias     grad: MISSING
last.weight     grad norm: 0.184076
last.bias       grad norm: 0.079574

The map reads like the pipe diagram: everything behind the detach is dry. Note the direction — gradients flow backward, so a cut after middle starves middle and first, while last (between the cut and the loss) stays healthy.

Remove the .detach() and rerun:

first.weight    grad norm: 0.083987
first.bias      grad norm: 0.019576
middle.weight   grad norm: 0.142110
middle.bias     grad norm: 0.038178
last.weight     grad norm: 0.184076
last.bias       grad norm: 0.079574

All soil wet. Notice last's numbers are identical in both runs — its water never depended on the broken joint.

Reading the three verdicts#

MISSING (grad is None) — no gradient ever arrived. The wire is cut: a detach() or .data in the path, requires_grad=False on those parameters, an operation autograd cannot pass through (argmax, round, comparisons, integer casts, a trip through NumPy), or the layer genuinely is not in the path — defined in __init__ but never called in forward.

Norm exactly 0.0 — a gradient arrived and it is zero. Different disease: dead ReLUs (every input negative), multiplication by an exact zero mask, or a loss term weighted to nothing.

Norms present but vanishingly small in early layers — each layer passes a fraction of the signal and the fractions compound; ten layers of halving leaves a thousandth. That is the vanishing gradient problem — see backpropagation — treated with normalisation, residual connections, or gentler depth.

Common mistakes#

Running the check before any backward(). Every grad is None before the first backward pass, and the printout screams catastrophe. The check only means something after a real loss has flowed.

Forgetting grads accumulate. Checking after several backwards without zero_grad reads summed history, not this step's flow. One clean zero_grad, one forward, one backward, then read.

Confusing intentional freezes with bugs. Frozen backbone layers in transfer learning should read MISSING — the check reports flow, and you supply the judgement about which dryness is deliberate.

Using .data to dodge an autograd error. It silences the error by cutting the graph — manufacturing exactly the bug this lesson hunts. The honest alternatives are covered in detach vs no_grad.

Try it yourself#

Break the model three new ways, predicting the gradient map before each run: set self.first.requires_grad_(False); replace the middle activation with (h > 0).float(); add a fourth layer to __init__ without calling it in forward. Then check what named_parameters shows for the fourth case and note which verdict it produces.

What to learn next#

• Detach vs no_grad — the legitimate uses of the wire-cutter that caused today's bug.
• Memory that grows every epoch — the opposite failure: graphs kept alive too long.
• Backpropagation — the flow being diagnosed, from first principles.