Linear regression: normal equation vs SGD

README.md

@@ -16,10 +16,16 @@ Every deep learning framework is, at its core, a graph that records operations a
 - 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
+- Ordinary least squares in closed form via the normal equation
+- Ridge regularization that shrinks slopes without penalizing the intercept
+- Minibatch stochastic gradient descent with feature standardization for conditioning
+- The equivalence of the closed-form and iterative solutions on a convex loss
+- Pseudoinverse fallback for rank-deficient design matrices
 
 ## What's implemented
 
 - **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.
+- **Linear regression, two ways**: `fit_normal_equation` solves ordinary least squares exactly by setting the gradient to zero and solving `(AᵀA + λR)θ = Aᵀy`, with optional ridge and a least-norm fallback when `AᵀA` is singular. `fit_sgd` fits the same model with minibatch gradient descent on standardized features, then maps the weights back to raw feature space. Both are checked to recover the true coefficients and to agree with each other, so the exact-vs-iterative tradeoff is measurable.
 
 ## Usage
 
@@ -32,6 +38,23 @@ minimize(lambda: (x - 3.0) ** 2, [x], steps=200, lr=0.1)
 print(x.data)  # ~3.0
 ```
 
+Fit a line two ways and compare:
+
+```python
+import numpy as np
+from src.linreg import fit_normal_equation, fit_sgd, mse
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-2, 2, size=(200, 3))
+y = X @ np.array([2.0, -3.0, 0.5]) + 1.5
+
+exact = fit_normal_equation(X, y)              # closed form, no tuning
+approx, history = fit_sgd(X, y, lr=0.1, epochs=300)  # iterative, streams over data
+
+print(exact.weights, exact.bias)   # ~[2, -3, 0.5], 1.5
+print(mse(approx, X, y))           # ~0, matches the closed form
+```
+
 Differentiate an arbitrary scalar expression:
 
 ```python

src/linreg.py

@@ -0,0 +1,137 @@
+"""Linear regression two ways: the normal equation and minibatch SGD.
+
+Ordinary least squares has a closed form, so you can solve it exactly by
+inverting a small matrix. It also has a convex loss, so gradient descent finds
+the same answer iteratively. This module builds both, using numpy for the
+linear algebra but writing out the gradient and the normal-equation solve by
+hand, so the tradeoff is visible: the closed form is exact and needs no tuning
+but costs O(d^3) in the feature count, while SGD is O(nd) per step, streams
+over data, and is what actually scales to large models.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+Array = NDArray[np.float64]
+
+
+@dataclass(frozen=True)
+class LinearModel:
+    weights: Array
+    bias: float
+
+    def predict(self, X: Array) -> Array:
+        return _as_2d(X) @ self.weights + self.bias
+
+
+def mse(model: LinearModel, X: Array, y: Array) -> float:
+    resid = model.predict(X) - np.asarray(y, dtype=np.float64)
+    return float(np.mean(resid**2))
+
+
+def _as_2d(X: Array) -> Array:
+    arr = np.asarray(X, dtype=np.float64)
+    return arr.reshape(-1, 1) if arr.ndim == 1 else arr
+
+
+def _check_xy(X: Array, y: Array) -> tuple[Array, Array]:
+    Xm = _as_2d(X)
+    yv = np.asarray(y, dtype=np.float64).reshape(-1)
+    if Xm.shape[0] == 0:
+        raise ValueError("need at least one sample")
+    if Xm.shape[0] != yv.shape[0]:
+        raise ValueError(
+            f"X has {Xm.shape[0]} rows but y has {yv.shape[0]} entries"
+        )
+    return Xm, yv
+
+
+def fit_normal_equation(X: Array, y: Array, l2: float = 0.0) -> LinearModel:
+    """Solve OLS in closed form via (A^T A + l2 R) theta = A^T y.
+
+    A augments X with a bias column of ones. R is the identity with a zero in
+    the bias slot so ridge never shrinks the intercept. When A^T A is singular
+    (rank-deficient X, no regularization) we fall back to the least-norm
+    pseudoinverse solution instead of raising.
+    """
+    if l2 < 0.0:
+        raise ValueError("l2 must be non-negative")
+    Xm, yv = _check_xy(X, y)
+    n, d = Xm.shape
+    A = np.hstack([np.ones((n, 1)), Xm])
+
+    gram = A.T @ A
+    if l2 > 0.0:
+        reg = np.eye(d + 1)
+        reg[0, 0] = 0.0
+        gram = gram + l2 * reg
+
+    rhs = A.T @ yv
+    try:
+        theta = np.linalg.solve(gram, rhs)
+    except np.linalg.LinAlgError:
+        theta = np.linalg.lstsq(A, yv, rcond=None)[0]
+
+    return LinearModel(weights=theta[1:], bias=float(theta[0]))
+
+
+def fit_sgd(
+    X: Array,
+    y: Array,
+    lr: float = 0.05,
+    epochs: int = 100,
+    batch_size: int = 32,
+    l2: float = 0.0,
+    seed: int = 0,
+) -> tuple[LinearModel, list[float]]:
+    """Fit the same model by minibatch gradient descent on standardized X.
+
+    Standardizing each feature to zero mean and unit variance before stepping
+    keeps the loss well-conditioned, so a single learning rate works regardless
+    of feature scale. The fitted weights are mapped back to the raw feature
+    space at the end, so the returned model consumes unscaled X exactly like
+    the closed-form one. Returns the model and the per-epoch training MSE.
+    """
+    if lr <= 0.0:
+        raise ValueError("lr must be positive")
+    if epochs <= 0:
+        raise ValueError("epochs must be positive")
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive")
+    if l2 < 0.0:
+        raise ValueError("l2 must be non-negative")
+
+    Xm, yv = _check_xy(X, y)
+    n, d = Xm.shape
+
+    mu = Xm.mean(axis=0)
+    sigma = Xm.std(axis=0)
+    sigma = np.where(sigma == 0.0, 1.0, sigma)  # constant features carry no info
+    Xs = (Xm - mu) / sigma
+
+    rng = np.random.default_rng(seed)
+    w = np.zeros(d)
+    b = 0.0
+    history: list[float] = []
+
+    for _ in range(epochs):
+        order = rng.permutation(n)
+        for start in range(0, n, batch_size):
+            idx = order[start : start + batch_size]
+            xb, yb = Xs[idx], yv[idx]
+            resid = xb @ w + b - yb
+            m = xb.shape[0]
+            grad_w = (2.0 / m) * (xb.T @ resid) + 2.0 * l2 * w
+            grad_b = (2.0 / m) * float(resid.sum())
+            w -= lr * grad_w
+            b -= lr * grad_b
+        resid_full = Xs @ w + b - yv
+        history.append(float(np.mean(resid_full**2)))
+
+    weights = w / sigma
+    bias = b - float(np.sum(w * mu / sigma))
+    return LinearModel(weights=weights, bias=bias), history

tests/test_linreg.py

@@ -0,0 +1,118 @@
+import dataclasses
+
+import numpy as np
+import pytest
+
+from src.linreg import LinearModel, fit_normal_equation, fit_sgd, mse
+
+
+def make_data(n=200, w=(2.0, -3.0, 0.5), b=1.5, noise=0.0, seed=0):
+    rng = np.random.default_rng(seed)
+    X = rng.uniform(-2.0, 2.0, size=(n, len(w)))
+    y = X @ np.array(w) + b + noise * rng.standard_normal(n)
+    return X, y
+
+
+def test_normal_equation_recovers_exact_weights_noiseless():
+    X, y = make_data(noise=0.0)
+    model = fit_normal_equation(X, y)
+    np.testing.assert_allclose(model.weights, [2.0, -3.0, 0.5], atol=1e-9)
+    assert model.bias == pytest.approx(1.5, abs=1e-9)
+    assert mse(model, X, y) == pytest.approx(0.0, abs=1e-18)
+
+
+def test_sgd_recovers_weights_noiseless():
+    X, y = make_data(noise=0.0)
+    model, history = fit_sgd(X, y, lr=0.1, epochs=300, batch_size=32)
+    np.testing.assert_allclose(model.weights, [2.0, -3.0, 0.5], atol=1e-2)
+    assert model.bias == pytest.approx(1.5, abs=1e-2)
+    assert history[-1] < history[0]
+
+
+def test_normal_equation_and_sgd_agree():
+    # convex loss + noiseless data: both methods land on the same minimum
+    X, y = make_data(noise=0.05, seed=7)
+    exact = fit_normal_equation(X, y)
+    approx, _ = fit_sgd(X, y, lr=0.1, epochs=500, batch_size=64)
+    np.testing.assert_allclose(approx.weights, exact.weights, atol=2e-2)
+    assert approx.bias == pytest.approx(exact.bias, abs=2e-2)
+
+
+def test_single_feature_and_1d_input():
+    rng = np.random.default_rng(1)
+    x = rng.uniform(-1.0, 1.0, size=50)
+    y = 4.0 * x - 2.0
+    model = fit_normal_equation(x, y)
+    assert model.weights.shape == (1,)
+    assert model.weights[0] == pytest.approx(4.0, abs=1e-9)
+    assert model.bias == pytest.approx(-2.0, abs=1e-9)
+    assert model.predict(np.array([0.5]))[0] == pytest.approx(0.0, abs=1e-9)
+
+
+def test_sgd_history_is_monotone_enough():
+    X, y = make_data(noise=0.1, seed=3)
+    _, history = fit_sgd(X, y, lr=0.05, epochs=200, batch_size=32)
+    assert len(history) == 200
+    # not strictly monotone with minibatches, but the trend must fall hard
+    assert history[-1] < 0.1 * history[0]
+
+
+def test_ridge_shrinks_weights_but_not_intercept():
+    X, y = make_data(noise=0.0, seed=2)
+    plain = fit_normal_equation(X, y, l2=0.0)
+    ridged = fit_normal_equation(X, y, l2=50.0)
+    assert np.linalg.norm(ridged.weights) < np.linalg.norm(plain.weights)
+    # intercept stays near its true value; only slopes are penalized
+    assert ridged.bias == pytest.approx(1.5, abs=0.2)
+
+
+def test_rank_deficient_falls_back_to_pseudoinverse():
+    # duplicate column makes X^T X singular; solve would blow up
+    rng = np.random.default_rng(4)
+    base = rng.uniform(-1.0, 1.0, size=(30, 1))
+    X = np.hstack([base, base])  # perfectly collinear
+    y = (3.0 * base[:, 0] + 1.0)
+    model = fit_normal_equation(X, y)
+    # split across the two identical columns, but predictions still fit
+    np.testing.assert_allclose(model.predict(X), y, atol=1e-6)
+
+
+def test_predict_shapes():
+    X, y = make_data(n=10, seed=5)
+    model = fit_normal_equation(X, y)
+    preds = model.predict(X)
+    assert preds.shape == (10,)
+
+
+def test_empty_input_raises():
+    with pytest.raises(ValueError):
+        fit_normal_equation(np.zeros((0, 3)), np.zeros(0))
+    with pytest.raises(ValueError):
+        fit_sgd(np.zeros((0, 3)), np.zeros(0))
+
+
+def test_mismatched_lengths_raise():
+    with pytest.raises(ValueError):
+        fit_normal_equation(np.zeros((5, 2)), np.zeros(4))
+
+
+def test_invalid_hyperparams_raise():
+    X, y = make_data(n=20)
+    for kwargs in ({"lr": 0.0}, {"epochs": 0}, {"batch_size": 0}, {"l2": -1.0}):
+        with pytest.raises(ValueError):
+            fit_sgd(X, y, **kwargs)
+    with pytest.raises(ValueError):
+        fit_normal_equation(X, y, l2=-1.0)
+
+
+def test_batch_larger_than_dataset_runs_full_batch():
+    X, y = make_data(n=20, noise=0.0, seed=6)
+    model, history = fit_sgd(X, y, lr=0.1, epochs=400, batch_size=256)
+    np.testing.assert_allclose(model.weights, [2.0, -3.0, 0.5], atol=1e-2)
+    assert len(history) == 400
+
+
+def test_linear_model_is_frozen():
+    m = LinearModel(weights=np.array([1.0]), bias=0.0)
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        m.bias = 2.0