Computing Gradients with loss.backward in PyTorch

The forward pass and the backward pass are separate steps, and the backward pass is where gradient values come into existence.

Programming

Making a Prediction Before the Trace

When a deep network learns, it adjusts its weights using gradients, which are numbers that describe how sensitive the output error is to each weight. This article uses a small measurement function to look at those gradients layer by layer. Instead of starting with equations about the chain rule, you will run the measurement on a deep network, read the numbers it prints, and find where the learning signal weakens.

The question driving everything that follows is simple: when a deep network learns, how strong is the gradient at each layer, and does that strength hold up as the signal travels back toward the input? Before you run the trace, commit to a prediction. Imagine a deep network with no shortcut connections. You measure one gradient value per weight layer, starting at the layer nearest the output and moving toward the layer nearest the input.

What do you think happens?

In a deep network without skip connections, what pattern do you expect in the mean absolute gradient as you move from the layer nearest the output toward the layer nearest the input?

It grows steadily larger toward the first layer
It stays roughly the same across all layers
It becomes smaller toward the first layer

Reveal answer

Answer:

It becomes smaller toward the first layer.

Without skip connections, the gradients become smaller as the path moves from the last layer toward the first. This sustained shrinking across several layers is the vanishing gradient problem. Hold your prediction loosely: the next sections check it against what the measurement reports.

The Three Stages of the Measurement

The measurement is a function that runs a forward pass, computes a loss against a fixed target, triggers the backward pass, and then reports one gradient value for each weight layer. Treat the printed list as data to interrogate: your job is to decide what it says and to notice when it says something you did not expect.

The first stage is the forward pass. The model receives a sample input and produces an output. The second stage compares that output with a target. Here the target is simply the value zero, and the comparison uses a mean squared error loss. That comparison produces a single scalar loss that says how far the output is from the target. The third stage is calling backward on the loss. This makes PyTorch compute the gradient of the loss with respect to each parameter that needs one. The forward pass and the backward pass are separate steps, and the backward pass is where gradient values come into existence.

What does each stage of the measurement produce for the next one?

What Each Weight Layer Contributes

After the backward pass, the function walks through the model's named parameters and keeps only those whose names contain the word weight. Bias parameters are skipped. For each kept parameter, it takes the gradient tensor, computes the absolute value of every entry, and averages them. A weight matrix of size 3 by 3 has nine gradient values, and the mean of their absolute values collapses those nine numbers into one value for that layer.

Absolute values matter because the sign of a gradient indicates direction, while this measurement asks about strength. Averaging gives one comparable number per layer, which is what makes layer-to-layer comparison possible. The sketch below shows the same logic in PyTorch. It is an illustrative reconstruction of the steps described above, not a listing from the source.

python

import torch
import torch.nn as nn

loss_fn = nn.MSELoss()
output = model(sample_input)
target = torch.zeros_like(output)
loss = loss_fn(output, target)
loss.backward()

for name, param in model.named_parameters():
    if 'weight' in name:
        print(name, param.grad.abs().mean().item())

import torch
import torch.nn as nn

loss_fn = nn.MSELoss()
output = model(sample_input)
target = torch.zeros_like(output)
loss = loss_fn(output, target)
loss.backward()

for name, param in model.named_parameters():
    if 'weight' in name:
        print(name, param.grad.abs().mean().item())

Notes on the code

model and sample_input are assumed to be defined already
zeros_like builds a target tensor of zeros with the same shape as the output
loss.backward() fills in the .grad tensor of each parameter that needs one
Only parameters whose names contain weight are kept, as in the measurement
abs().mean() reduces each gradient tensor to one strength value per layer

Following the Gradient Backward

The printed list is ordered from the first layer to the last, and the natural habit is to read it top to bottom. The signal that matters, however, travels in the opposite direction. Gradients are produced at the output end and carried backward toward the input. So when you read the list, read it as a trail that starts at the last layer and ends at the first.

After loss.backward(), in what order does the gradient travel through the layers, and which layer does it reach last?

When you read the printed values along this path, ask at each step whether the strength is holding up or thinning out. The layer nearest the output is the starting point for the trail, and the layer nearest the input is where the trail ends. Each entry in the list is a summary of one layer's weights, so the list as a whole is a series of readings taken along the backward path.

How does each layer's mean absolute gradient change as the layer index moves from the output back toward the input?

Reading a Printed Gradient List

Worked example

Deciding what a printed gradient list says

You have run the measurement on a deep network without skip connections. The printed list has one mean absolute gradient value per weight layer, ordered from the first layer to the last. What does the list say about the learning signal?

1
Reorder the reading direction
The list starts at the first layer, but gradients are produced at the output end. Read the list as a trail that begins at the last layer and ends at the first.
2
Read the value at the output end
Start with the layer nearest the output. This is the strength of the signal where the trail begins.
3
Follow the trend toward the input
At each layer, ask whether the strength is holding up or thinning out. Look at the trend across several layers, not at one step.
4
Name the pattern
If the values shrink steadily toward the first layer, the source calls this the vanishing gradient problem. The name describes the sustained trend, not one dramatic drop.

Result

The list shows the learning signal weakening as it travels from the last layer toward the first. Without skip connections, the first layers receive a much weaker signal than the layers near the output.

Why Shrinking Gradients Matter

In the model without skip connections, the gradients become smaller as the path moves from the last layer toward the first. That shrinking, sustained across several layers, is the vanishing gradient problem. The name describes the trend, not a single dramatic drop, so you should look for a steady weakening across the whole list rather than one unusually small number.

The learning signal that reaches the early layers is the signal that the measurement is designed to expose. If that signal is weak at the first layer, the first layer receives a weak gradient to learn from. This is why the trend toward the input matters more than the absolute size of the number at the output.

Adding Skip Connections to the Same Measurement

The second model has the same layer sizes, but its shortcut option is turned on through the use_shortcut flag. The random seed is fixed before construction so the two runs begin from comparable starting conditions. The measurement function is then applied to the new model with the same sample input. Keeping the function, the input, and the layer sizes constant isolates one difference: whether skip connections are present.

How does the per-layer gradient pattern differ when the same measurement is run on a network with skip connections versus one without?

With skip connections, the gradient value stabilizes toward the first layer rather than shrinking to a vanishingly small value. Skip connections are described in the source as a technique that stabilizes gradient values by creating shortcuts between layers. The measurement lets you see that stabilization in the printed numbers rather than take it on trust.

Mistakes When Reading the Measurement

Common mistake
Reading the printed list top to bottom as if the signal flows that way
Why it fails
The list runs from the first layer to the last, but gradients are produced at the output and carried backward toward the input.
Do this instead
Read the list as a trail that starts at the last layer and ends at the first.
Common mistake
Treating a large gradient at the output end as a sign of a problem
Why it fails
The last layer tends to have a larger gradient than its neighbors even when skip connections are present.
Do this instead
Focus on the trend toward the first layer, not on the size of the output-end value alone.
Common mistake
Reading the sign of a gradient as its strength
Why it fails
The sign indicates direction. The measurement takes absolute values so that the average reflects strength.
Do this instead
Compare the mean absolute values across layers when judging how strong the signal is.
Common mistake
Assuming the zero target tells you how well the model performs on a real task
Why it fails
The zero target is a convenience. The gradients measure sensitivity to one error signal, not task accuracy.
Do this instead
Use the measurement to study how the signal reaches earlier layers, not to judge learning quality.
Common mistake
Changing the function, input, or layer sizes between the two runs, which removes the controlled comparison

Locating Where Your Prediction Breaks

MEDIUM

You run the measurement on two deep networks with identical layer sizes and the same sample input, one with skip connections and one without. Your prediction was that the gradient shrinks steadily from the last layer toward the first in both networks. Read the printed trail for the network with skip connections. Identify the layer and the direction where the result departs from your prediction, and explain what that departure tells you about the role of the shortcut connections.

Hints

Start at the layer nearest the output and move toward the input, since that is the direction the signal travels
Check whether the trend toward the first layer keeps shrinking or levels off
Remember that a large value at the output end is expected even with skip connections

Your attempt (not saved)

Key Takeaways

The gradient measurement runs a forward pass, computes a loss against a zero target, calls loss.backward(), and reports one mean absolute gradient value per weight layer.
Gradients are produced at the output and carried backward, so the trend that matters runs from the last layer toward the first layer, even though the printed list runs the other way.
Without skip connections, gradients become smaller toward the first layer, which is the vanishing gradient problem.
With skip connections, the gradient value stabilizes toward the first layer rather than shrinking to a vanishingly small value, and the same measurement makes that visible.
A large gradient at the output end is not itself a warning sign; the trend toward the first layer is where to look for a weakening signal.

Related concepts

Vanishing Gradient Problem

More concepts starting with C