From 5eb49bf73f913941143eb01987bb092e374e724b Mon Sep 17 00:00:00 2001 From: Thomas Hart Date: Mon, 27 Jul 2026 21:27:58 +0000 Subject: [PATCH] feat: Add SGD, Momentum, and Adam with same-net comparison --- README.md | 24 ++++++ src/mlp.py | 32 +++---- src/optimizers.py | 178 +++++++++++++++++++++++++++++++++++++++ tests/test_optimizers.py | 140 ++++++++++++++++++++++++++++++ 4 files changed, 356 insertions(+), 18 deletions(-) create mode 100644 src/optimizers.py create mode 100644 tests/test_optimizers.py diff --git a/README.md b/README.md index 3f7e97b..90dfc41 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,11 @@ Every deep learning framework is, at its core, a graph that records operations a - He/Kaiming initialization: variance `2/fan_in` restoring the half of the signal ReLU drops - Forward variance profiling across depth: naive `N(0,1)` weights explode; correct scales stay `O(1)` - Vanishing activations under mismatched init (Xavier on a deep ReLU stack) measured, not just described +- First-order optimizers as pure parameter updates decoupled from the loss and backprop +- Vanilla SGD: θ ← θ − ηg +- Heavy-ball / Polyak momentum: velocity accumulation that carries steps through elongated valleys +- Adam (adaptive moment estimation): bias-corrected first and second moments, per-parameter step sizes +- Fair convergence comparison: same MLP init, same minibatches, only the update rule changes ## What's implemented @@ -43,6 +48,7 @@ Every deep learning framework is, at its core, a graph that records operations a - **Logistic regression + cross-entropy, decision boundary demo**: `src/logreg.py` trains a binary classifier by minibatch SGD, using the fact that the gradient of the mean cross-entropy with respect to the logits is exactly `sigmoid(z) - y`, the same residual form linear regression has. The sigmoid branches on the sign of the logit so `exp` never overflows, and the loss is computed as `softplus(z) - y·z` through `numpy.logaddexp` so a confidently-wrong prediction gives a large finite loss instead of `inf`. `decision_boundary` returns the straight line where the model sits at 50 percent for a two-feature problem, the level set `w·x + b = 0`. - **Multilayer perceptron with hand-derived backprop**: `src/mlp.py` stacks affine layers with `tanh` or `relu` and a softmax head, and trains a multiclass classifier by minibatch SGD. The backward pass is written out by hand as one recursion on the per-layer delta rather than delegated to an autodiff engine: the output delta is the `softmax - onehot` residual, each hidden delta is `(delta_next @ W_nextᵀ) ⊙ act'(z)`, and the parameter gradients are `dW = a_prevᵀ @ delta` and `db = Σ delta`. Weights use He init for `relu` and Xavier for `tanh` so the signal variance holds across depth. The gradients are verified against central finite differences to a tight tolerance, and the model learns XOR and a three-arm spiral, targets a single hyperplane provably cannot separate. Ships with `make_xor` and `make_spiral` toy generators. - **Activation functions + weight initialization (Xavier/He) and why they matter**: `src/activations.py` is the dedicated treatment of the nonlinearity and the initial scale. Each activation (`linear`, `tanh`, `sigmoid`, `relu`, `leaky_relu`) exposes `forward` and a local `backward(z, grad_out)` that multiplies by `act'(z)`, so a hand-written backprop step can drop it in. Xavier/Glorot draws `N(0, 2/(fan_in+fan_out))` (or the matching uniform bound) to keep both forward and backward variance stable for symmetric activations; He/Kaiming draws `N(0, 2/fan_in)` so a ReLU stack does not quietly die after a few layers. `forward_variance_profile` stacks affine+activation layers from unit-variance noise and returns the per-layer activation variance: naive `N(0,1)` weights explode, He keeps a ReLU stack `O(1)`, and Xavier on the same ReLU stack fades, which is the usual silent failure mode when the scheme and the nonlinearity disagree. +- **SGD, Momentum, Adam from scratch, convergence compared on the same net**: `src/optimizers.py` implements the three standard first-order update rules as numpy-only classes that own their state (velocity for momentum, bias-corrected moments for Adam) and mutate a flat list of parameter arrays in place. The MLP training loop calls `optimizer.step(params, grads)` after the hand-written backprop pass, so swapping the rule never touches the gradient math. `compare_optimizers` retrains the same architecture on the same data with the same seed for each factory, which keeps init and minibatch order fixed and isolates the update rule; on XOR, all three cut loss, and Adam typically pulls ahead of plain SGD early because its per-parameter rates absorb the uneven scale of the gradient. ## Usage @@ -104,6 +110,24 @@ print(history[0], history[-1]) # cross-entropy falls over training print(model.predict_proba(X[:3])) # per-class probabilities that sum to 1 ``` +Train the same net under SGD, momentum, and Adam and compare loss curves: + +```python +from src.mlp import make_xor, fit, accuracy +from src.optimizers import SGD, Momentum, Adam, compare_optimizers + +X, y = make_xor(n=400, seed=0) + +# plug any optimizer into the existing MLP trainer +model, history = fit(X, y, hidden=(16,), epochs=200, optimizer=Adam(lr=0.01)) +print(accuracy(model, X, y), history[0], history[-1]) + +# same init + minibatches; only the update rule changes +curves = compare_optimizers(X, y, epochs=150, seed=0) +for name, h in curves.items(): + print(name, h[0], "->", h[-1]) +``` + Compare init schemes by watching activation variance with depth: ```python diff --git a/src/mlp.py b/src/mlp.py index d7be74f..9ca4753 100644 --- a/src/mlp.py +++ b/src/mlp.py @@ -147,14 +147,17 @@ def fit( epochs: int = 200, batch_size: int = 32, seed: int = 0, + optimizer: object | None = None, ) -> tuple[MLP, list[float]]: - """Train an MLP classifier by minibatch SGD, returning it and the loss curve. + """Train an MLP classifier by minibatch descent, returning it and the loss curve. Inputs are standardized to zero mean and unit variance so one learning rate works across features, and the standardizer is stored on the model so it consumes raw X at predict time. Labels may be any hashable values; they are - mapped to softmax columns and remembered on the model. Returns the trained - model and the per-epoch training cross-entropy. + mapped to softmax columns and remembered on the model. Pass an `optimizer` + from `src.optimizers` (SGD, Momentum, Adam) to swap the update rule; when + omitted, plain SGD at `lr` is used. Returns the trained model and the + per-epoch training cross-entropy. """ if activation not in ACTIVATIONS: raise ValueError(f"unknown activation {activation!r}") @@ -167,6 +170,11 @@ def fit( if any(h <= 0 for h in hidden): raise ValueError("hidden layer sizes must be positive") + # local import keeps the mlp module free of a hard optimizers dependency + from src.optimizers import SGD + + opt = optimizer if optimizer is not None else SGD(lr=lr) + Xm, yv = _check_xy(X, y) n, d = Xm.shape @@ -192,7 +200,9 @@ def fit( for start in range(0, n, batch_size): idx = order[start : start + batch_size] xb, yb = Xs[idx], y_onehot[idx] - _sgd_step(weights, biases, act, xb, yb, lr) + grad_w, grad_b = _backprop(weights, biases, act, xb, yb) + # flat list so one optimizer step covers every free parameter + opt.step([*weights, *biases], [*grad_w, *grad_b]) proba = _softmax(_forward_logits(weights, biases, act, Xs)) history.append(cross_entropy(proba, y_onehot)) @@ -254,20 +264,6 @@ def _backprop( return grad_w, grad_b -def _sgd_step( - weights: list[Array], - biases: list[Array], - act: _Activation, - xb: Array, - yb: Array, - lr: float, -) -> None: - grad_w, grad_b = _backprop(weights, biases, act, xb, yb) - for i in range(len(weights)): - weights[i] -= lr * grad_w[i] - biases[i] -= lr * grad_b[i] - - def make_xor(n: int = 400, noise: float = 0.15, seed: int = 0) -> tuple[Array, Array]: """XOR: four gaussian blobs at the corners, labeled by parity of the corner. diff --git a/src/optimizers.py b/src/optimizers.py new file mode 100644 index 0000000..65fadbd --- /dev/null +++ b/src/optimizers.py @@ -0,0 +1,178 @@ +"""First-order optimizers as pure numpy parameter updates. + +Each optimizer owns its step rule and any state (velocity, moments). The +training loop still owns the loss and the gradients; it hands a flat list of +parameter arrays and matching gradients, and the optimizer mutates the +parameters in place. That split is what lets the same MLP train under SGD, +momentum, or Adam without changing the backprop code. + +Formulas (all elementwise): + +- SGD: θ ← θ − η · g +- Momentum: v ← β v + g; θ ← θ − η · v (PyTorch-style, no dampening) +- Adam: m ← β₁ m + (1−β₁) g + v ← β₂ v + (1−β₂) g² + m̂ = m / (1−β₁ᵗ), v̂ = v / (1−β₂ᵗ) + θ ← θ − η · m̂ / (√v̂ + ε) +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + +Array = NDArray[np.float64] + + +class Optimizer(Protocol): + def step(self, params: Sequence[Array], grads: Sequence[Array]) -> None: ... + + +def _validate_pairs(params: Sequence[Array], grads: Sequence[Array]) -> None: + if len(params) != len(grads): + raise ValueError( + f"params has {len(params)} tensors but grads has {len(grads)}" + ) + for i, (p, g) in enumerate(zip(params, grads, strict=True)): + if p.shape != g.shape: + raise ValueError( + f"param/grad shape mismatch at index {i}: {p.shape} vs {g.shape}" + ) + + +class SGD: + """Plain gradient descent: θ ← θ − η g.""" + + def __init__(self, lr: float = 0.1) -> None: + if lr <= 0.0: + raise ValueError("lr must be positive") + self.lr = float(lr) + + def step(self, params: Sequence[Array], grads: Sequence[Array]) -> None: + _validate_pairs(params, grads) + for p, g in zip(params, grads, strict=True): + p -= self.lr * g + + +class Momentum: + """Heavy-ball momentum: accumulates a velocity so steps keep going downhill. + + Helps on elongated valleys where pure SGD zig-zags. Velocity is initialized + lazily from the first grad shapes so callers need not pre-register params. + """ + + def __init__(self, lr: float = 0.05, beta: float = 0.9) -> None: + if lr <= 0.0: + raise ValueError("lr must be positive") + if not 0.0 <= beta < 1.0: + raise ValueError("beta must be in [0, 1)") + self.lr = float(lr) + self.beta = float(beta) + self._v: list[Array] | None = None + + def step(self, params: Sequence[Array], grads: Sequence[Array]) -> None: + _validate_pairs(params, grads) + if self._v is None: + self._v = [np.zeros_like(p) for p in params] + elif len(self._v) != len(params): + raise ValueError("param list length changed between steps") + for i, (p, g) in enumerate(zip(params, grads, strict=True)): + self._v[i] = self.beta * self._v[i] + g + p -= self.lr * self._v[i] + + +class Adam: + """Adaptive moment estimation (Kingma & Ba). + + Per-parameter learning rates from the second moment, with bias correction so + the early steps are not dominated by zero-initialized m and v. Default + β₁=0.9, β₂=0.999, ε=1e-8 match the paper and common library defaults. + """ + + def __init__( + self, + lr: float = 0.01, + beta1: float = 0.9, + beta2: float = 0.999, + eps: float = 1e-8, + ) -> None: + if lr <= 0.0: + raise ValueError("lr must be positive") + if not 0.0 <= beta1 < 1.0: + raise ValueError("beta1 must be in [0, 1)") + if not 0.0 <= beta2 < 1.0: + raise ValueError("beta2 must be in [0, 1)") + if eps <= 0.0: + raise ValueError("eps must be positive") + self.lr = float(lr) + self.beta1 = float(beta1) + self.beta2 = float(beta2) + self.eps = float(eps) + self._m: list[Array] | None = None + self._v: list[Array] | None = None + self.t = 0 + + def step(self, params: Sequence[Array], grads: Sequence[Array]) -> None: + _validate_pairs(params, grads) + if self._m is None: + self._m = [np.zeros_like(p) for p in params] + self._v = [np.zeros_like(p) for p in params] + assert self._v is not None + if len(self._m) != len(params): + raise ValueError("param list length changed between steps") + + self.t += 1 + # bias correction grows toward 1; without it early steps are tiny + bc1 = 1.0 - self.beta1**self.t + bc2 = 1.0 - self.beta2**self.t + for i, (p, g) in enumerate(zip(params, grads, strict=True)): + self._m[i] = self.beta1 * self._m[i] + (1.0 - self.beta1) * g + self._v[i] = self.beta2 * self._v[i] + (1.0 - self.beta2) * (g * g) + m_hat = self._m[i] / bc1 + v_hat = self._v[i] / bc2 + p -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps) + + +def compare_optimizers( + X: Array, + y: Array, + *, + factories: dict[str, Callable[[], Optimizer]] | None = None, + hidden: tuple[int, ...] = (16,), + activation: str = "tanh", + epochs: int = 150, + batch_size: int = 32, + seed: int = 0, +) -> dict[str, list[float]]: + """Train the same MLP under several optimizers; return per-epoch loss curves. + + Each factory builds a fresh optimizer so moment/velocity state never leaks + across runs. The same `seed` is reused so init and minibatch order match, + isolating the update rule as the only difference. + """ + # late import: optimizers is usable without pulling the full MLP module + from src.mlp import fit + + if factories is None: + factories = { + "sgd": lambda: SGD(lr=0.1), + "momentum": lambda: Momentum(lr=0.05, beta=0.9), + "adam": lambda: Adam(lr=0.01), + } + histories: dict[str, list[float]] = {} + for name, factory in factories.items(): + _, history = fit( + X, + y, + hidden=hidden, + activation=activation, + epochs=epochs, + batch_size=batch_size, + seed=seed, + optimizer=factory(), + ) + histories[name] = history + return histories diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py new file mode 100644 index 0000000..6652786 --- /dev/null +++ b/tests/test_optimizers.py @@ -0,0 +1,140 @@ +import numpy as np +import pytest + +from src.mlp import accuracy, fit, make_xor +from src.optimizers import SGD, Adam, Momentum, compare_optimizers + + +def test_sgd_step_matches_closed_form(): + p = np.array([1.0, -2.0, 3.0]) + g = np.array([0.5, 1.0, -1.0]) + orig = p.copy() + SGD(lr=0.2).step([p], [g]) + assert np.allclose(p, orig - 0.2 * g) + + +def test_momentum_accumulates_velocity(): + p = np.array([0.0]) + opt = Momentum(lr=1.0, beta=0.5) + opt.step([p], [np.array([2.0])]) # v = 2, p = -2 + assert p[0] == pytest.approx(-2.0) + opt.step([p], [np.array([2.0])]) # v = 0.5*2 + 2 = 3, p = -5 + assert p[0] == pytest.approx(-5.0) + + +def test_adam_first_step_bias_corrected(): + # t=1, m = (1-β1)g, v = (1-β2)g², m̂ = g, v̂ = g² + # update = lr * g / (|g| + eps) → sign(g) * lr / (1 + eps) for |g|=1 + p = np.array([0.0]) + g = np.array([1.0]) + opt = Adam(lr=0.1, beta1=0.9, beta2=0.999, eps=1e-8) + opt.step([p], [g]) + expected = -0.1 * 1.0 / (np.sqrt(1.0) + 1e-8) + assert p[0] == pytest.approx(expected) + assert opt.t == 1 + + +def test_adam_second_moment_shrinks_noisy_steps(): + # a large spike in g should produce a smaller step once v has warmed up + p = np.zeros(1) + opt = Adam(lr=0.1) + for _ in range(20): + opt.step([p], [np.array([0.1])]) + pos_after_steady = p[0] + p2 = np.zeros(1) + opt2 = Adam(lr=0.1) + for _ in range(19): + opt2.step([p2], [np.array([0.1])]) + before_spike = p2[0] + opt2.step([p2], [np.array([10.0])]) + spike_delta = abs(p2[0] - before_spike) + steady_step = abs(pos_after_steady) / 20 + # spike is 100x the steady grad, but Adam's √v damps the step well below 100x + assert spike_delta < 100 * steady_step + + +def test_invalid_hyperparams(): + with pytest.raises(ValueError): + SGD(lr=0.0) + with pytest.raises(ValueError): + Momentum(lr=0.1, beta=1.0) + with pytest.raises(ValueError): + Momentum(lr=0.1, beta=-0.1) + with pytest.raises(ValueError): + Adam(lr=0.01, beta1=1.0) + with pytest.raises(ValueError): + Adam(lr=0.01, beta2=-0.1) + with pytest.raises(ValueError): + Adam(lr=0.01, eps=0.0) + + +def test_shape_and_length_mismatch_raise(): + p = np.zeros(3) + with pytest.raises(ValueError): + SGD(lr=0.1).step([p], [np.zeros(2)]) + with pytest.raises(ValueError): + SGD(lr=0.1).step([p], [np.zeros(3), np.zeros(3)]) + + +def test_empty_params_is_noop(): + for opt in (SGD(lr=0.1), Momentum(lr=0.1), Adam(lr=0.01)): + opt.step([], []) + + +def test_mlp_fit_with_each_optimizer_learns_xor(): + X, y = make_xor(n=400, seed=0) + configs = [ + SGD(lr=0.1), + Momentum(lr=0.05, beta=0.9), + Adam(lr=0.01), + ] + for opt in configs: + model, history = fit( + X, y, hidden=(16,), activation="tanh", epochs=200, seed=0, optimizer=opt + ) + assert history[-1] < history[0] + assert accuracy(model, X, y) > 0.9 + + +def test_compare_optimizers_all_descend_and_differ(): + X, y = make_xor(n=300, seed=1) + histories = compare_optimizers( + X, y, hidden=(16,), epochs=120, seed=1, batch_size=32 + ) + assert set(histories) == {"sgd", "momentum", "adam"} + for name, h in histories.items(): + assert len(h) == 120 + assert h[-1] < h[0], f"{name} did not reduce loss" + # same init/data, different rules → not identical curves + assert histories["sgd"] != histories["adam"] + assert histories["sgd"] != histories["momentum"] + + +def test_adam_outpaces_sgd_early_on_xor(): + # same seed and architecture; Adam's adaptive rates usually pull ahead early + X, y = make_xor(n=400, seed=2) + histories = compare_optimizers( + X, + y, + factories={ + "sgd": lambda: SGD(lr=0.05), + "adam": lambda: Adam(lr=0.02), + }, + hidden=(16,), + epochs=40, + seed=2, + ) + # mid-run checkpoint: Adam should be clearly ahead of plain SGD + mid = 20 + assert histories["adam"][mid] < histories["sgd"][mid] + + +def test_default_fit_still_uses_sgd(): + # regression: omit optimizer → same behavior as before (plain SGD) + X, y = make_xor(n=200, seed=3) + m1, h1 = fit(X, y, hidden=(8,), epochs=50, lr=0.1, seed=3) + m2, h2 = fit( + X, y, hidden=(8,), epochs=50, seed=3, optimizer=SGD(lr=0.1) + ) + assert h1 == h2 + assert np.allclose(m1.weights[0], m2.weights[0])