D2.6 · LESSON 6 OF 6

A small autograd engine

Apply the chain rule through a graph and implement the core of a scalar reverse-mode autodiff engine.

THE IDEA

A whole autograd engine in a few lines

Everything so far fits in one small class. Each node stores its value, a gradient that starts at 0, its parents, and a function that knows how to send its gradient to those parents. Calling backward() on the output orders the graph, sets the output gradient to 1, then runs each node's rule in reverse order.

a minimal scalar engine (Python)

class Value:
    def __init__(self, data, parents=()):
        self.data, self.grad = data, 0.0
        self._parents = parents
        self._backward = lambda: None

    def __add__(self, other):
        out = Value(self.data + other.data, (self, other))
        def backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = backward
        return out

    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other))
        def backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = backward
        return out

    def backward(self):
        order, seen = [], set()
        def visit(v):
            if v not in seen:
                seen.add(v)
                for p in v._parents:
                    visit(p)
                order.append(v)
        visit(self)
        self.grad = 1.0
        for v in reversed(order):
            v._backward()

Notice the +=. Every rule adds to the gradient of its inputs, which is exactly how values used twice collect contributions. The same += has a side effect: calling backward() again keeps adding.

clear the gradients before every backward() call

call	∂L/∂a	∂L/∂b	∂L/∂c
1	7	2	2
2	21	6	6
3	42	12	12

Without clearing, every stored gradient keeps its old value and the new contributions are added on top, including the gradients of the in-between nodes p and q. So the totals grow faster than once per call (7, 21, 42 for a). This is why a training loop needs zero_grad before each backward().

Three calls to backward() on the same graph, with and without clearing between calls.

The training loop around it

A training step needs four things in the right order: compute fresh values (forward), clear old gradients (zero_grad), compute new gradients (backward), then move the weights (step). Skip the clearing and each update includes the stale gradients of earlier passes.

loop body	w after 3 iterations	loss
correct: forward, zero_grad, backward, step	1.984	0.001
forgot zero_grad: forward, backward, step	4.224	19.785
zero_grad after backward: forward, backward, zero_grad, step	0	16

The same tiny problem trained for 3 iterations with three different loop bodies.

Predict first. What happens if old gradients are unintentionally retained between updates?

Chapter context

Backpropagation organizes repeated applications of the chain rule. The graph stores how values depend on earlier values. Reverse traversal accumulates how the final loss changes with each intermediate quantity.

TRY IT YOURSELF

Part 1. Do it hands-on

Assemble a loop body from the four calls. Build the correct one, then reproduce the two classic bugs on purpose.

model: prediction = w · 2, target y = 4, loss = (w·2 − 4)², start w = 0, learning rate 0.1, 3 iterations

Loop body:empty

iteration	gradient used	w after	loss after
1	0	0	16
2	0	0	16
3	0	0	16

To do: Build a loop body so that after 3 iterations the loss is below 0.01
To do: Now cause the classic bug on purpose: a loop that steps with backward() but never clears, ending with the loss above 10
To do: Find the loop body that gets the update but runs zero_grad AFTER backward. Does it still train? Make the final loss above 10

Part 2. Do it in code

Finish the scalar Value class. __add__ is done as a model. Write __mul__, tanh, relu and __truediv__ so each returns a new Value that remembers its parents and a _backward closure that ADDS (+=) into the parents' grads.
Write backward(): build the topological order of the graph that ends at self (each node once), set self.grad = 1.0, then call each node's _backward in reverse order. It must NOT clear any gradients.
Write zero_grad(): set grad to 0.0 on self and everything it depends on. Numbers mixed with Values (2 * v, v + 1, v - 1) are already handled by the helpers at the bottom.

import math

class Value:
    def __init__(self, data, parents=(), op=""):
        self.data = float(data)
        self.grad = 0.0
        self._parents = parents
        self._backward = lambda: None
        self.op = op

    def __repr__(self):
        return "Value(data=" + str(self.data) + ", grad=" + str(self.grad) + ")"

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

    def __mul__(self, other):
        pass

    def tanh(self):
        pass

    def relu(self):
        pass

    def __truediv__(self, other):
        pass

    def backward(self):
        pass

    def zero_grad(self):
        pass

    # given: let plain numbers and Values mix
    def __neg__(self):
        return self * -1.0
    def __sub__(self, other):
        return self + (-other)
    def __radd__(self, other):
        return self + other
    def __rmul__(self, other):
        return self * other

a, b = Value(2.0), Value(3.0)
c = a * b + a
c.backward()
print(a.grad, b.grad)   # should print 4.0 2.0

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

What happens if old gradients are unintentionally retained?

ANothing; old gradients are cleared automatically before each backward().BLater updates include contributions from earlier passes, because backward() adds onto whatever is already stored.COnly the first update is affected; later updates ignore them.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

If $u=ab+c$ and $L=u^2$ , then

\frac{\partial L}{\partial a}=2ub,\quad \frac{\partial L}{\partial b}=2ua,\quad \frac{\partial L}{\partial c}=2u.

At a shared node x, add all downstream contributions: $\partial L/\partial x=\sum_j(\partial L/\partial u_j)(\partial u_j/\partial x)$ .

Worked example

Take a=2,b=3,c=1. Then u=7 and L=49. The upstream gradient into u is 14, giving dL/da=14×3=42, dL/db=14×2=28 and dL/dc=14. These are sensitivities, not new parameter values; an optimizer still needs a step size.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import isclose
a,b,c = 2.,3.,1.
u = a*b+c
assert u*u == 49
assert (2*u*b,2*u*a,2*u) == (42,28,14)
h = 1e-5
f = lambda x: (x*b+c)**2
assert isclose((f(a+h)-f(a-h))/(2*h), 42, rel_tol=1e-7)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Forward values and gradients are different objects. Shared branches require addition of gradient contributions. Backpropagation computes derivatives; the optimizer changes parameters.