Math for CS and ML · 30 min

Backpropagation: the chain rule on a graph

Watch gradients flow back through a computation graph, then write the engine yourself, micrograd style.

01Why it matters

How does a network know which weight to blame?

A neural network can have billions of weights. After it gets an answer wrong, training nudges every one of them in the direction that would have made the error smaller. That needs the gradient of the loss with respect to every weight, and computing them one at a time would take forever.

Backpropagation gets all of them in one backward sweep, for about the cost of running the network once. It is nothing more than the chain rule from calculus, applied node by node on a graph. In this lesson you will watch it, then write it.

02Key parts

The pieces you will see

computation graph
A formula broken into small steps (add, multiply, tanh), each a box fed by its inputs.
forward pass
Computing every box's value from left to right.
local derivative
How much a box's output changes per unit change in one input, ignoring everything else: 1 for an addition, the other factor for a product, 1−tanh⁡21 - \tanh^2 for tanh.
gradient
How much the final output changes per unit change in a box: what training needs.

03See it work

Forward, then backward

Values flow left to right, then gradients flow right to left. Each step multiplies a local derivative by the gradient arriving from above.

Each box only needs to know its own local derivative. The chain rule glues them together:

∂L∂a=∂L∂d⋅∂d∂e⋅∂e∂a\frac{\partial L}{\partial a} = \frac{\partial L}{\partial d}\cdot\frac{\partial d}{\partial e}\cdot\frac{\partial e}{\partial a}

04Predict

Your turn: a neuron

A single neuron computes tanh⁡(x1w1+x2w2+b)\tanh(x_1 w_1 + x_2 w_2 + b). Predict each value going forward and each gradient coming back.

Make 6 predictions: step forward and answer each question before the next step plays.

0 of 6 made

05Try it

Compute the gradients yourself

First by hand: write the three gradients of a neuron using the chain rule.

Backpropagate through a neuron by hand

A neuron computes z = w1·x1 + w2·x2 + b and outputs out = tanh(z). Write neuron_grads(x, w, b, out, grad) (neuronGrads in JavaScript and TypeScript), where x and w are lists of two numbers and out is the output, already worked out for you.

Fill grad with how out changes with each parameter: grad[0] = d out/d w1, grad[1] = d out/d w2, grad[2] = d out/d b. Return grad.

Chain two local gradients: tanh’s, which is 1 − out², and z’s with respect to each parameter, which is the input it multiplies (or 1 for the bias).

def neuron_grads(x, w, b, out, grad):
    # out = tanh(w[0] * x[0] + w[1] * x[1] + b), already worked out
    # Fill grad with d out/d w1, d out/d w2 and d out/d b.
    return grad

Then build the engine: a Value that remembers its inputs, and the backward step for each operation.

Give each operation its local gradient

Below is the core of an autograd engine like micrograd. A Value remembers the values it was made from, and out._backward must pass out.grad back to them, each scaled by its local gradient.

Fill in the three _backward bodies: + (both inputs get out.grad), × (each input gets the other input’s data × out.grad) and tanh (its input gets (1 − t²) × out.grad). Always add to grad with +=: a value used twice collects a gradient from both uses.

grads_of(expr, vals, grad) (gradsOf in JavaScript and TypeScript) is written for you: it builds an expression from vals, runs backward, and copies each leaf’s gradient into grad. For neuron it is tanh(x1·w1 + x2·w2 + b); for reuse it is a·b + a·a (a is used three times); for chain it is tanh applied to out × w ten times over, starting from x.

import math


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

    def __add__(self, other):
        out = Value(self.data + other.data, (self, other))

        def _backward():
            pass  # TODO

        out._backward = _backward
        return out

    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other))

        def _backward():
            pass  # TODO

        out._backward = _backward
        return out

    def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,))

        def _backward():
            pass  # TODO

        out._backward = _backward
        return out


def backward(out):
    order = []
    seen = set()

    def visit(v):
        if v not in seen:
            seen.add(v)
            for child in v._prev:
                visit(child)
            order.append(v)  # after everything it was made from

    visit(out)
    out.grad = 1.0
    for v in reversed(order):
        v._backward()


def build(expr, vals):
    leaves = [Value(v) for v in vals]
    if expr == 'neuron':
        x1, w1, x2, w2, b = leaves
        out = (x1 * w1 + x2 * w2 + b).tanh()
    elif expr == 'reuse':
        a, b = leaves
        out = a * b + a * a
    else:
        x, w = leaves
        out = x
        for _ in range(10):
            out = (out * w).tanh()
    return leaves, out


def grads_of(expr, vals, grad):
    leaves, out = build(expr, vals)
    backward(out)
    for i in range(len(leaves)):
        grad[i] = round(leaves[i].grad, 9)  # same digits in every language
    return grad

Finally, the loop that makes it automatic: sort the graph so every value comes after its inputs, then walk it backwards.

Run backpropagation in the right order

Every Value below already knows its local gradients: calling v._backward() passes v.grad on to the values v was made from (in v._prev). What is missing is backward(out), which runs them all.

Set out.grad to 1 (out changes by 1 per unit of itself). Then call _backward() on every value in the graph, starting at out and working back to the leaves, so a value has received all of its gradient before it passes it on.

A topological order does it: visit a value’s inputs before the value itself (a depth-first search over _prev, remembering what you’ve seen), then go through that list backwards. For neuron it is tanh(x1·w1 + x2·w2 + b); for reuse it is a·b + a·a (a is used three times); for chain it is tanh applied to out × w ten times over, starting from x.

import math


class Value:
    def __init__(self, data, children=()):
        self.data = data
        self.grad = 0.0
        self._prev = children
        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 tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,))

        def _backward():
            self.grad += (1 - t * t) * out.grad

        out._backward = _backward
        return out


def backward(out):
    # Set out.grad to 1, then call _backward() on every value,
    # starting from out and working back towards the leaves.
    pass


def build(expr, vals):
    leaves = [Value(v) for v in vals]
    if expr == 'neuron':
        x1, w1, x2, w2, b = leaves
        out = (x1 * w1 + x2 * w2 + b).tanh()
    elif expr == 'reuse':
        a, b = leaves
        out = a * b + a * a
    else:
        x, w = leaves
        out = x
        for _ in range(10):
            out = (out * w).tanh()
    return leaves, out


def grads_of(expr, vals, grad):
    leaves, out = build(expr, vals)
    backward(out)
    for i in range(len(leaves)):
        grad[i] = round(leaves[i].grad, 9)  # same digits in every language
    return grad

06Check yourself

Check yourself

  1. Question 1 of 5

    e = a × b with a = 2, b = −3, and the gradient arriving at e is −2. What is the gradient of a?

    Your answer comes back in Review in a few days, so it sticks.

  2. Question 2 of 5

    What does an addition node do to the gradient flowing back through it?

  3. Question 3 of 5

    tanh(n) = 0.7 and the gradient arriving at its output is 1. What is the gradient of n?

  4. Question 4 of 5

    Why does backpropagation visit nodes in reverse topological order?

  5. Question 5 of 5

    a feeds two boxes. The gradient it gets back through the first is 3 and through the second is 5. What is the gradient of a?

07Explain it back

Say it in your own words

How does backpropagation find the gradient of every input in one pass, and why does it walk the graph in reverse order?

08Recap

What to remember

  • A computation graph splits a formula into simple steps, each with an easy local derivative.
  • Backward pass: output gradient 1, then each input's gradient += local derivative × upstream gradient, in reverse topological order.
  • One backward pass gives the gradient for every input; this is how PyTorch and JAX train networks.
  • Next: use these gradients for gradient descent on a real model, starting with linear regression.
Back to path