Automatic differentiation from scratch
By the end: You have written a tiny automatic differentiation engine, used it to train a neuron, and checked it against PyTorch.
- A value that remembers how it was made
- Local derivatives for each operation
- Walking the graph backwards
- Why gradients accumulate
- Checking against PyTorch
New to backpropagation? A freshly built network has random weights, so its answers are nonsense. Backpropagation is how that changes. The network makes a prediction, the prediction is compared with the right answer, and the size of the mistake is pushed backwards through the layers, working out how much each weight contributed to it. Every weight then moves a little in the direction that would have made the mistake smaller, and the whole thing repeats millions of times. Why it matters
Every deep learning framework rests on one mechanism: it records the operations you perform, then computes the derivative of the result with respect to any number that went into it. This section builds that mechanism in about a hundred lines and trains a neuron with it. It is the shortest route to understanding what loss.backward() is doing, why gradients have to be zeroed between steps, and where the memory in a training run actually goes.
The idea#
Every number produced during the forward pass remembers which numbers it came from and how. Backpropagation walks that record in reverse, multiplying local derivatives along the way, which is the chain rule. PyTorch does exactly this with tensors; here it is with single numbers, so every step is visible. The design follows the well-known micrograd project, which has a video walkthrough.
One design decision in the code below carries more weight than the rest, so it is worth naming before you read it: the reverse pass has to visit nodes in a particular order. A node's gradient is only final once every node that used it has contributed, so it cannot be processed until all of its consumers have been. Sorting the graph so that every node appears after its parents, then walking that list backwards, guarantees exactly that. This is a topological sort, and it is the only algorithm in the whole engine — everything else is arithmetic.
import math
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):
other = other if isinstance(other, Value) else Value(other)
out = Value(self.data + other.data, (self, other))
def backward():
self.grad += out.grad # d(a+b)/da = 1
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 # d(ab)/da = b
other.grad += self.data * out.grad
out._backward = backward
return out
__radd__, __rmul__ = __add__, __mul__
def tanh(self):
t = math.tanh(self.data)
out = Value(t, (self,))
def backward():
self.grad += (1 - t * t) * out.grad # d tanh(x)/dx = 1 - tanh(x)^2
out._backward = backward
return out
def backward(self):
order, seen = [], set()
def visit(v): # every node after its parents
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()Why += and not =#
A value can feed into several later operations. Its total gradient is the sum of the gradients flowing back from each use. Overwriting instead of adding is the classic bug, and the same reason PyTorch requires zeroing gradients between steps: they accumulate by design.
The smallest case that proves it: let y be x plus x. The derivative is obviously 2. With the plus sign, the add node sends a gradient of 1 back to its first parent and 1 to its second, and both parents are the same object, so x ends with 2. Replace the two plus-equals with plain assignment and x ends with 1, because the second write overwrote the first. Half the answer, no error, and a network that trains slowly for reasons nobody can find.
The other thing the graph explains is memory. Every intermediate value stays alive until the backward pass is done, because its local derivative needs the numbers from the forward pass. That is why a training step needs several times the memory of an inference step at the same batch size, why activation checkpointing works by throwing those intermediates away and recomputing them, and why keeping a tensor in a Python list across steps quietly holds the whole graph that produced it.
Check it#
x, w, b = Value(2.0), Value(-3.0), Value(1.0)
y = (x * w + b).tanh()
y.backward()
print(y.data, w.grad) # -0.9999092042625951 0.0003631664618877206
import torch
xt, wt, bt = (torch.tensor(v, requires_grad=True) for v in (2.0, -3.0, 1.0))
torch.tanh(xt * wt + bt).backward()
print(wt.grad.item()) # 0.0003630950232036412, the same to four digitsThe same idea on arrays#
The Value class is one number per node, which is clear and hopelessly slow. Real frameworks run the identical algorithm on whole arrays, and writing the backward pass of a two-layer network by hand in NumPy takes ten lines. It is worth doing once, because every line is one of the local derivatives you already know, and because the shapes are the part that actually confuses people.
import numpy as np
rng = np.random.default_rng(0)
X = rng.standard_normal((16, 4)) # 16 examples, 4 features
y = rng.standard_normal((16, 1))
params = {
"W1": rng.standard_normal((4, 8)) * 0.5,
"b1": np.zeros(8),
"W2": rng.standard_normal((8, 1)) * 0.5,
"b2": np.zeros(1),
}
def forward(p):
z1 = X @ p["W1"] + p["b1"] # [16, 8]
h = np.maximum(0, z1) # [16, 8] ReLU
out = h @ p["W2"] + p["b2"] # [16, 1]
loss = ((out - y) ** 2).mean()
return loss, (z1, h, out)
def backward(p, cache):
z1, h, out = cache
n = len(y)
d_out = 2 * (out - y) / n # [16, 1] derivative of the mean squared error
grads = {}
grads["W2"] = h.T @ d_out # [8, 16] @ [16, 1] -> [8, 1], matches W2
grads["b2"] = d_out.sum(axis=0) # sum over the batch
d_h = d_out @ p["W2"].T # [16, 1] @ [1, 8] -> [16, 8]
d_z1 = d_h * (z1 > 0) # ReLU: pass through where the unit was active
grads["W1"] = X.T @ d_z1 # [4, 16] @ [16, 8] -> [4, 8], matches W1
grads["b1"] = d_z1.sum(axis=0)
return grads
loss, cache = forward(params)
grads = backward(params, cache)
print(round(loss, 6)) # 1.5287Now the part that makes this trustworthy rather than plausible. A numerical gradient nudges one parameter, recomputes the loss, and divides — no calculus, no chance of repeating your own mistake. Comparing it against the analytic gradient for every parameter is the gradient check, and it is how anyone writing a new layer knows the backward pass is right.
def numerical(p, name, i, h=1e-5):
flat = p[name].reshape(-1) # a view, so writing to it edits the parameter
original = flat[i]
flat[i] = original + h
up = forward(p)[0]
flat[i] = original - h
down = forward(p)[0]
flat[i] = original # always put it back
return (up - down) / (2 * h)
def check(p, grads):
worst = 0.0
for name in p:
analytic = grads[name].reshape(-1)
for i in range(len(analytic)):
num = numerical(p, name, i)
scale = max(1e-8, abs(num) + abs(analytic[i]))
worst = max(worst, abs(num - analytic[i]) / scale)
return worst
print(check(params, grads)) # 4.5294809909942426e-10Then break it deliberately, because a check you have never seen fail is not a check you trust. Delete the ReLU mask from the backward pass — one of the most common real mistakes, since the forward pass still looks right — and run the same comparison.
def backward_buggy(p, cache):
z1, h, out = cache
n = len(y)
d_out = 2 * (out - y) / n
g = {}
g["W2"] = h.T @ d_out
g["b2"] = d_out.sum(axis=0)
d_h = d_out @ p["W2"].T
d_z1 = d_h # BUG: the ReLU mask is missing
g["W1"] = X.T @ d_z1
g["b1"] = d_z1.sum(axis=0)
return g
print(check(params, backward_buggy(params, cache))) # 1.0Finally, use the gradients for what they are for. Fifty lines of NumPy, no framework, and the loss falls from 1.5287 to 0.0739 in 500 steps.
for step in range(500):
loss, cache = forward(params)
grads = backward(params, cache)
for name in params:
params[name] -= 0.05 * grads[name]
print(round(loss, 6)) # 0.073931When to reach for a gradient check#
- You wrote a custom loss or a custom layer. Check it before training anything, on a tiny input.
- A model trains but plateaus higher than it should. A wrong gradient often still reduces the loss, just not to where it could go.
- You changed a reduction, a mask or an axis argument. These are the edits that break a backward pass silently.
- Never in a training loop. It costs two forward passes per parameter, which is why it is run once on a four by eight network rather than on the real one.
- Use double precision while checking. In 32-bit arithmetic the numerical estimate itself is only good to about five digits, so a correct gradient can look like a failed check.
Where people go wrong
Treating backpropagation as a mysterious algorithm. It is the chain rule applied in a sensible order, and writing it once removes the mystery.
If this comes up in an interview
You have written a custom layer. How do you know its backward pass is right?
A gradient check. Compute the analytic gradient, then for each parameter nudge it by about 1e-5 in both directions, recompute the loss, and divide the difference by twice the step. Compare the two as a relative error rather than an absolute difference: below 1e-7 is correct, above 1e-4 is a bug. Run it on the smallest input that exercises the layer, because it costs two forward passes per parameter.
Then they push further
Why not use a much smaller step for the numerical gradient?
Below about 1e-8 the two loss values are so close that floating point subtraction throws away the digits carrying the answer. At 1e-15 the estimate is wrong in the first decimal place. There is a window, and 1e-5 sits in it.
What kinds of bug does this actually catch?
A missing activation mask, a reduction over the wrong axis, a sign error, and a factor left out. In my experience the mask is the most common, because the forward pass still looks correct and the model still trains, just worse than it should.
What gets marked up. Marked up for relative rather than absolute comparison, for knowing the step size has a lower limit, and for saying it runs once on a tiny model rather than in the training loop.
How does reverse-mode automatic differentiation work?
During the forward pass each operation records its inputs and how to compute its local derivative. Backward starts with a gradient of 1 at the output and visits nodes in reverse topological order, each passing its gradient times its local derivative to its inputs, summing contributions where a value was used more than once.
Then they push further
Why is reverse mode used for training rather than forward mode?
Training has one scalar output, the loss, and millions of inputs, the parameters. Reverse mode gets all parameter gradients in one backward pass; forward mode would need one pass per parameter.
Why must gradients be zeroed between training steps in PyTorch?
Gradients accumulate by design, because a value used in several places needs the sum of its gradients. Without zeroing, each step would add to the previous step's gradients.
What gets marked up. Marked up for topological order, summing gradients, and the reverse versus forward mode argument.
Build this: Extend the Value class with subtraction, division, powers and exp, build a two-layer network from Value neurons, train it on a tiny dataset, and check every gradient against PyTorch.
Watch it explained
This is one section of Backpropagation, 17 hours across 4 sections. Read it in full to get the whole topic in order.