Autograd from scratch, and how small ε should be

Why start here

This is the first module of llm-lab, where I'm building the parts of an LLM from scratch with NumPy and nothing else: a tokenizer, attention, a small GPT, and later the training and serving tricks. All of it needs gradients, so gradients came first.

How it works

Every operation knows only its own local derivative. Multiplication knows that the derivative of x * y with respect to x is y, and nothing about the rest of the model. The forward pass records which values produced which, and backward() walks that graph in reverse, multiplying the incoming gradient by each local derivative and adding the result into the inputs. It adds rather than assigns because a value used twice has to receive both contributions.

That part is short. Most of the real work was shapes. If a bias of shape (M,) is added to a batch of shape (N, M), it was used N times, so its gradient has to be summed back down to (M,). Every op that can broadcast does that in its backward.

Matmul is where the shapes do the thinking for you:

def _backward():
    self.grad += unbroadcast(out.grad @ np.swapaxes(other.data, -1, -2), self.shape)
    other.grad += unbroadcast(np.swapaxes(self.data, -1, -2) @ out.grad, other.shape)

I used swapaxes instead of .T so the same code works on batched inputs, which attention will need later.

Softmax and cross-entropy are built from these primitives, so their gradients come for free. Softmax subtracts the row max first, because exp(1000) overflows to inf, and inf / inf is nan. The shift doesn't change the result at all, so no gradient has to flow through the max.

Proving it's right

The engine was the easy half. The tests are what make it trustworthy: 29 of them, every op checked against finite differences (relative error between 3e-11 and 1.4e-9), cross-entropy compared with its closed-form gradient, a two-layer MLP checked end to end, and backward() run through a chain of 5,000 operations. The first version crashed at 1,000 because it walked the graph recursively and hit Python's recursion limit, so it now uses an explicit stack.

Two details matter more than the count:

• The checker weights the output with random numbers before summing it. With a plain sum(), every output gets a gradient of 1, and a gradient that forgot to transpose a square matrix still passes.
• One test hands the checker a deliberately wrong gradient (1.9x instead of 2x) and expects it to fail. A checker that can't fail doesn't prove anything.

How small should ε be?

To check a gradient numerically you nudge each input by a small step ε and see how much the loss moves. I assumed a smaller ε was always more accurate. It isn't. Two errors pull against each other: the approximation error shrinks as ε shrinks, but the rounding error from subtracting two nearly equal numbers grows. So the error is U-shaped, and where the bottom sits depends on the precision.

On a small MLP with 67 parameters, using central differences:

The float32 column is the one that surprised me. At ε = 1e-6, a value you'll see as the default in a lot of code, a correct gradient comes out 25% wrong. Even at its best ε, float32's noise hides any bug of 1e-4 or smaller.

So every gradient check in the rest of the lab runs in float64, with central differences and ε between 1e-6 and 1e-5. A correct gradient lands around 1e-10, and anything above 1e-6 is treated as a bug until proven otherwise.