Automatic Differentiation Engine
Build a micrograd-style reverse-mode autodiff engine and train a neural network in pure Python
What You Will Build
Modern deep learning frameworks like PyTorch and JAX are built around one core abstraction: Reverse-Mode Automatic Differentiation (`autograd`). Whenever you execute a forward calculation, the engine dynamically constructs a Directed Acyclic Graph (DAG) of operations and closures. Calling `.backward()` on the final scalar loss traverses that graph in reverse topological order, applying the multivariable chain rule to populate `.grad` on every weight in the network. Inspired by Andrej Karpathy's `micrograd`, you will build a complete scalar autodiff engine in pure Python and use it to train a multilayer perceptron classifier.
Why This Project Matters in the AI/LLM Stack
Once you build `Value` and `.backward()` in ~150 lines of pure Python, backpropagation ceases to be magic. You will understand why gradients accumulate with `+=` when a node is reused across multiple paths, why `.zero_grad()` is required every step, and how PyTorch's `autograd` engine works under the hood.
Learning Objectives
- Implement a scalar `Value` node that wraps a float, tracks parent operands, and stores its local `.grad` and `_backward` closure.
- Overload arithmetic operators (`+`, `-`, `*`, `/`, `**`) and nonlinear activations (`tanh`, `relu`, `exp`) with exact local chain-rule updates.
- Implement DFS topological sorting and reverse-mode `.backward()` with correct `+=` gradient accumulation for multi-use nodes.
- Build `Neuron`, `Layer`, and `MLP` classes on top of `Value`, verify gradients numerically, and train a neural network to minimize loss.
Prerequisites
- Phase 1 Multivariable Calculus (derivatives, partial derivatives, chain rule)
- Object-oriented Python (operator overloading: `__add__`, `__mul__`, `__pow__`, `__radd__`, `__neg__`)
Required Mathematical Concepts
- Directed Acyclic Graphs (DAGs) of elementary operations
- Single and multivariable chain rule: ∂L/∂x = ∑_i (∂L/∂z_i)(∂z_i/∂x)
- Local elementary derivatives: d(a+b)/da = 1, d(a·b)/da = b, d(aⁿ)/da = n·aⁿ⁻¹
- Nonlinear activation derivatives: d(tanh x)/dx = 1 - tanh²(x), d(ReLU x)/dx = 𝟙(x > 0)
- Reverse topological ordering of graph nodes
- Central finite-difference gradient verification
Step-by-Step Implementation Tasks (15 Steps)
Check off each engineering task as your code passes its acceptance criteria. Progress is saved automatically in your browser.
Scalar Value Node & Elementary Arithmetic Operations
Create the `Value` class, track parent pointers, and overload addition, multiplication, subtraction, division, and power operations.
Every scalar number in the computational graph is wrapped in a `Value` object that tracks both its forward value and its backward sensitivity.
Expressions often mix `Value` objects with plain Python `int` or `float` literals (e.g., `x + 2.0` or `3.0 * x`).
To traverse the graph backward from the output loss to the leaf inputs, each output node must remember which child `Value` nodes produced it.
An addition gate has local derivative `1.0` for both inputs, so it routes the upstream gradient `out.grad` unchanged to both parents.
By the product rule, for `out = a * b`, `∂out/∂a = b` and `∂out/∂b = a`. A multiplication gate swaps and scales gradients by the opposite operand's value.
Supporting `x ** n` (for `int` or `float` exponent `n`) enables division `a / b = a * (b ** -1)` and squared error loss `(y_pred - y) ** 2`.
Nonlinear Activations & Computational Graph Construction
Implement `tanh`, `relu`, and `exp` operations and inspect multi-step computational graphs.
Without nonlinear activation functions, a stack of linear layers collapses into a single linear transformation.
Composing multiple operations builds a multi-layer expression graph ready for reverse traversal.
Topological Traversal, Backpropagation & Numerical Validation
Implement topological sort, reverse-mode `.backward()`, `+=` gradient accumulation, and numerical gradient verification.
Before computing `node._backward()`, every downstream node that depends on `node` must have already finished accumulating its gradient into `node.grad`.
Calling `loss.backward()` seeds `loss.grad = 1.0` (`∂loss/∂loss = 1`) and invokes `_backward()` on every node in reverse topological order.
When a variable `a` is used more than once (e.g., `b = a + a` or `c = a * a + a`), the multivariable chain rule requires summing the gradients from all outgoing branches.
Comparing `Value.grad` against finite differences `(f(x+h) - f(x-h)) / (2h)` (and optionally PyTorch `torch.tensor(..., requires_grad=True)`) proves your autodiff engine is exact.
Multilayer Perceptron, Training Loop & Testing
Construct `Neuron`, `Layer`, and `MLP` modules, train on a dataset, and write unit tests.
With `Value` complete, a neural network is simply a hierarchy of `Module`, `Neuron`, `Layer`, and `MLP` classes that hold `Value` parameters.
Training the MLP requires a loop of: forward pass -> scalar loss -> `zero_grad()` -> `loss.backward()` -> parameter update `p.data -= lr * p.grad`.
Documenting how computational graphs and reverse-mode autodiff work makes your repository a reference-grade engineering artifact.
Starter & Verification Code Examples
Topological Sort & Reverse-Mode Backpropagation on `Value`
Illustrates how `backward()` orders the computational DAG topologically and executes local chain-rule closures.
class Value:
def __init__(self, data: float, _children: tuple["Value", ...] = (), _op: str = "") -> None:
self.data = float(data)
self.grad = 0.0
self._backward = lambda: None
self._prev = set(_children)
self._op = _op
def backward(self) -> None:
topo: list[Value] = []
visited: set[Value] = set()
def build_topo(v: Value) -> None:
if v not in visited:
visited.add(v)
for child in v._prev:
build_topo(child)
topo.append(v)
build_topo(self)
self.grad = 1.0
for node in reversed(topo):
node._backward()How to Test Your Implementation
Verify multi-path gradient accumulation, finite-difference agreement across all operators, and MLP training convergence.
Diamond Graph Gradient Accumulation (`+=` Test)
Catch the classic bug where `_backward` overwrites `self.grad = ...` instead of accumulating `self.grad += ...`.
Steps: Compute x = Value(3.0); y = (x * x) + (x + x); y.backward().
Pass Criteria: y.data == 15.0 and x.grad == 8.0 (since dy/dx = 2x + 2 = 8 at x=3).
Central Finite-Difference Verification on Compound DAG
Confirm chain rule accuracy through composed division, power, tanh, and ReLU gates.
Steps: Evaluate f(a, b) = ((a * b + b**2).tanh() / (a - b)).relu() and compare a.grad, b.grad with finite differences (eps=1e-6).
Pass Criteria: Absolute difference < 1e-6 for all leaf inputs.
MLP Convergence & Zero-Grad Check
Verify that an MLP reduces loss over 40 steps of gradient descent.
Steps: Run 40 steps of forward, zero_grad(), backward(), and step(lr=0.05) on a 4-sample dataset.
Pass Criteria: Final loss < 0.05 and all 4 binary targets classified accurately.
Overall Project Acceptance Criteria
- `Value` class supports `+`, `-`, `*`, `/`, `**`, `tanh`, `relu`, and `exp` with parent tracking.
- Topological sort and reverse-mode `.backward()` accurately propagate gradients through arbitrary DAGs.
- Multi-branch gradient accumulation (`+=`) verified on reused variables.
- `Neuron`, `Layer`, and `MLP` classes built and trained to convergence on a small dataset.
Expandable Hints & Debugging Tips
Try implementing each stage on your own first. Expand a hint below if you want to verify tensor shapes, mathematical identities, or common edge cases.
Optional Stretch Challenges
Computational Graph Visualizer (DOT / SVG)
Write a utility that exports any `Value` expression graph into Graphviz DOT format showing `.data`, `.grad`, and `_op` gates at every node.
What You Learn: How to visually inspect forward activations and backward gradient flow through a network.
Cross-Entropy Loss & Softmax in Scalar Autodiff
Using `.exp()`, `.log()`, and summation over `Value` logits, implement multi-class Softmax Cross-Entropy loss and verify its gradients against `probs - one_hot`.
What You Learn: How composite losses automatically differentiate through elementary `exp` and `log` primitives.
What Your Completed Repository Should Contain
A working autodiff engine capable of training a small neural network.
Deliverable Artifacts
- Pure-Python `Value` reverse-mode autodiff engine (`autograd/engine.py`)
- `Neuron`, `Layer`, and `MLP` neural network modules (`autograd/nn.py`)
- Numerical gradient verification and MLP training test suite
After Completing This Project, You Should Be Able To...
- Explain how dynamic computational graphs record forward operations and parent dependencies.
- Implement reverse-mode automatic differentiation using DFS topological sorting and chain-rule closures.
- Handle multi-path gradient accumulation (`+=`) and explain why `.zero_grad()` is required between optimization steps.
- Build and train a multilayer perceptron from scratch on top of scalar autodiff primitives.
GitHub Repository Submission
Link your public or private GitHub repository URL to keep a structured portfolio of your from-scratch AI builds. Stored locally in your browser (self-reported; no GitHub tokens or remote code execution).
Suggested Repository: micrograd-autodiff-engine
Recommended Directory Layout:
autograd/engine.py autograd/nn.py tests/test_autograd.py README.md
Recommended README Documentation Checklist
- Diagram or explanation of forward DAG construction and reverse topological pass
- Instructions to run finite-difference gradient tests
- Example training log of the pure-Python MLP