scripts/claim5_kl_dro.py

"""Claim 5 -- Sec. 3.2 / Fig. 3: KL-regularised DRO on California Housing.

Objective (17):  L(theta) = lam * log( (1/n) sum_i exp( l_i(theta)/lam ) ),
                 l_i(theta) = (y_i - theta^T x_i)^2,  lam = 1.
Baseline (18) : batch-softmax gradient estimator of Levy et al. (2020).
Proposed (19) : safe-KL / SoftPlus estimator with the extra scalar alpha.
Optimiser     : SGD with Nesterov momentum 0.9 (PyTorch semantics), init = least squares.

The crux tested here is the STEPSIZE THRESHOLD: the largest eta in the paper's grid
{1e-8,...,1e-4} for which a method still converges, per batch size 10 / 100 / 1000.
"""

import argparse
import itertools
import json
import os
import time
from multiprocessing import Pool

import numpy as np
from scipy.special import expit, logsumexp
from sklearn.datasets import fetch_california_housing
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "outputs")
os.makedirs(OUT, exist_ok=True)
LAM = 1.0
N_EPOCHS = 50
CHECK_EPOCHS = [1, 2, 5, 10, 20, 30, 40, 50]


def get_data():
    raw = fetch_california_housing()
    X = StandardScaler().fit_transform(raw.data)
    y = raw.target.astype(np.float64)
    X = np.hstack([X, np.ones((X.shape[0], 1))])  # intercept column
    lr = LinearRegression(fit_intercept=True).fit(X[:, :-1], y)
    theta0 = np.concatenate([lr.coef_, [lr.intercept_]])
    return X, y, theta0


def objective(X, y, theta, lam=LAM):
    l = (X @ theta - y) ** 2
    return float(lam * (logsumexp(l / lam) - np.log(len(y))))


def run(method, batch, eta, rho, seed, X, y, theta0, lam=LAM, n_epochs=N_EPOCHS):
    rng = np.random.default_rng(seed)
    n, d = X.shape
    theta = theta0.copy()
    alpha = 0.0
    buf = np.zeros(d)
    buf_a = 0.0
    mom = 0.9
    L0 = objective(X, y, theta, lam)
    traj = {0: L0}
    diverged_at = None
    log_rho = np.log(rho) if rho is not None else None

    for ep in range(1, n_epochs + 1):
        perm = rng.permutation(n)
        for s in range(0, n - batch + 1, batch):
            idx = perm[s : s + batch]
            Xb, yb = X[idx], y[idx]
            r = Xb @ theta - yb
            l = r**2
            if method == "baseline":  # Eq. (18)
                p = np.exp(l / lam - logsumexp(l / lam))
                g = 2.0 * (Xb.T @ (p * r))
                ga = None
            else:  # Eq. (19)
                w = expit((l - alpha) / lam + log_rho) / rho
                g = 2.0 * (Xb.T @ (w * r)) / batch
                ga = 1.0 - float(np.mean(w))
            if not np.all(np.isfinite(g)):
                diverged_at = ep
                break
            buf = mom * buf + g
            theta = theta - eta * (g + mom * buf)
            if ga is not None:
                buf_a = mom * buf_a + ga
                alpha = alpha - eta * (ga + mom * buf_a)
        if diverged_at is not None:
            break
        if ep in CHECK_EPOCHS:
            v = objective(X, y, theta, lam)
            traj[ep] = v
            if not np.isfinite(v) or v > L0 + 50:
                diverged_at = ep
                break
    finite = np.isfinite(list(traj.values())).all() and diverged_at is None
    final = traj[max(traj)]
    return {
        "method": method,
        "batch": batch,
        "eta": eta,
        "rho": rho,
        "seed": seed,
        "L0": L0,
        "final": float(final),
        "traj": {str(k): float(v) for k, v in traj.items()},
        "diverged_at": diverged_at,
        "converged": bool(finite and final <= L0 + 1e-9),
        "improvement": float(L0 - final) if finite else None,
        "alpha": float(alpha),
    }


_G = {}


def _init():
    _G["data"] = get_data()


def _worker(t):
    X, y, th = _G["data"]
    return run(*t, X, y, th)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--procs", type=int, default=60)
    ap.add_argument("--seeds", type=int, default=10)
    args = ap.parse_args()
    etas = [1e-8, 1e-7, 1e-6, 1e-5, 1e-4]
    batches = [10, 100, 1000]
    rhos = [1e-4, 1e-3, 1e-2, 1e-1, 0.5, 0.9]
    seeds = list(range(args.seeds))
    tasks = [
        ("baseline", b, e, None, s)
        for b, e, s in itertools.product(batches, etas, seeds)
    ]
    tasks += [
        ("proposed", b, e, r, s)
        for b, e, r, s in itertools.product(batches, etas, rhos, seeds)
    ]
    print(f"{len(tasks)} runs", flush=True)
    t0 = time.time()
    with Pool(args.procs, initializer=_init) as pool:
        recs = pool.map(_worker, tasks, chunksize=1)
    print("elapsed", time.time() - t0, flush=True)

    X, y, th = get_data()
    summary = {
        "L0": objective(X, y, th),
        "n": int(X.shape[0]),
        "d": int(X.shape[1]),
        "lam": LAM,
        "n_epochs": N_EPOCHS,
        "seeds": args.seeds,
        "runs": recs,
    }
    with open(os.path.join(OUT, "claim5_kl_dro.json"), "w") as fh:
        json.dump(summary, fh)

    # ---- stepsize threshold table -------------------------------------------
    table = []
    for b in batches:
        for meth, r in [("baseline", None)] + [("proposed", rr) for rr in rhos]:
            row = {"method": meth, "rho": r, "batch": b, "per_eta": {}}
            for e in etas:
                rs = [
                    x
                    for x in recs
                    if x["method"] == meth
                    and x["batch"] == b
                    and x["eta"] == e
                    and x["rho"] == r
                ]
                ok = sum(x["converged"] for x in rs)
                fin = [x["final"] for x in rs if x["converged"]]
                row["per_eta"][str(e)] = {
                    "converged_seeds": ok,
                    "n_seeds": len(rs),
                    "mean_final": float(np.mean(fin)) if fin else None,
                    "std_final": float(np.std(fin)) if fin else None,
                }
            good = [
                e
                for e in etas
                if row["per_eta"][str(e)]["converged_seeds"] == len(seeds)
            ]
            row["max_stable_eta"] = max(good) if good else None
            best = None
            for e in etas:
                m = row["per_eta"][str(e)]["mean_final"]
                if m is not None and (best is None or m < best[1]):
                    best = (e, m)
            row["best_eta"] = best[0] if best else None
            row["best_final_objective"] = best[1] if best else None
            table.append(row)
    with open(os.path.join(OUT, "claim5_threshold_table.json"), "w") as fh:
        json.dump(table, fh, indent=2)
    for row in table:
        print(
            row["method"],
            row["rho"],
            "batch",
            row["batch"],
            "max_stable_eta",
            row["max_stable_eta"],
            "best_eta",
            row["best_eta"],
            "best_obj",
            (
                None
                if row["best_final_objective"] is None
                else round(row["best_final_objective"], 4)
            ),
            flush=True,
        )


if __name__ == "__main__":
    main()