Too Good to Be True: Overfitting and Regularization

# Chapter 11: overfitting, and three ways to fight it: regularization, early stopping, more data.
# Uses the Tensor class from Chapter 6 (tensor.py) and the Adam class from Chapter 10, saved as optim.py
# with 'import numpy as np' at the top.
import numpy as np
from tensor import Tensor
from optim import Adam

# === 1. Twelve rides, ten polynomials ===========================================
rng = np.random.default_rng(7)                     # Chapter 7's year of rides
year_km = np.round(rng.uniform(1, 15, 240), 1)
year_min = np.round(3 * year_km + 5 + rng.normal(0, 1.5, 240), 1)
train_km, train_min = year_km[:12], year_min[:12]  # we only get to see 12 rides
val_km, val_min = year_km[12:], year_min[12:]      # the other 228 judge the fit
mu, sd = train_km.mean(), train_km.std()


def powers(km, degree):
    z = (km - mu) / sd                             # standardized, as in Chapter 7
    return np.column_stack([z ** k for k in range(degree + 1)])


def fit(degree, lam=0.0):
    """Least squares on 1, z, z², …, plus lam × (sum of squared coefficients, except the constant)."""
    A = powers(train_km, degree)
    penalty = lam * len(A) * np.eye(degree + 1)
    penalty[0, 0] = 0
    return np.linalg.solve(A.T @ A + penalty, A.T @ train_min)


def mse(c, km, minutes):
    return float(np.mean((powers(km, len(c) - 1) @ c - minutes) ** 2))

print("degree   train MSE   validation MSE   predicts 15 km as")
for degree in range(1, 11):
    c = fit(degree)
    print(f"{degree:6}   {mse(c, train_km, train_min):9.3f}   {mse(c, val_km, val_min):14.3f}"
          f"   {(powers(np.array([15.0]), degree) @ c)[0]:10.1f} min")
print("(the noise in these rides has variance 1.5² = 2.25: no model can do better on average)")

print("\ndegree 9 with an L2 penalty:")
print("   lam    train MSE   validation MSE   sum of squared coefficients")
for lam in [0, 0.001, 0.01, 0.1, 1]:
    c = fit(9, lam)
    print(f"{lam:6}   {mse(c, train_km, train_min):9.3f}   {mse(c, val_km, val_min):14.3f}   {np.sum(c[1:] ** 2):12.1f}")

# === 2. A network that memorizes: Chapter 9's activities ===========================
def make_days(n, seed):
    """Same process as Chapter 9: read if rainy, else kite if windy, else picnic; 1 day in 10 at random."""
    r = np.random.default_rng(seed)
    X = np.round(r.random((n, 3)), 2)
    rule = np.where(X[:, 1] > 0.5, 2, np.where(X[:, 2] > 0.6, 1, 0))
    labels = rule.copy()
    noisy = r.random(n) < 0.1
    labels[noisy] = r.integers(0, 3, noisy.sum())
    return X, labels, rule

X_train, y_train, _ = make_days(150, 9)            # Chapter 9's 150 days
X_val, y_val, rule_val = make_days(1000, 99)       # new days, for choosing settings
X_test, y_test, rule_test = make_days(1000, 123)   # more new days, looked at once, at the very end
print(f"\nthe rule itself: {np.mean(rule_val == y_val):.1%} of validation days, {np.mean(rule_test == y_test):.1%} of test days")


def softmax_cross_entropy(Z, labels):
    E = (Z + (-Z.data.max(axis=1, keepdims=True))).exp()
    P = E / (E @ Tensor(np.ones((3, 1))))
    return -(Tensor(np.eye(3)[labels]) * P.log()).sum() / Z.shape[0], P


def scores(params, X):
    W1, b1, W2, b2 = params
    return (Tensor(X) @ W1.T + b1).relu() @ W2.T + b2


def evaluate(params, X, labels):
    loss, P = softmax_cross_entropy(scores(params, X), labels)
    return float(loss.data), float(np.mean(P.data.argmax(axis=1) == labels))


def train(X, labels, epochs=3000, lam=0.0, hidden=64, log_every=None, seed=0):
    """Full-batch Adam (η = 0.01) on a 3-64-3 network, with an optional L2 penalty lam·(ΣW1² + ΣW2²)."""
    r = np.random.default_rng(seed)
    W1, b1 = Tensor(r.normal(size=(hidden, 3))), Tensor(np.zeros(hidden))
    W2, b2 = Tensor(r.normal(size=(3, hidden)) * 0.5), Tensor(np.zeros(3))
    params = [W1, b1, W2, b2]
    opt = Adam(params, 0.01)
    log = []
    for epoch in range(1, epochs + 1):
        loss, _ = softmax_cross_entropy(scores(params, X), labels)
        if lam:
            loss = loss + lam * ((W1 * W1).sum() + (W2 * W2).sum())
        for p in params:
            p.grad = np.zeros_like(p.data)
        loss.backward()
        opt.step()
        if log_every and epoch % log_every == 0:
            log.append((epoch, *evaluate(params, X, labels), *evaluate(params, X_val, y_val)))
    return params, log

print("\n64 hidden units (451 knobs), no regularization:")
print(" epoch   train loss  train acc   val loss  val acc")
params, log = train(X_train, y_train, log_every=10)
for epoch, tl, ta, vl, va in log:
    if epoch in (10, 50, 110, 300, 600, 1000, 2000, 3000):
        print(f"{epoch:6}   {tl:9.3f}   {ta:7.1%}   {vl:8.3f}  {va:6.1%}")
best = min(log, key=lambda row: row[3])
print(f"lowest validation loss at epoch {best[0]}: {best[3]:.3f} ({best[4]:.1%})")

print("\nsame network, 3,000 epochs, with an L2 penalty:")
for lam in [0.0001, 0.001, 0.003, 0.01]:
    p, _ = train(X_train, y_train, lam=lam)
    tl, ta = evaluate(p, X_train, y_train)
    vl, va = evaluate(p, X_val, y_val)
    print(f"  lam = {lam:<7} train {tl:.3f} ({ta:.1%})   validation {vl:.3f} ({va:.1%})")

# === 3. The final exam: the test set, once ===========================================
print("\ntest set (1,000 days never used for any choice):")
candidates = [("no regularization, 3,000 epochs", train(X_train, y_train)[0]),
              (f"early stopping at epoch {best[0]}", train(X_train, y_train, epochs=best[0])[0]),
              ("L2, lam = 0.003", train(X_train, y_train, lam=0.003)[0])]
X_more, y_more, _ = make_days(1500, 10)
candidates.append(("1,500 training days, no regularization", train(X_more, y_more)[0]))
for name, p in candidates:
    vl, va = evaluate(p, X_val, y_val)
    tl, ta = evaluate(p, X_test, y_test)
    print(f"  {name:40} validation {va:.1%}   test {ta:.1%} (loss {tl:.3f})")