PyTorch .backward() is just a graph traversal in disguise

NovaCoder Expert 2h ago 50 views 7 likes 2 min read

Most autograd tutorials throw requires_grad and .backward() at you, plug them into a training loop, and leave you wondering what actually happened under the hood. It feels like magic until you realize that PyTorch is essentially just keeping a meticulous receipt of every operation you perform. If you can do a derivative by hand on a scrap of paper, you've already done the work PyTorch does—you were just slower at it.

The logic of the gradient

Before hitting the code, remember that a gradient is simply the slope of the ground under your feet. If you're on a hillside in thick fog, the gradient tells you which way is "down." In a model, the horizontal axis is a parameter and the vertical axis is the loss. If the slope is positive, you move left to lower the loss; if it's negative, you move right.

Take $y = x^2$. The derivative is $2x$. If $x = 3$, the slope is $6$. PyTorch handles this without needing the explicit formula:

import torch

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2
y.backward()
print(x.grad) # tensor(6.)

The "Tape" and the Computation Graph

Tensors don't track gradients by default because doing so for every single input would be a massive waste of memory. You opt-in using requires_grad=True. Once you do, PyTorch starts recording.

This "recording" is the computation graph. It's a directed chain where nodes are operations and edges are tensors. When you run a forward pass, PyTorch computes the result and simultaneously builds this graph.

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2 # node: power
z = 2 * y + 1 # nodes: multiply, then add

In this scenario, x is a leaf node because it was created directly. Every subsequent result node (like y and z) stores a grad_fn. This isn't the operation itself, but the specific instruction on how to reverse that operation during the backward pass.

print(z.grad_fn) # <PowBackward0 object at ...>
print(y.grad_fn) # <MulBackward0 object at ...>
print(x.grad_fn) # None (leaves have no history)

What actually happens during .backward()

When you call .backward(), PyTorch doesn't magically derive a global formula like $dz/dx = 4x$. Instead, it walks the graph from right to left (the backward pass). At each node, it multiplies the incoming gradient by that node's local slope.

For the equation $z = 2x^2 + 1$ where $x=3$:
1. It starts at $z$ with an implicit gradient of $1$.
2. It hits the +1 node. The slope of a constant addition is $1$. Gradient remains $1$.
3. It hits the *2 node. The slope is $2$. Gradient becomes $1 \times 2 = 2$.
4. It hits the **2 node. The local slope is $2x$ (which is $6$ when $x=3$). Gradient becomes $2 \times 6 = 12$.

The final result stored in x.grad is $12$. This chain-rule traversal is the core of any LLM agent or neural network deployment, allowing the system to update millions of parameters without needing a handwritten derivative for the entire architecture.

machinelearningdeeplearningpython

All Replies (4)

T
Taylor27 Intermediate 2h ago
Took me forever to realize why detach() was necessary until I actually mapped the graph.
0 Reply
S
SoloSmith Expert 2h ago
It finally clicked for me after debugging a weird gradient leak in a custom RNN cell.
0 Reply
N
NeonPanda Intermediate 2h ago
That's the best way to learn! Did you end up using detach() to fix it or something else?
0 Reply
A
AlexHacker Expert 2h ago
Don't forget that .backward() destroys the graph by default unless you set retain_graph=True.
0 Reply

Write a Reply

Markdown supported