DEV Community

How does PyTorch find every gradient by itself?

Autograd, step by step: PyTorch records the forward pass as a graph, then replays it backward.

In the last article we did backpropagation by hand. We sent the error back through a tiny network, one node at a time, and found each weight's gradient. It worked, but nobody does that for a real network with millions of weights.

In PyTorch you write the forward pass, call L.backward(), and every gradient is there. This article shows what happens in between. The engine that does it is called autograd, short for automatic differentiation.

We'll use the same network and the same numbers as last time, from the CS878 Deep Learning slides.

The forward pass builds the graph

Here is the network in PyTorch:

x  = torch.tensor(1.0)
w1 = torch.tensor(1/3, requires_grad=True)
w2 = torch.tensor(2.0, requires_grad=True)
y  = torch.tensor(0.0)
z  = x * w1
s  = z.relu()
yh = s * w2
L  = (yh - y)**2
Enter fullscreen mode Exit fullscreen mode

The first four lines make the leaves of the graph: the input x, the target y and the two weights. Only the weights have requires_grad=True. That flag tells autograd we will want their gradients.

Every line after that runs one operation, and autograd records it. x * w1 makes a multiply node that holds 1/3. z.relu() adds a ReLU node, still 1/3. Multiplying by w2 gives 2/3, and the squared error gives the loss, 4/9.

You can see the record in what each result prints. z prints as tensor(0.3333, grad_fn=<MulBackward0>). The value is the number we expected. The grad_fn part is new: it is the function that made z.

Typing the forward pass builds the graph, one node per line: the loss is 4/9.

Typing the forward pass builds the graph, one node per line: the loss is 4/9.

grad_fn links make the backward graph

Every result keeps its grad_fn, and every grad_fn keeps links to the functions that made its inputs. Follow those links from the loss and you have the whole graph, pointing backward.

One line of code can hold more than one operation. L = (yh - y)**2 is a subtraction and then a power, so PyTorch records two nodes: L.grad_fn is PowBackward0, and its input comes from SubBackward0.

The weights asked for gradients, so each gets an extra node at the very end of the backward graph, called AccumulateGrad. Think of it as a bucket: it is where the weight's .grad will be stored. The input x and the target y did not ask for gradients, so autograd skips them.

The backward graph: every operation is now its backward function, the arrows point down, and each weight has an AccumulateGrad bucket.

The backward graph: every operation is now its backward function, the arrows point down, and each weight has an AccumulateGrad bucket.

L.backward() walks the graph in reverse

Calling L.backward() starts with a gradient of 1 at the loss and visits the nodes in reverse order. Each node knows one small rule: take the gradient from above and multiply it by its own local derivative.

  • Pow: the local derivative of d² is 2d, so d gets 1 × 2d = 4/3.
  • Sub: the local derivative of ŷ − y with respect to ŷ is 1, so ŷ gets 4/3. The path to y is skipped.
  • Mul: a multiply node swaps values. w2 gets 4/3 × s = 4/9, and s gets 4/3 × w2 = 8/3.
  • ReLU: z was positive, so the gate is open and passes 8/3 through.
  • Mul: w1 gets 8/3 × x = 8/3. The path to x is skipped.

The buckets fill: w1.grad is tensor(2.6667), or 8/3, and w2.grad is tensor(0.4444), or 4/9. Those are the numbers we found by hand.

After L.backward(): the buckets hold 8/3 and 4/9, the same numbers as the hand-worked backprop.

After L.backward(): the buckets hold 8/3 and 4/9, the same numbers as the hand-worked backprop.

When a value is used twice, gradients add

Take y = w * w with w = 3. The multiply node has two inputs, and both are w. On the way back, each path sends 3 (the other input's value). The two contributions meet in w's bucket and add up to 6, which is the derivative of w², 2w. This is the multivariable chain rule: when a value feeds several paths, its gradient is the sum over the paths.

The bucket keeps adding across calls too. Run the forward and backward pass again and w.grad goes from 6 to 12. That's why every training loop calls optimizer.zero_grad() (or sets the gradients to None) before each backward pass.

Running backward twice: the bucket jumps from 6 to 12, so training loops call zero_grad.

Running backward twice: the bucket jumps from 6 to 12, so training loops call zero_grad.

Why go backward?

A network has one loss and millions of weights. You could also push derivatives forward, from one weight up to the loss. That's forward-mode differentiation, and it needs one pass per weight: a million weights, a million passes.

Reverse mode, the one autograd uses, starts from the single loss and reaches every weight in one sweep. It costs about as much as a couple of forward passes, whatever the number of weights. With one output and many inputs, going backward is the only sensible direction.

Reverse mode: one wavefront sweeps down from the loss through every layer and reaches all 36 weights in a single pass.

Reverse mode: one wavefront sweeps down from the loss through every layer and reaches all 36 weights in a single pass.

The graph is rebuilt every step

After backward(), PyTorch frees the graph. The next forward pass records a new one. Because the graph is built while your code runs, ordinary Python works inside a model: an if can add a node on one step and leave it out on the next, and a loop can run a different number of times.

The same code records two different graphs: with s > 0 the  raw `if` endraw  adds a ×2 node.

The same code records two different graphs: with s > 0 the if adds a ×2 node.

Try it in PyTorch

import torch
w = torch.tensor(3.,
  requires_grad=True)
(w * w).backward()
(w * w).backward()
print(w.grad)
Enter fullscreen mode Exit fullscreen mode

It prints tensor(12.), not 6, because the second backward pass adds to the first. Add w.grad = None between the two backward calls and it prints tensor(6.).

Key takeaways

  • Autograd records every operation of the forward pass as a graph. Each result's grad_fn is the function that made it.
  • L.backward() walks that graph in reverse. Each node multiplies the gradient from above by its local derivative.
  • Weights with requires_grad=True collect their gradient in .grad through AccumulateGrad. Tensors without the flag are skipped.
  • Gradients from several paths add up, and .grad keeps adding across steps, so training loops call zero_grad().
  • Reverse mode gets every gradient in one sweep, and the graph is rebuilt on every forward pass, so plain Python control flow works.

This article is part of Deep Learning, Visually by Logit Solutions AI: short animated explainers that cover one idea at a time, each with a written version like this one.

Top comments (0)