7.2. Checking Gradients and Vanishing Gradients

Why gradients matter when debugging

Training works by adjusting parameters using gradients, so when a model does not learn, gradients are one of the first things to inspect. If gradients are consistently near zero, learning will be extremely slow or stop entirely. If gradients explode to very large values, loss can become NaN, updates can destabilize training, and learning curves can look chaotic. Vanishing gradients are especially common in deep networks, saturating activations, and some sequence models, and they often show up as a loss that barely changes while accuracy stays flat.

A quick checklist before you inspect gradients

Make sure you are actually computing gradients. The model must be in training mode when appropriate, you must not wrap the forward pass in torch.no_grad(), and parameters must have requires_grad=True. Ensure you call loss.backward() and that you are not accidentally detaching the loss or intermediate tensors with .detach(), .item(), or conversions to NumPy before backpropagation. Also confirm you are zeroing gradients each step, since stale accumulated gradients can hide issues by making values look larger than they should.

Checking that parameters receive gradients

After loss.backward(), every trainable parameter that influences the loss should have a non-None .grad. A None gradient means the parameter did not participate in the graph, was frozen, or the graph was broken. A gradient that exists but is all zeros indicates either vanishing gradients or a logic issue such as dead activations.

A simple inspection pattern is to iterate parameters and print basic statistics. You want to see that gradients are finite and not trivially zero.

def grad_report(model):
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.grad is None:
            print(f"{name:40s} grad=None")
            continue
        g = p.grad
        print(
            f"{name:40s} "
            f"mean={g.mean().item():+.2e} "
            f"std={g.std(unbiased=False).item():.2e} "
            f"max={g.abs().max().item():.2e}"
        )

If many parameters show grad=None, focus on where the computation graph is being cut. If gradients exist but are extremely tiny, focus on vanishing gradients and scaling issues.

Detecting NaNs and infinities in gradients

NaNs and infinities can appear in gradients even when the forward loss is finite. You should explicitly check for them and stop early when detected.

def assert_finite_grads(model):
    for name, p in model.named_parameters():
        if p.grad is None:
            continue
        if not torch.isfinite(p.grad).all():
            raise RuntimeError(f"Non-finite gradient in {name}")

When non-finite gradients appear, the usual culprits are an excessively large learning rate, unstable loss formulations, numerical overflow in mixed precision, or extreme activations. Fixing this often starts with reducing learning rate and verifying inputs and targets are in expected ranges.

Measuring gradient scale with global norm

A compact way to summarize gradient magnitude is the global $L_2$ norm across all parameters. This is useful for monitoring vanishing or exploding behavior over time.

$$\lVert g \rVert_2 = \sqrt{\sum_i \lVert g_i \rVert_2^2}$$

In PyTorch you can compute it without modifying gradients:

def global_grad_norm(model):
    total = 0.0
    for p in model.parameters():
        if p.grad is None:
            continue
        total += p.grad.detach().pow(2).sum().item()
    return total ** 0.5

If this norm quickly drops toward zero as depth increases or as training progresses, you may be experiencing vanishing gradients. If it spikes upward and correlates with loss spikes or NaNs, gradients are likely exploding.

Understanding vanishing gradients in practice

Vanishing gradients mean that earlier layers, usually closer to the input, receive very small gradient signals. This makes them learn extremely slowly relative to later layers. Common reasons include long chains of multiplications by derivatives whose magnitudes are less than 1, saturating nonlinearities where derivatives approach 0, and poor initialization.

Practically, you will observe that later layers have reasonable gradient magnitudes, while early layers have gradients orders of magnitude smaller. Your grad_report output might show early layers with max around 1e-10 while later layers are around 1e-3 to 1e-2. Loss often decreases a little at first, then plateaus.

Using hooks to inspect gradients layer by layer

Printing .grad on parameters is helpful, but sometimes you want to observe gradients of activations. Hooks can capture gradient norms at key points.

grad_norms = {}
def save_grad_norm(name):
    def hook(grad):
        grad_norms[name] = grad.detach().norm().item()
    return hook
# Example: attach to an activation tensor inside forward
# h is a tensor you want to monitor
# h.register_hook(save_grad_norm("hidden_1"))

To use this, you typically store intermediate tensors in the forward pass and register hooks on them. If activation gradient norms collapse toward zero deep in the network, it indicates gradient signal is not propagating effectively.

Common causes of vanishing gradients you can verify quickly

If you are using saturating activations like sigmoid or tanh in deep stacks, gradients can vanish when activations saturate. You can often confirm saturation by checking activation statistics, for example many values near 0 or 1 for sigmoid, or near -1 or 1 for tanh, combined with small gradients.

Another cause is input and target scaling. Extremely large or small input magnitudes can push activations into saturation or lead to ill-conditioned optimization. You can often detect this by printing mean and standard deviation of inputs per batch and ensuring they look plausible.

Finally, depth plus poor initialization can shrink signals. Even with reasonable activations, very deep networks without modern design choices can struggle.

Practical responses when gradients vanish

Start with the simplest corrective actions that do not require redesigning everything. Try switching to non-saturating activations in hidden layers, such as ReLU family activations, if your current design uses sigmoid or tanh heavily. Consider adding normalization layers where appropriate, since they can stabilize activation scales. Ensure initialization is sensible for your activation choice. If you are using mixed precision, verify that gradient scaling is configured, because underflow can make small gradients become zero.

You can also check that your learning rate is not too small. Vanishing gradients and too small learning rate can look similar from the outside, but gradient inspection will show whether the signal is tiny before the optimizer step.

Distinguishing vanishing gradients from a broken training step

Vanishing gradients produce small but nonzero gradients that often vary across parameters. A broken training step often produces grad=None widely, or produces exactly zero gradients due to a detach or a hard non-differentiable operation in the path to the loss.

Two quick differentiators are whether gradients exist at all, and whether the last layer has meaningful gradients. If even the last layer has grad=None or all zeros, suspect the loss computation, target formatting, or an accidental no_grad context. If the last layer has reasonable gradients but earlier layers do not, suspect vanishing gradients or saturation.

A minimal gradient debugging routine to reuse

A practical routine for a single training iteration is to run forward, compute loss, backward, then immediately check gradient finiteness and print a short report including global norm and a few key layers.

model.train()
optimizer.zero_grad(set_to_none=True)
y_pred = model(x)
loss = criterion(y_pred, y)
loss.backward()
assert_finite_grads(model)
print("global_grad_norm:", global_grad_norm(model))
grad_report(model)
optimizer.step()

If you run this for one batch and the gradients already look wrong, you have saved yourself many hours of guessing based on learning curves alone.

Views: 99