How Autodiff Actually Works — Unpacking the PyTorch Magic

Why does writing loss.backward() hand you derivatives for millions of parameters? We build up computation graphs, the chain rule, and forward vs. reverse mode from zero — then write a working 40-line autograd engine.

The line that shouldn't work

In any deep learning script, one line quietly does something remarkable:

loss = F.mse_loss(model(x), t)
loss.backward()          # every parameter now has a gradient

Inside model there are matrix products, activations and normalization layers stacked dozens deep, and easily millions of parameters. Yet nowhere in the file did anyone write a derivative. Nobody derived d(loss)/d(w_483921) by hand. Run it anyway, and every w.grad holds the right number.

The short version: there is no magic. There is only the chain rule, applied mechanically and without forgetting anything. But the "mechanically" part is where all the engineering lives.

An analogy: cooking with receipts

Imagine costing out a dish in a restaurant kitchen. You could look at the finished plate and estimate — but if you want precision, it's safer to keep a receipt at every step. With a record saying "used two onions, trimmed half of one away," you can walk backwards afterwards and answer exactly: if onions go up ten cents, how much more does this plate cost?

That is autodiff. Compute forward while recording each step — which value came from which values, through which operation. When the computation finishes, walk that record backwards, redistributing influence. The record is called a computation graph; walking it backwards is the backward pass.

Three ways to get a derivative

Why does this need a dedicated technique at all?

Symbolic differentiation hands $f(x)$ to a computer algebra system and gets back the expression for $f'(x)$. Exact — but under deep composition the expression blows up (the classic "expression swell"), and a fifty-layer network drowns you before you finish writing it down.

Numerical differentiation nudges the input and takes a difference, straight from the definition.

That says: move the input by $h$, see how far the output moves, divide by $h$. It's three lines of code and unusable for two reasons. Accuracy: shrinking $h$ improves the approximation in Eq. (1), but $f(x+h)$ and $f(x)$ become nearly identical, so the subtraction destroys your significant digits — catastrophic cancellation. That tug-of-war is the subject of Numerical Pitfalls — Cancellation, Rounding, and logsumexp. Count: with $n$ parameters you must run the model $n+1$ times; at seventy million parameters that's seventy million forward passes per gradient. Not a contest.

Automatic differentiation neither expands the formula nor takes differences. It applies the chain rule mechanically to the program as it executes.

One rule, and only one

Everything in autodiff rides on this:

If $x$ makes $y$ and $y$ makes $z$, the influence of $x$ on $z$ is the product of the influences along the way. Here $dz/dy$ is "how far $z$ moves when $y$ moves by one," and $dy/dx$ is "how far $y$ moves when $x$ moves by one." Two meshed gears: multiply the gear ratios and you get the overall ratio.

The important part is that this extends to any number of stages. And each stage's ratio is known in advance from the operation performed there — addition, multiplication, exp, matrix product. So: register a derivative rule once per operation type, and the machine handles any combination of them. That is the core idea.

Computation graphs: chopping a program into parts

Take the loss of a one-variable linear regression:

$$L = (wx + b - t)^2$$

Break it into primitives — operations you refuse to subdivide further.

u = w * x        # multiply
v = u + b        # add
e = v - t        # subtract
L = e * e        # multiply

Each line is a node; each value handoff is an edge. That's the computation graph, assembled as a side effect of computing forward. Every node knows the derivative of its own output with respect to its own inputs — its local derivative. For $u = w \cdot x$, $\partial u/\partial w = x$; for $v = u + b$, $\partial v/\partial u = 1$. It's a ratio that looks only at its immediate neighbors, so each node produces it alone.

Then the chain rule strings them together end to end: $\partial L/\partial w$ comes out as $2e \times 1 \times 1 \times x$.

Forward mode and reverse mode

When you multiply ratios from one end of the graph to the other, you get to choose which end to start from. This is the most interesting fork in the road.

Forward mode starts at the inputs. Seed with "suppose $w$ moves by one" and travel forward, computing at each node how far it moves. Carry derivatives alongside values using dual numbers ($a + b\varepsilon$ with $\varepsilon^2 = 0$) and ordinary arithmetic needs only rewriting, not restructuring — and nothing has to be stored.

Reverse mode starts at the output. Run forward to build the graph, then seed with "for $L$ to move by one" and travel backwards, distributing influence to every node.

Both give the same answer. They differ in cost. Consider a function with $n$ inputs and $m$ outputs.

• One forward-mode pass tells you how all outputs change when you move one input — so you need $n$ passes
• One reverse-mode pass tells you how all inputs must move to change one output — so you need $m$ passes

A deep learning loss has millions to billions of inputs (parameters) and exactly one output (a scalar loss). Huge $n$, $m = 1$. Reverse mode wins by a landslide. Better still, computing one gradient in reverse mode is known to cost a constant multiple of evaluating the function once — the Baur–Strassen result, often called the "cheap gradient principle." Ten times the parameters does not mean ten times the gradient cost.

The price is memory. Local derivatives on the way back need values from the forward pass, so you travel backwards while holding onto intermediate results. The GPU memory training eats is not mostly weights; it's these retained intermediate activations, which is also why raising the batch size hits OOM so fast. The escape hatch is gradient checkpointing (torch.utils.checkpoint): throw intermediates away and recompute them during the backward pass, trading compute for memory.

That said, forward mode is not useless. When inputs are few and outputs are many — sensitivity analysis, for instance — it is the better deal, and PyTorch ships it as torch.func.jvp / jacfwd.

What backward() is actually doing

With that in hand, every PyTorch behavior becomes explainable.

1. Any operation touching requires_grad=True leaves a record. The resulting tensor carries a grad_fn, an object pointing back at the operation that produced it. Run y = w * x and y.grad_fn is a MulBackward0. That's the receipt.

2. .backward() seeds the output with a 1. Since $\partial L/\partial L = 1$, that's the spark. It's exactly why backward() complains when the loss isn't a scalar — it has no idea where to put the 1.

3. It walks the graph in reverse. Each node computes "gradient arriving from upstream × my local derivative" and hands the result downstream. The visiting order is topological, so a node is processed only once every gradient flowing into it has arrived. When a tensor was used in several places, gradients from those paths are summed.

4. It accumulates into .grad on leaf tensors. At nn.Parameter and friends the result lands in .grad — accumulated, not assigned. That's why every step needs optimizer.zero_grad(). It isn't an oversight; it's what makes gradient accumulation (several small batches imitating one large one) fall out naturally.

5. It frees the graph once used. By default the buffers are discarded after the backward pass. Calling backward() twice on the same graph gets you "the graph has already been freed"; if you meant it, pass retain_graph=True.

So backward() is a debt collector reading a stack of receipts in reverse, apportioning responsibility. Backpropagation is nothing more than reverse-mode autodiff applied to the special case of a neural network — the layer-by-layer formulas are in Backpropagation from Scratch — It Is All Just the Chain Rule.

Feeling a local derivative in your hands

Those per-stage ratios matter most dramatically at the activation functions. Switch functions below and drag the input. The slope of the curve at that point is the local derivative, and it's exactly the number the backward pass multiplies your gradient by.

You just watched the slope collapse at both ends of the sigmoid. Stack layers and that small ratio gets multiplied again and again, shrinking the gradient exponentially — vanishing gradients. The same picture explains ReLU's popularity: its slope on the positive side is exactly 1, so no amount of multiplying shrinks anything.

ReLU also carries a trap that only autodiff users meet. It has a kink at $x=0$, so mathematically no derivative exists there; implementations pick one subgradient and return it (PyTorch's ReLU returns 0). Usually harmless — but differentiate torch.sqrt(0) or x.norm() at the origin and the local derivative behaves like $1/(2\sqrt{x})$, blowing up: forward values look fine while the gradient alone turns to nan. A class of bug you cannot see by inspecting forward values.

Forty lines of your own autodiff

Once the idea is clear, writing it yourself is the fastest way to own it. Scalar reverse-mode autodiff fits in this much:

class Var:
    def __init__(self, value, parents=(), local=()):
        self.value = value
        self.parents = parents     # upstream Vars
        self.local = local         # local derivative w.r.t. each parent
        self.grad = 0.0

    def __add__(self, o):          # d(a+b)/da = 1, /db = 1
        return Var(self.value + o.value, (self, o), (1.0, 1.0))

    def __mul__(self, o):          # d(ab)/da = b, /db = a
        return Var(self.value * o.value, (self, o), (o.value, self.value))

    def backward(self):
        order, seen = [], set()
        def visit(v):              # build a topological order
            if id(v) in seen: return
            seen.add(id(v))
            for p in v.parents: visit(p)
            order.append(v)
        visit(self)

        self.grad = 1.0            # dL/dL = 1 is the spark
        for v in reversed(order):  # travel from the output back
            for p, d in zip(v.parents, v.local):
                p.grad += v.grad * d   # chain rule; multiple paths sum

Let's use it — the gradient of $L = (wx + b - t)^2$ with respect to $w$:

w, x, b, t = Var(2.0), Var(3.0), Var(1.0), Var(5.0)
e = w * x + b + Var(-1.0) * t
L = e * e
L.backward()
print(w.grad)      # 2 * e * x = 2 * 2.0 * 3.0 = 12.0

PyTorch's autograd is this, extended to tensors, with a few hundred operation types registered, plus GPU kernels and memory management. The skeleton is the twenty lines above.

What the gradient was for

The gradient you worked so hard for exists to move parameters: step opposite the gradient, scaled by the learning rate. That's gradient descent. Raise the learning rate below and the updates diverge even though the gradients stay exact. Keeping that distinction — a correct derivative and a working training run are two different things — saves a lot of confusion in practice. What a gradient means in the first place is covered in Calculus for AI — The Gradient Is an Arrow Saying Which Way Is Better.

How this shows up on the job

Autodiff looks like something you delegate to the framework and forget. In practice ML engineers and researchers reach in and touch the graph several times a week.

Situation 1: cutting the graph on purpose. Wrap inference and evaluation in with torch.no_grad():. Nothing is recorded, so no intermediates are held, and memory and speed both change substantially. Even mid-training you'll use .detach() to stop gradient flowing backwards — updating a GAN's discriminator, generating pseudo-labels. Forget it and parameters you meant to freeze start moving, and training breaks quietly. The symptom is the nasty kind: loss goes down while evaluation refuses to improve.

Situation 2: buying memory with recomputation. When you OOM, wrap blocks in torch.utils.checkpoint.checkpoint to discard intermediate activations and recompute them on the way back. You pay roughly one extra forward pass and keep your batch size.

Situation 3: gradients turned to nan. Start with torch.autograd.set_detect_anomaly(True), which traces back to the operation that produced the nan (slow — debugging only). If you wrote a custom op by subclassing torch.autograd.Function, check it with torch.autograd.gradcheck against numerical differentiation. Numerical diff is far too slow for computing gradients, but it is still the standard tool for checking them.

Three pitfalls worth memorizing:

• .grad accumulates. Skip zero_grad() and gradients keep adding to those of previous steps. The loss wanders up and down and the cause stays invisible for hours.
• Only leaf tensors get a .grad. To inspect an intermediate tensor's gradient you must call retain_grad() explicitly. Most "why is it None?" reports are this.
• In-place operations corrupt the graph. x += 1 or relu_(x) can overwrite values the backward pass still needs, producing "a variable needed for gradient computation has been modified by an inplace operation." One character (_) added to save memory is usually the culprit, so remove it first when diagnosing.

Phrased as an interview question, this becomes: why does deep learning use reverse mode rather than forward mode? A complete answer has two halves — many inputs, one output, and the price paid for it is memory.

Summary

• Autodiff is neither symbolic nor numerical differentiation. It registers a local derivative per operation and chains them mechanically
• The forward pass builds the computation graph (the receipts); backward() walks it in reverse and apportions gradient
• With many inputs and one output, reverse mode is overwhelmingly the right choice — paid for in memory holding intermediate results
• no_grad / detach / zero_grad / in-place ops — nearly every real-world incident traces back to how the graph was handled

Next we'll follow this graph onto the GPU and look at where those intermediate results actually live in the memory hierarchy.