Build an autograd engine

Idea 2: gradients add up

Every _backward uses +=, never =. Here's why.

Picture a café where milk goes into both the latte and the cappuccino. If milk gets pricier, your costs rise through both drinks. The total change is the two changes added together. A Value used in several places is the same. It affects the answer along several paths, so its gradient must collect every one.

With =, each new path would overwrite the last one. You'd keep only one contribution and lose the rest. With +=, they all add up. This is the rule from the computation-graph lesson, now built into the code.

Idea 3: work backward, in the right order

A node must collect all of its gradient before it passes any of it on. Think of a team captain reporting votes. The captain must hear from every team member before sending the total up the line. Report too early, and the late votes are lost.

The backward() method makes sure this happens. It works in two steps:

1. It lists every node in topological order. That's an order where each node comes after the nodes it was made from. You met it in the computation-graph lesson.
2. It reverses that list, so the final answer comes first. Then it calls each node's _backward.

Before the sweep, it sets the output's own gradient to 1. Why 1? If the answer moves by 1, then the answer has moved by 1. So its gradient with respect to itself is always 1.

In reverse order, every node that uses a value runs before that value does. So each value's gradient is complete by the time it passes gradient further back.

Enough operations to train a network

Add a few more operations and you can build a whole network out of Values:

• power (__pow__), so you can square an error;
• tanh, the activation;
• negation and subtraction, built from what's already there. $-a$ is $a \times -1$, and $a - b$ is $a + (-b)$.

Next come two small classes you'll recognize. A Neuron holds one weight per input and a bias. It adds up weight times input, adds the bias, then applies tanh. An MLP stacks layers of neurons. The __call__ method lets you use one like a function: model(x).

MLP([3, 4, 4, 1]) means 3 inputs, two hidden layers of 4 neurons, and 1 output. The network learns to output +1 or −1 for four small examples. The loss adds up the squared misses. Each step then does four things: work out the loss, reset every gradient to 0, run backward(), and move each parameter a little downhill. The step size, 0.05, is the learning rate.

import math
import random

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

    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):
        other = other if isinstance(other, Value) else Value(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 __pow__(self, k):
        out = Value(self.data ** k, (self,))
        def _backward():
            self.grad += k * self.data ** (k - 1) * 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

    __radd__ = __add__
    __rmul__ = __mul__
    def __neg__(self): return self * -1
    def __sub__(self, other): return self + (-other)

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

class Neuron:
    def __init__(self, n_in):
        self.w = [Value(random.uniform(-1, 1)) for _ in range(n_in)]
        self.b = Value(0.0)
    def __call__(self, x):
        return sum((wi * xi for wi, xi in zip(self.w, x)), self.b).tanh()
    def parameters(self):
        return self.w + [self.b]

class MLP:
    def __init__(self, sizes):
        self.layers = [[Neuron(a) for _ in range(b)] for a, b in zip(sizes, sizes[1:])]
    def __call__(self, x):
        for layer in self.layers:
            x = [n(x) for n in layer]
        return x[0] if len(x) == 1 else x
    def parameters(self):
        return [p for layer in self.layers for n in layer for p in n.parameters()]

random.seed(1)
model = MLP([3, 4, 4, 1])
xs = [[2.0, 3.0, -1.0], [3.0, -1.0, 0.5], [0.5, 1.0, 1.0], [1.0, 1.0, -1.0]]
ys = [1.0, -1.0, -1.0, 1.0]
print("parameters:", len(model.parameters()))

for step in range(40):
    preds = [model(x) for x in xs]
    loss = sum((p - y) ** 2 for p, y in zip(preds, ys))
    for p in model.parameters():
        p.grad = 0.0                         # gradients accumulate, so reset them first
    loss.backward()
    for p in model.parameters():
        p.data -= 0.05 * p.grad              # step downhill
    if step % 10 == 0 or step == 39:
        print(f"step {step:>2}  loss {loss.data:.4f}  predictions {[round(p.data, 2) for p in preds]}")

random.seed(1) fixes the starting weights, so you get the same run every time.

First it prints parameters: 41. Count them layer by layer. Each neuron has one weight per input, plus a bias. The first layer has 4 neurons with 3 inputs each: $4 \times (3 + 1) = 16$. The second has 4 neurons with 4 inputs: $4 \times (4 + 1) = 20$. The output neuron has 4 inputs: $1 \times (4 + 1) = 5$. That's $16 + 20 + 5 = 41$.

Then watch the loss. It starts at 4.5167 and falls to 0.0202 by step 39. The predictions start near random, at 0.03, −0.71, 0.41 and −0.23. They end at 0.94, −0.92, −0.93 and 0.93, close to the targets of +1, −1, −1 and +1. Every gradient came from your own engine.

Watch out: the reset to zero before each backward() matters. Gradients add up with +=. Skip the reset, and this step's gradients pile on top of the last step's. PyTorch behaves the same way. That's why every PyTorch training loop calls zero_grad().

a, b = 3.0, -2.0

# forward
p = a * b        # multiply: -6.0
f = p + a        # add: -3.0

# backward, output first
f_grad = 1.0
p_grad = f_grad * 1          # add passes gradient straight through
a_grad = p_grad * b          # multiply: the other input
b_grad = p_grad * a
a_grad += f_grad * 1         # a is used again in the add, so add its share
print(f, a_grad, b_grad)

From scalars to tensors

A scalar is a single number. This engine handles one scalar per Value. A network with millions of weights would build millions of tiny nodes, which is far too slow.

PyTorch works with tensors, its name for whole arrays of numbers. It makes one node per array operation, such as a whole matrix multiply. The backward rules from the backpropagation lesson are written once, in fast code. Everything else is the same as your engine: record the graph as the code runs, give each operation a local backward function, add gradients with +=, and sweep from the output back.