Scalar reverse-mode autograd engine + SGD

README.md

@@ -4,16 +4,45 @@ Core machine-learning and deep-learning building blocks implemented from scratch
 
 ## What this demonstrates
 
-Every deep learning framework is, at its core, a graph that records operations and replays them backward to compute gradients. This repo builds that machinery by hand and then uses it to train models, so the chain rule, backpropagation, and gradient descent stay visible in a few hundred lines you can read end to end. Each concept lands as its own tested module with the intuition and the math written out.
+Every deep learning framework is, at its core, a graph that records operations and replays them backward to compute gradients. This repo builds that machinery by hand and then uses it to train models, so the chain rule, backpropagation, and gradient descent are all visible in a few hundred lines you can read end to end. Each concept lands as its own tested module with the intuition and the math written out, working up to a from-scratch attention block that learns a toy task.
 
 ## Concepts demonstrated
 
-- Reverse-mode automatic differentiation (backpropagation)
-- Gradient descent and numerical gradient checking
+- Reverse-mode automatic differentiation (backpropagation) over a dynamically built computation graph
+- The chain rule applied per-operation via local backward closures
+- Reverse topological ordering of the graph so each node's gradient is complete before it is used
+- Gradient accumulation for nodes reused multiple times in a graph (a real source of bugs)
+- Gradient descent: vanilla SGD and SGD with momentum
+- Numerical gradient checking (central finite differences) as a correctness oracle
+- Softmax with numerically-checked cross-entropy loss, including the `softmax - onehot` gradient identity
+- Linear regression trained purely through the autograd engine
 
 ## What's implemented
 
-Nothing yet. First module lands next.
+- **Gradient descent + autodiff-lite**: a scalar reverse-mode autograd engine (`Value`) with `+`, `*`, `**`, division, and `relu`/`tanh`/`exp`/`log`/`sigmoid`, plus an `SGD` optimizer (with momentum) and a `minimize` training loop. Gradients are checked against finite differences and against numpy for softmax cross-entropy.
+
+## Usage
+
+```python
+from src.autograd import Value, minimize
+
+# minimize f(x) = (x - 3)^2 with gradient descent
+x = Value(-4.0)
+minimize(lambda: (x - 3.0) ** 2, [x], steps=200, lr=0.1)
+print(x.data)  # ~3.0
+```
+
+Differentiate an arbitrary scalar expression:
+
+```python
+from src.autograd import Value
+
+a = Value(2.0)
+b = Value(-3.0)
+c = (a * b + b).tanh()
+c.backward()
+print(a.grad, b.grad)  # d c / d a, d c / d b
+```
 
 ## Running the tests
 

pyproject.toml

@@ -13,6 +13,7 @@ dev = ["pytest>=8.0", "ruff>=0.6"]
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = "-q"
+pythonpath = ["."]
 
 [tool.ruff]
 line-length = 88

src/autograd/__init__.py

@@ -1,3 +1,4 @@
-from .engine import Value
+from .engine import Value, value_sum
+from .optim import SGD, minimize
 
-__all__ = ["Value"]
+__all__ = ["Value", "value_sum", "SGD", "minimize"]

src/autograd/engine.py

@@ -1,12 +1,162 @@
-"""Scalar reverse-mode autograd. Placeholder until the engine lands."""
+"""Scalar reverse-mode automatic differentiation.
+
+A `Value` wraps a single float and records the operation that produced it,
+so calling `.backward()` walks the graph in reverse topological order and
+accumulates gradients via the chain rule. This is the mechanism every deep
+learning framework uses, stripped down to scalars so the math stays visible.
+"""
 
 from __future__ import annotations
 
+import math
+from collections.abc import Callable, Iterable
+
 
 class Value:
-    def __init__(self, data: float) -> None:
+    __slots__ = ("data", "grad", "_backward", "_prev", "_op")
+
+    def __init__(
+        self,
+        data: float,
+        _children: tuple[Value, ...] = (),
+        _op: str = "",
+    ) -> None:
         self.data = float(data)
         self.grad = 0.0
+        # local backprop closure; leaves keep the no-op
+        self._backward: Callable[[], None] = lambda: None
+        self._prev: set[Value] = set(_children)
+        self._op = _op
+
+    def __add__(self, other: Value | float) -> Value:
+        other = other if isinstance(other, Value) else Value(other)
+        out = Value(self.data + other.data, (self, other), "+")
+
+        def _backward() -> None:
+            self.grad += out.grad
+            other.grad += out.grad
+
+        out._backward = _backward
+        return out
+
+    def __mul__(self, other: Value | float) -> Value:
+        other = other if isinstance(other, Value) else Value(other)
+        out = Value(self.data * other.data, (self, other), "*")
+
+        def _backward() -> None:
+            self.grad += other.data * out.grad
+            other.grad += self.data * out.grad
+
+        out._backward = _backward
+        return out
+
+    def __pow__(self, other: int | float) -> Value:
+        if not isinstance(other, (int, float)):
+            raise TypeError("only int/float exponents are supported")
+        out = Value(self.data**other, (self,), f"**{other}")
+
+        def _backward() -> None:
+            self.grad += (other * self.data ** (other - 1)) * out.grad
+
+        out._backward = _backward
+        return out
+
+    def relu(self) -> Value:
+        out = Value(self.data if self.data > 0 else 0.0, (self,), "relu")
+
+        def _backward() -> None:
+            self.grad += (out.data > 0) * out.grad
+
+        out._backward = _backward
+        return out
+
+    def tanh(self) -> Value:
+        t = math.tanh(self.data)
+        out = Value(t, (self,), "tanh")
+
+        def _backward() -> None:
+            self.grad += (1.0 - t * t) * out.grad
+
+        out._backward = _backward
+        return out
+
+    def exp(self) -> Value:
+        e = math.exp(self.data)
+        out = Value(e, (self,), "exp")
+
+        def _backward() -> None:
+            self.grad += e * out.grad
+
+        out._backward = _backward
+        return out
+
+    def log(self) -> Value:
+        if self.data <= 0.0:
+            raise ValueError("log is only defined for positive values")
+        out = Value(math.log(self.data), (self,), "log")
+
+        def _backward() -> None:
+            self.grad += (1.0 / self.data) * out.grad
+
+        out._backward = _backward
+        return out
+
+    def sigmoid(self) -> Value:
+        s = 1.0 / (1.0 + math.exp(-self.data))
+        out = Value(s, (self,), "sigmoid")
+
+        def _backward() -> None:
+            self.grad += s * (1.0 - s) * out.grad
+
+        out._backward = _backward
+        return out
+
+    def backward(self) -> None:
+        # reverse topological order so every child's grad is final before use
+        topo: list[Value] = []
+        visited: set[Value] = set()
+
+        def build(v: Value) -> None:
+            if v in visited:
+                return
+            visited.add(v)
+            for child in v._prev:
+                build(child)
+            topo.append(v)
+
+        build(self)
+        self.grad = 1.0
+        for node in reversed(topo):
+            node._backward()
+
+    def __neg__(self) -> Value:
+        return self * -1.0
+
+    def __radd__(self, other: float) -> Value:
+        return self + other
+
+    def __sub__(self, other: Value | float) -> Value:
+        return self + (-other if isinstance(other, Value) else Value(-other))
+
+    def __rsub__(self, other: float) -> Value:
+        return Value(other) + (-self)
+
+    def __rmul__(self, other: float) -> Value:
+        return self * other
+
+    def __truediv__(self, other: Value | float) -> Value:
+        other = other if isinstance(other, Value) else Value(other)
+        return self * other**-1
+
+    def __rtruediv__(self, other: float) -> Value:
+        return Value(other) * self**-1
 
     def __repr__(self) -> str:
         return f"Value(data={self.data:.6g}, grad={self.grad:.6g})"
+
+
+def value_sum(values: Iterable[Value]) -> Value:
+    acc = Value(0.0)
+    for v in values:
+        acc = acc + v
+    return acc

src/autograd/optim.py

@@ -0,0 +1,61 @@
+"""Gradient descent over `Value` parameters.
+
+`SGD` holds references to the leaf `Value`s it optimizes and steps each one
+against its accumulated gradient. `momentum` adds a velocity term, the same
+update PyTorch's SGD uses. `zero_grad` must run before each backward pass
+because gradients accumulate rather than overwrite.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from .engine import Value
+
+
+class SGD:
+    def __init__(
+        self,
+        params: Sequence[Value],
+        lr: float = 0.01,
+        momentum: float = 0.0,
+    ) -> None:
+        if lr <= 0.0:
+            raise ValueError("lr must be positive")
+        if not 0.0 <= momentum < 1.0:
+            raise ValueError("momentum must be in [0, 1)")
+        self.params = list(params)
+        self.lr = lr
+        self.momentum = momentum
+        self._velocity = [0.0 for _ in self.params]
+
+    def zero_grad(self) -> None:
+        for p in self.params:
+            p.grad = 0.0
+
+    def step(self) -> None:
+        for i, p in enumerate(self.params):
+            if self.momentum:
+                self._velocity[i] = self.momentum * self._velocity[i] + p.grad
+                p.data -= self.lr * self._velocity[i]
+            else:
+                p.data -= self.lr * p.grad
+
+
+def minimize(
+    loss_fn,
+    params: Sequence[Value],
+    steps: int = 100,
+    lr: float = 0.01,
+    momentum: float = 0.0,
+) -> list[float]:
+    """Run `steps` of gradient descent, returning the loss at each step."""
+    opt = SGD(params, lr=lr, momentum=momentum)
+    history: list[float] = []
+    for _ in range(steps):
+        opt.zero_grad()
+        loss = loss_fn()
+        loss.backward()
+        opt.step()
+        history.append(loss.data)
+    return history