feat: Add SGD, Momentum, and Adam with same-net comparison

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

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.
 

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