Project 44: Autograd Engine

A self-contained deep-dive from the canonical Math from Foundations to Machine Learning curriculum.

  • Difficulty: Level 4: Expert
  • Time: 2 weeks
  • Language: Python
  • Prerequisites: Projects 2 and 42-43
  • Source: NN P05

What You Will Build

Implement a small define-by-run automatic differentiation engine. A Value or tensor object stores data, accumulated gradient, parent nodes, operation metadata, and a local backward closure. Arithmetic and nonlinear operators construct a directed acyclic computation graph dynamically; backward() topologically sorts reachable nodes and executes local rules in reverse. Support repeated operands and branching, where gradients must accumulate rather than overwrite. Add graph visualization, zeroing semantics, finite-difference verification, and enough operations to train a tiny multilayer perceptron.

Real World Outcome

A user can compose an ordinary expression such as a tiny neural loss, inspect its computation graph, call one backward method, and see correct derivatives for every leaf. The engine passes numerical checks on branched expressions and trains a small classifier while logging graph size, loss, and gradients.

Core Question

How can a program record arbitrary calculations and automatically apply the chain rule backward through them?

Concepts You Must Understand First

  • Directed acyclic graphs, reachability, and topological ordering.
  • Reverse-mode automatic differentiation and local chain-rule composition — Baydin et al., Automatic Differentiation in Machine Learning: a Survey.
  • Gradient accumulation when one value contributes through multiple paths.
  • Operator overloading and closure capture — Fluent Python by Luciano Ramalho.
  • Numerical gradient checking and non-differentiable points.

Build Milestones

  1. Implement scalar nodes and overloaded addition, multiplication, power, negation, and division.
  2. Add tanh, ReLU, exponential, and logarithm with local backward rules.
  3. Implement deterministic topological traversal, reverse execution, accumulation, and gradient reset.
  4. Verify nested, branched, repeated-node, and shared-subexpression graphs against finite differences.
  5. Add neuron/layer/module abstractions and train a small binary classifier using the engine.

Hints in Layers

  1. Draw the graph for a*a + a; it exposes both repeated-use and accumulation requirements.
  2. The backward seed for a scalar output is 1 because its derivative with respect to itself is 1.
  3. Capture operands, not mutable loop variables, inside each backward closure.

Common Pitfalls and Debugging

  • A reused variable has a gradient that is too small: local backward rules overwrite .grad. Add contributions with +=.
  • Some nodes never receive gradients: traversal order is wrong or parents are omitted. Build a postorder list, then reverse it.
  • A second training step explodes: old gradients remain. Make zero_grad() explicit and test repeated backward behavior.

Definition of Done

  • Supported operators pass analytical and finite-difference unit tests.
  • Repeated operands and branched graphs accumulate every gradient contribution.
  • Topological execution is deterministic and handles shared subexpressions once.
  • Graph visualization labels values, operations, and gradients.
  • A small MLP trained through the engine reduces loss on a nonlinear dataset.
  • Gradient reset and repeated-backward semantics are documented and tested.

Previous: Project 43 · Complete learning path · Next: Project 45