scripts/exp_composite.py

"""Claim 6 -- ULMC vs composite overdamped LMC when tr(H) << d
(ridge-separable targets).

Ridge-separable target of Freund et al. (2022) / Liu et al. (2023):
    V(x) = sum_{j=1}^m (beta/2) <w_j, x>^2  +  (alpha/2) ||x||^2 ,  w_j orthonormal
=> spectrum  a = beta+alpha  (m times),  alpha  (d-m times)
=> H = sum_j beta w_j w_j^T + alpha I,  tr(H) = m beta + alpha d,  kappa=(beta+alpha)/alpha
The "effective dimension" tr(H)/beta = m + d/kappa stays O(m) while d grows, if
kappa grows with d.  This is exactly the tr(H) << d regime of Section 3.3.

Competitors, all counted in gradient evaluations:
  * ULMC            (this paper, Thm 4.3)      Otilde(kappa^{3/2} beta^{-1/2} trH^{1/2}/eps)
  * RMD             (this paper, Thm 5.2)      Otilde(kappa (trH/beta)^{1/3} eps^{-2/3})
                    -- 2 gradient evaluations per iteration
  * composite LMC   (Freund et al. 2022)       Otilde(kappa^2 beta^{-1} trH / eps^2)
                    implemented as the exponential integrator that treats the
                    alpha-strongly-convex quadratic part exactly and the
                    remainder explicitly
  * plain LMC       (Euler-Maruyama of OLD)    Otilde(kappa^2 d / eps^2)
"""

import numpy as np

import common as C
import ulmc_core as U

NGRID = U.n_grid(int(4e10), 1.04)
res = {"target": "ridge separable: m stiff directions at beta+alpha, d-m at alpha"}

BETA, M_RIDGE, EPS = 1.0, 4, 0.05


def ridge(d, kappa, m=M_RIDGE, beta=BETA):
    alpha = beta / kappa
    a = np.array([beta + alpha, alpha])
    mult = np.array([float(m), float(d - m)])
    return a, mult


def one(scheme, a, mult, beta, eps, hs=None):
    S0, m0 = C.init_cold(a, beta)
    r = C.sweep_neps(
        scheme,
        a,
        mult,
        beta,
        eps,
        S0,
        m0,
        hs=hs if hs is not None else C.hgrid(beta, n=30, ratio=1.35),
        ngrid=NGRID,
    )
    r.pop("table")
    r["grad_evals"] = (
        (2 * r["N_eps"] if scheme == "rmd" else r["N_eps"]) if r["N_eps"] else None
    )
    return r


# --- main head-to-head: kappa = d, so tr(H)/(beta d) -> 0 --------------------
rows = []
for d in (32, 64, 128, 256, 512, 1024, 2048):
    kappa = float(d)
    a, mult = ridge(d, kappa)
    trH = float(M_RIDGE * BETA + (BETA / kappa) * d)
    row = {"d": d, "kappa": kappa, "trH": trH, "trH_over_beta_d": trH / (BETA * d)}
    for sch in ("ulmc", "rmd", "composite", "lmc"):
        row[sch] = one(sch, a, mult, BETA, EPS)
    row["speedup_ulmc_over_composite"] = (
        row["composite"]["N_eps"] / row["ulmc"]["N_eps"]
    )
    row["speedup_ulmc_over_lmc"] = row["lmc"]["N_eps"] / row["ulmc"]["N_eps"]
    row["speedup_rmd_over_composite"] = (
        row["composite"]["N_eps"] / row["rmd"]["grad_evals"]
    )
    rows.append(row)
    print(
        "d",
        d,
        {k: row[k]["N_eps"] for k in ("ulmc", "rmd", "composite", "lmc")},
        "ULMC/composite speedup",
        round(row["speedup_ulmc_over_composite"], 2),
        flush=True,
    )

res["head_to_head"] = {
    "rows": rows,
    "eps": EPS,
    "beta": BETA,
    "m_ridge": M_RIDGE,
    "kappa_equals_d": True,
    "ulmc_beats_composite_everywhere": bool(
        all(r["speedup_ulmc_over_composite"] > 1 for r in rows)
    ),
    "rmd_beats_composite_everywhere": bool(
        all(r["speedup_rmd_over_composite"] > 1 for r in rows)
    ),
    "max_speedup_ulmc_over_composite": float(
        max(r["speedup_ulmc_over_composite"] for r in rows)
    ),
}
for sch in ("ulmc", "rmd", "composite", "lmc"):
    s, r2 = C.fit_exponent([r["d"] for r in rows], [r[sch]["N_eps"] for r in rows])
    res["head_to_head"][f"{sch}_d_exponent"] = s
    res["head_to_head"][f"{sch}_r2"] = r2

# --- is composite LMC itself dimension-free?  fixed kappa, tr(H); vary d -----
rows2 = []
for d in (10, 20, 40, 80, 160, 320, 640, 900):
    a, m = C.spectrum(d, 0.01, 1.0, 10.0) if d >= 11 else (None, None)
    if a is None:
        continue
    row = {"d": d, "trH": 10.0}
    for sch in ("ulmc", "composite", "lmc"):
        row[sch] = one(sch, a, m, 1.0, 0.02)
    rows2.append(row)
res["dimension_freeness_of_baselines"] = {"rows": rows2, "eps": 0.02}
for sch in ("ulmc", "composite", "lmc"):
    s, r2 = C.fit_exponent([r["d"] for r in rows2], [r[sch]["N_eps"] for r in rows2])
    res["dimension_freeness_of_baselines"][f"{sch}_d_exponent"] = s

# --- composite LMC scaling law vs Freund et al. Otilde(kappa^2 beta^-1 trH/eps^2)
comp = {}
rows3 = []
for eps in np.geomspace(0.01, 0.1, 7):
    a, m = np.array([1.0, 0.01]), np.array([20.0, 2.0])
    rows3.append({"eps": float(eps), **one("composite", a, m, 1.0, float(eps))})
s, r2 = C.fit_exponent([r["eps"] for r in rows3], [r["N_eps"] for r in rows3])
comp["eps"] = {"rows": rows3, "fit_exponent": s, "predicted": -2.0, "r2": r2}
rows4 = []
for kappa in (2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0):
    a, m = np.array([1.0, 1.0 / kappa]), np.array([20.0, 2.0])
    rows4.append({"kappa": kappa, **one("composite", a, m, 1.0, 0.05)})
s, r2 = C.fit_exponent([r["kappa"] for r in rows4], [r["N_eps"] for r in rows4])
comp["kappa"] = {"rows": rows4, "fit_exponent": s, "predicted": 2.0, "r2": r2}
rows5 = []
for k in (2, 4, 8, 16, 32, 64, 128):
    a, m = np.array([1.0, 0.01]), np.array([float(k), 2.0])
    rows5.append({"trH": float(k + 0.02), **one("composite", a, m, 1.0, 0.02)})
s, r2 = C.fit_exponent([r["trH"] for r in rows5], [r["N_eps"] for r in rows5])
comp["trH"] = {"rows": rows5, "fit_exponent": s, "predicted": 1.0, "r2": r2}
res["composite_scaling_vs_freund"] = comp
print(
    "composite exponents:",
    {k: round(v["fit_exponent"], 3) for k, v in comp.items()},
    "predicted",
    {"eps": -2, "kappa": 2, "trH": 1},
)

C.dump("composite", res)