diff --git a/README.md b/README.md index 9b6ca20..62089cf 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index e616e09..7dee9a6 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/autograd/__init__.py b/src/autograd/__init__.py index c154139..4a3122c 100644 --- a/src/autograd/__init__.py +++ b/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"] diff --git a/src/autograd/engine.py b/src/autograd/engine.py index ee522b7..4444c98 100644 --- a/src/autograd/engine.py +++ b/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 diff --git a/src/autograd/optim.py b/src/autograd/optim.py new file mode 100644 index 0000000..3dc5b8a --- /dev/null +++ b/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 diff --git a/tests/test_engine.py b/tests/test_engine.py new file mode 100644 index 0000000..cd2f399 --- /dev/null +++ b/tests/test_engine.py @@ -0,0 +1,149 @@ +import math + +import numpy as np +import pytest + +from src.autograd import Value, value_sum + + +def numeric_grad(f, x, eps=1e-6): + return (f(x + eps) - f(x - eps)) / (2 * eps) + + +def test_add_and_mul_grads(): + a = Value(2.0) + b = Value(-3.0) + c = a * b + b + c.backward() + assert c.data == pytest.approx(-9.0) + assert a.grad == pytest.approx(-3.0) + assert b.grad == pytest.approx(a.data + 1.0) + + +def test_shared_node_accumulates(): + # a used twice: gradients must add, not overwrite + a = Value(3.0) + b = a + a + b.backward() + assert b.data == pytest.approx(6.0) + assert a.grad == pytest.approx(2.0) + + +def test_pow_grad_matches_numeric(): + def f(x): + return x**3 + + x0 = 1.7 + x = Value(x0) + y = x**3 + y.backward() + assert x.grad == pytest.approx(numeric_grad(f, x0), rel=1e-4) + + +def test_division(): + a = Value(6.0) + b = Value(3.0) + c = a / b + c.backward() + assert c.data == pytest.approx(2.0) + assert a.grad == pytest.approx(1.0 / 3.0) + assert b.grad == pytest.approx(-6.0 / 9.0) + + +def test_rsub_and_rtruediv(): + x = Value(4.0) + y = 10.0 - x + z = 8.0 / x + y.backward() + assert y.data == pytest.approx(6.0) + assert x.grad == pytest.approx(-1.0) + x.grad = 0.0 + z.backward() + assert z.data == pytest.approx(2.0) + assert x.grad == pytest.approx(-8.0 / 16.0) + + +@pytest.mark.parametrize("fn_name", ["relu", "tanh", "exp", "sigmoid"]) +def test_activation_grads_match_numeric(fn_name): + def forward(x): + v = Value(x) + return getattr(v, fn_name)() + + for x0 in (-1.3, 0.4, 2.1): + v = Value(x0) + out = getattr(v, fn_name)() + out.backward() + expected = numeric_grad(lambda x: forward(x).data, x0) + assert v.grad == pytest.approx(expected, rel=1e-4, abs=1e-6) + + +def test_relu_at_zero_is_dead(): + x = Value(0.0) + y = x.relu() + y.backward() + assert y.data == 0.0 + assert x.grad == 0.0 + + +def test_log_grad_and_domain(): + x = Value(2.5) + y = x.log() + y.backward() + assert y.data == pytest.approx(math.log(2.5)) + assert x.grad == pytest.approx(1.0 / 2.5) + with pytest.raises(ValueError): + Value(0.0).log() + with pytest.raises(ValueError): + Value(-1.0).log() + + +def test_pow_rejects_value_exponent(): + with pytest.raises(TypeError): + Value(2.0) ** Value(3.0) + + +def test_softmax_cross_entropy_matches_numpy(): + # a real composite: softmax over 3 logits then negative-log-likelihood. + # gradient of loss wrt the correct logit should be softmax_i - 1. + logits = [Value(1.0), Value(2.0), Value(0.5)] + target = 1 + exps = [z.exp() for z in logits] + denom = value_sum(exps) + probs = [e / denom for e in exps] + loss = -probs[target].log() + loss.backward() + + raw = np.array([1.0, 2.0, 0.5]) + sm = np.exp(raw) / np.exp(raw).sum() + expected = sm.copy() + expected[target] -= 1.0 + got = np.array([z.grad for z in logits]) + np.testing.assert_allclose(got, expected, rtol=1e-6, atol=1e-8) + assert loss.data == pytest.approx(-math.log(sm[target])) + + +def test_deep_chain_gradient(): + # long product chain stresses topo ordering: d/dx x^n = n x^(n-1) + x = Value(1.1) + out = x + for _ in range(9): + out = out * x + out.backward() + assert out.data == pytest.approx(1.1**10, rel=1e-9) + assert x.grad == pytest.approx(10 * 1.1**9, rel=1e-6) + + +def test_backward_twice_needs_manual_zero(): + # gradients accumulate across backward calls unless reset + a = Value(3.0) + b = Value(4.0) + + def build(): + return a * b + + out = build() + out.backward() + first = a.grad + out2 = build() + out2.backward() + assert a.grad == pytest.approx(2 * first) diff --git a/tests/test_optim.py b/tests/test_optim.py new file mode 100644 index 0000000..30192df --- /dev/null +++ b/tests/test_optim.py @@ -0,0 +1,74 @@ +import numpy as np +import pytest + +from src.autograd import SGD, Value, minimize, value_sum + + +def test_sgd_step_updates_data(): + p = Value(5.0) + p.grad = 2.0 + opt = SGD([p], lr=0.1) + opt.step() + assert p.data == pytest.approx(5.0 - 0.1 * 2.0) + + +def test_zero_grad_resets(): + p = Value(1.0) + p.grad = 9.0 + opt = SGD([p], lr=0.1) + opt.zero_grad() + assert p.grad == 0.0 + + +def test_minimize_quadratic_finds_minimum(): + # f(x) = (x - 3)^2 has its minimum at x = 3 + x = Value(-4.0) + history = minimize(lambda: (x - 3.0) ** 2, [x], steps=200, lr=0.1) + assert x.data == pytest.approx(3.0, abs=1e-3) + assert history[-1] < history[0] + assert history[-1] == pytest.approx(0.0, abs=1e-4) + + +def test_momentum_reaches_minimum(): + x = Value(-4.0) + history = minimize( + lambda: (x - 3.0) ** 2, [x], steps=300, lr=0.02, momentum=0.9 + ) + assert x.data == pytest.approx(3.0, abs=1e-2) + assert history[-1] < 1e-3 + # momentum should reach the basin faster than plain SGD at the same lr + y = Value(-4.0) + plain = minimize(lambda: (y - 3.0) ** 2, [y], steps=300, lr=0.02) + assert history[-1] < plain[-1] + + +def test_linear_regression_recovers_weights(): + # fit y = w*x + b to points generated from w=2, b=-1 via pure autograd + rng = np.random.default_rng(0) + xs = rng.uniform(-2, 2, size=40) + ys = 2.0 * xs + -1.0 + w = Value(0.0) + b = Value(0.0) + + def loss(): + errs = [(w * float(x) + b - float(y)) ** 2 for x, y in zip(xs, ys, strict=True)] + return value_sum(errs) * (1.0 / len(xs)) + + minimize(loss, [w, b], steps=400, lr=0.05) + assert w.data == pytest.approx(2.0, abs=1e-2) + assert b.data == pytest.approx(-1.0, abs=1e-2) + + +def test_invalid_hyperparams(): + with pytest.raises(ValueError): + SGD([Value(1.0)], lr=0.0) + with pytest.raises(ValueError): + SGD([Value(1.0)], lr=0.1, momentum=1.0) + with pytest.raises(ValueError): + SGD([Value(1.0)], lr=0.1, momentum=-0.1) + + +def test_empty_params_is_noop(): + opt = SGD([], lr=0.1) + opt.zero_grad() + opt.step() # must not raise diff --git a/tests/test_smoke.py b/tests/test_smoke.py deleted file mode 100644 index e8b2bae..0000000 --- a/tests/test_smoke.py +++ /dev/null @@ -1,5 +0,0 @@ -from src.autograd import Value # noqa: F401 - - -def test_import(): - assert Value is not None