S099 · Optimizers — SGD, Momentum, Adam, RMSprop

"""optimizer_lab.py — implement SGD, Momentum, RMSprop, Adam side-by-side."""
from __future__ import annotations
import numpy as np
 
RNG = np.random.default_rng(0)
 
 
# ---------------- optimizers ----------------
class SGD:
    def __init__(self, params: dict, lr: float = 0.1) -> None:
        self.params, self.lr = params, lr
    def step(self, grads: dict) -> None:
        for k in self.params:
            self.params[k] -= self.lr * grads[k]
 
 
class Momentum:
    def __init__(self, params: dict, lr: float = 0.1, beta: float = 0.9) -> None:
        self.params, self.lr, self.beta = params, lr, beta
        self.v = {k: np.zeros_like(v) for k, v in params.items()}
    def step(self, grads: dict) -> None:
        for k in self.params:
            self.v[k] = self.beta * self.v[k] + grads[k]
            self.params[k] -= self.lr * self.v[k]
 
 
class RMSprop:
    def __init__(self, params: dict, lr: float = 0.01, beta: float = 0.9, eps: float = 1e-8) -> None:
        self.params, self.lr, self.beta, self.eps = params, lr, beta, eps
        self.s = {k: np.zeros_like(v) for k, v in params.items()}
    def step(self, grads: dict) -> None:
        for k in self.params:
            self.s[k] = self.beta * self.s[k] + (1 - self.beta) * grads[k] ** 2
            self.params[k] -= self.lr * grads[k] / (np.sqrt(self.s[k]) + self.eps)
 
 
class Adam:
    def __init__(self, params: dict, lr: float = 1e-3, b1: float = 0.9, b2: float = 0.999,
                 eps: float = 1e-8) -> None:
        self.params, self.lr = params, lr
        self.b1, self.b2, self.eps = b1, b2, eps
        self.m = {k: np.zeros_like(v) for k, v in params.items()}
        self.v = {k: np.zeros_like(v) for k, v in params.items()}
        self.t = 0
    def step(self, grads: dict) -> None:
        self.t += 1
        for k in self.params:
            self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * grads[k]
            self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * grads[k] ** 2
            m_hat = self.m[k] / (1 - self.b1 ** self.t)   # bias correction
            v_hat = self.v[k] / (1 - self.b2 ** self.t)
            self.params[k] -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
 
 
# ---------------- race on narrow ellipsoid ----------------
def ellipsoid_race():
    """Minimise f(x, y) = 0.5·(a·x² + b·y²). Very stretched → SGD zig-zags."""
    a, b = 1.0, 20.0
    def loss(p):  return 0.5 * (a * p["x"] ** 2 + b * p["y"] ** 2)
    def grad(p):  return {"x": a * p["x"], "y": b * p["y"]}
 
    print("\n=== narrow-valley race (a=1, b=20) — steps to loss < 1e-3 ===")
    for name, opt_cls, lr in [
        ("SGD",      SGD,      0.05),
        ("Momentum", Momentum, 0.05),
        ("RMSprop",  RMSprop,  0.5),
        ("Adam",     Adam,     0.3),
    ]:
        p = {"x": np.array(-5.0), "y": np.array(-1.0)}
        opt = opt_cls(p, lr=lr) if opt_cls in (SGD,) else opt_cls(p, lr=lr)
        for step in range(1, 2001):
            opt.step(grad(p))
            if loss(p) < 1e-3:
                print(f"  {name:<10s}: converged in {step:>4d} steps (final={loss(p):.2e})")
                break
        else:
            print(f"  {name:<10s}: did NOT converge in 2000 (final={loss(p):.2e})")
 
 
# ---------------- MLP race on classification ----------------
def relu(z):      return np.maximum(0, z)
def relu_grad(z): return (z > 0).astype(z.dtype)
def softmax(z):
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z); return e / e.sum(axis=1, keepdims=True)
 
 
def init_mlp(sizes):
    p = {}
    for i in range(len(sizes) - 1):
        fin, fout = sizes[i], sizes[i + 1]
        p[f"W{i+1}"] = RNG.normal(0, np.sqrt(2.0 / fin), (fin, fout))
        p[f"b{i+1}"] = np.zeros(fout)
    return p
 
 
def forward_backward(p, X, y_oh):
    z1 = X @ p["W1"] + p["b1"]; a1 = relu(z1)
    z2 = a1 @ p["W2"] + p["b2"]; probs = softmax(z2)
    n = X.shape[0]
    loss = -float(np.mean(np.sum(y_oh * np.log(probs + 1e-12), axis=1)))
    dz2 = (probs - y_oh) / n
    grads = {
        "W2": a1.T @ dz2, "b2": dz2.sum(0),
    }
    da1 = dz2 @ p["W2"].T
    dz1 = da1 * relu_grad(z1)
    grads["W1"] = X.T @ dz1
    grads["b1"] = dz1.sum(0)
    return loss, grads, probs
 
 
def mlp_race():
    from sklearn.datasets import make_classification
    X, y = make_classification(n_samples=1500, n_features=10, n_informative=6,
                               n_classes=3, random_state=0)
    y_oh = np.zeros((y.shape[0], 3), dtype=np.float32); y_oh[np.arange(y.shape[0]), y] = 1.0
 
    print("\n=== MLP race — final loss after 50 epochs, batch=64 ===")
    for name, opt_cls, lr in [
        ("SGD",      SGD,      0.1),
        ("Momentum", Momentum, 0.1),
        ("RMSprop",  RMSprop,  0.005),
        ("Adam",     Adam,     0.003),
    ]:
        p = init_mlp([10, 32, 3])
        opt = opt_cls(p, lr=lr)
        for _ in range(50):
            idx = RNG.permutation(len(y))
            for start in range(0, len(y), 64):
                b = idx[start:start + 64]
                _, grads, _ = forward_backward(p, X[b], y_oh[b])
                opt.step(grads)
        loss, _, probs = forward_backward(p, X, y_oh)
        acc = float((probs.argmax(1) == y).mean())
        print(f"  {name:<10s}: loss={loss:.4f}  acc={acc:.4f}  (lr={lr})")
 
 
if __name__ == "__main__":
    ellipsoid_race()
    mlp_race()