A reverse-mode autodiff engine, written twice — once for clarity, once for speed

What makes this more than a micrograd clone

A scalar autodiff engine is a well-worn exercise. Three things here are not:

• A tensor engine with hand-written broadcasting adjoints. numpy supplies arrays and BLAS for the forward pass only. Every derivative, every chain-rule product and every broadcast reduction is written by hand — and broadcasting is where hand-rolled engines break silently, because a wrong bias gradient still trains.
• Second-order derivatives (double backpropagation). The scalar engine's gradients are Python floats with no history, so they cannot be differentiated again. A separate module expresses each operation's backward pass in the same algebra as its forward pass, making gradients themselves differentiable — to arbitrary order.
• Verification that pins the training loop, not just the gradients. Gradient checks say nothing about whether the optimizer solves anything. So gradient descent through this engine is asserted to land on the exact closed-form least-squares solution computed independently by numpy.linalg.solve.

Two engines, one design

Identical workload: MLP(2,[16,16,1]), 337 parameters, 200 spiral points, hinge loss, full-batch gradient descent. Node counts obtained by walking the actual graph, not estimated. The ~5,000× speedup is almost exactly the 5,649× reduction in node count — which is the point: the scalar engine is slow for one explainable reason, namely that it allocates a Python object, a closure and a set per arithmetic operation.

The scalar engine was deliberately kept unchanged. Its minimalism is the pedagogical asset; the tensor engine exists so the project can actually do something.

The broadcasting adjoint

This is the genuinely subtle part of tensor autodiff. Broadcasting is a linear map: evaluating x + b with x.shape == (64, 10) and b.shape == (10,) first applies a map that copies b down all 64 rows. Reverse-mode autodiff propagates gradients by applying the adjoint (the transpose) of each forward linear map — and the transpose of a 0/1 replication matrix is a 0/1 summation matrix.

Getting this wrong is nasty precisely because it still trains. A bias gradient that is 64× too small, or that comes back with the output's shape and silently re-broadcasts on the next update, still produces a loss curve that goes down. Only an independent check catches it — which is why every broadcast pattern in the test suite asserts both the gradient's values against finite differences and its shape.

Second order: double backpropagation

In a micrograd-style engine the backward pass writes raw floats (self.grad += other.data * out.grad). Those are values, not graph nodes — nothing records how a gradient was computed, so there is no graph to differentiate a second time. That is the ceiling micrograd hits.

The fix is to express each operation's vector-Jacobian product in the same algebra as the forward pass, so that _vjp = lambda g: (g * other, g * self) builds nodes rather than multiplying floats. The backward pass stops being a numeric sweep and becomes a graph transformation, whose output can be fed straight back in. Nothing is order-specific, so n-th derivatives fall out of the same code.

f(x, y) = x²y + sin(x)y³
  ∂f/∂x   = 2xy + cos(x)y³
  ∂²f/∂x² = 2y − sin(x)y³      ← what grad(grad(f)) returns, to 0.00e+00

This is the machinery behind gradient penalties (WGAN-GP), MAML-style meta-learning, and physics-informed losses — anything that optimizes through a derivative.

Verification — 210 checks all passing

Every check compares an engine-computed quantity against something computed independently of the engine. Nothing here is of the form “the loss went down” or “the output matches last time.”

The three that matter most

The (p − y) identity. Cross-entropy is deliberately composed from primitive ops rather than fused with a hand-written backward, so the gradient that emerges is the product of six or seven separate local derivatives chained together. That it collapses to exactly p − y — measured at 0.00e+00 — is a real algebraic identity the engine was never told.

The training loop, not just the gradients. Linear regression is trained by gradient descent through this engine and asserted to converge to the exact normal-equation solution, reached to 1.4e-16 relative error. Three further consequences are pinned: the final loss equals the closed-form residual; the engine's gradient at β* is zero (the normal equation restated); and the loss is monotonically non-increasing at lr = 1/L with L the largest Hessian eigenvalue from numpy.linalg.eigvalsh — a theorem, so any rise is a real defect.

The PyTorch oracle actually ran. PyTorch 2.13.0+cpu was installed and all 30 oracle checks pass, many at exactly 0.00e+00, including the Hessian comparison against torch.autograd.grad(..., create_graph=True). When torch is absent they skip cleanly with an explanatory message and the suite reports 180/180.

     17 checks  scalar engine gradcheck (finite differences)   [OK]
     79 checks  tensor engine gradcheck + broadcasting adjoints   [OK]
     40 checks  higher-order derivatives (analytic Hessians)   [OK]
     26 checks  softmax / cross-entropy closed-form gradient   [OK]
     18 checks  training-loop pins (closed-form least squares)   [OK]
     30 checks  PyTorch oracle [OPTIONAL]   [OK]
------------------------------------------------------------------------
  TOTAL: 210 checks, 210 passed, 0 failed

Results

Two-arm spiral

500 points, 350 train / 150 held-out test, MLP(2 → 64 → 64 → 2) with 4,482 parameters, 900 epochs of Adam with an annealed learning rate.

The dataset was deliberately made harder than it needed to be. At the gentler settings this project originally used, the same network scores 100% on both splits and the figure teaches nothing. At 2.5 revolutions with angular noise 0.5 the arms genuinely overlap, and the result is a network that memorizes the training set (100%) while reaching 90.0% on data it has never seen. That 10-point gap is the honest part:

Handwritten digits

sklearn's bundled 8×8 load_digits: 1,797 samples, 1,347 train / 450 held-out test (stratified), MLP(64 → 64 → 32 → 10) with 6,570 parameters, 60 epochs, minibatch 64, Adam. sklearn supplies the array and the split only — the model, gradients, softmax, cross-entropy, optimizer, training loop and even the confusion matrix are this project's own code.

Optimizer comparison

Same problem, same seed, the same initial weights restored before every run, full-batch so there is no sampling noise, and no schedule. Comparing at one shared learning rate would be rigged, so each optimizer gets its own sweep and is plotted at its own best setting; each grid brackets its optimum on both sides.

Adam wins decisively: it is the only one of the three to reach the threshold at all within 400 epochs, and its final loss is 33× lower than the runner-up. This is one small full-batch problem, not a general claim about optimizers.

The computational graph

The renderer degrades gracefully: Graphviz when the native dot binary is present, otherwise a built-in matplotlib layered renderer. Both figures above were produced by the fallback, because dot is not installed on the machine that built this page. DOT source is written next to every PNG regardless.

Honest limitations

Measured, not hedged.

At tiny sizes the two are indistinguishable, because per-op dispatch is the entire cost and a thin numpy wrapper is in the same class as PyTorch's dispatcher. At realistic sizes PyTorch wins. The gap is only ~2× rather than ~100× because both bottom out in the same float64 BLAS — PyTorch's real advantages (float32 kernels, operator fusion, GPUs) are things this engine does not have at all, rather than does slowly.

Adam does not converge to machine precision at a fixed learning rate. On a convex problem with a known exact answer, Adam at lr = 0.05 stalls ~4.6e-5 from the true optimum after 6,000 steps and is still at ~2.1e-5 after 20,000 — it has stopped converging, not slowed down. Its update stays O(lr) as the gradient shrinks, so it orbits the optimum in a ball of radius ~O(lr). Annealing closes it to 1.4e-16. This is asserted in the test suite, not merely narrated.

What this engine does not have: no GPU support, no operator fusion, no kernel specialization, no graph optimization, no float32 path, no in-place operations, no gradient checkpointing, no sparse tensors, no distributed anything, and no convolution — the digits model is an MLP over raw pixels, not a CNN, which is a large part of why it stops at 97.56% rather than the ~99% a small conv net reaches on the same data.

Scope of the results. The spiral and digits numbers are single runs at one seed, not means over repeated seeds with error bars. The optimizer comparison is one problem. Timing figures come from a busy Windows desktop and carry visible run-to-run spread, which is why ranges are quoted rather than false precision.

Running it

pip install -r requirements.txt

python walkthrough.py          # regenerates every figure (~2.5 min)
python walkthrough.py --quick  # skips the slow scalar benchmark

python tests/run_all.py        # all checks, with the exact count
pytest tests/ -q               # same suites under pytest

walkthrough.py writes every quoted number to results/measurements.json, so the README and this page can be checked against what the code actually produced.