Recycled-curvature proximal optimizer / experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json
  2import time
  3import numpy as np
  4
  5
  6def soft_threshold(x, t):
  7    return np.sign(x) * np.maximum(np.abs(x) - t, 0.0)
  8
  9
 10class Quadratic:
 11    def __init__(self, eigs, seed=0):
 12        self.eigs = np.asarray(eigs, dtype=float)
 13        self.n = len(eigs)
 14        rng = np.random.default_rng(seed)
 15        q, _ = np.linalg.qr(rng.normal(size=(self.n, self.n)))
 16        self.A = q @ np.diag(self.eigs) @ q.T
 17        self.b = rng.normal(size=self.n)
 18        self.calls = 0
 19
 20    def grad(self, x):
 21        self.calls += 1
 22        return self.A @ x - self.b
 23
 24    def value(self, x):
 25        return 0.5 * x @ self.A @ x - self.b @ x
 26
 27    def prox_exact(self, z, gamma):
 28        return np.linalg.solve(np.eye(self.n) + gamma * self.A, z + gamma * self.b)
 29
 30
 31def residual_from_grad(x, z, gamma, grad):
 32    return gamma * grad + x - z
 33
 34
 35def bfgs_inverse(H, s, q):
 36    ys = float(q @ s)
 37    if ys <= 1e-12 * max(1.0, np.linalg.norm(q) * np.linalg.norm(s)):
 38        return H
 39    rho = 1.0 / ys
 40    I = np.eye(len(s))
 41    return (I - rho * np.outer(s, q)) @ H @ (I - rho * np.outer(q, s)) + rho * np.outer(s, s)
 42
 43
 44def verify_math():
 45    rng = np.random.default_rng(12)
 46    f = Quadratic([0.4, 1.0, 2.5, 5.0], seed=4)
 47    gamma = 0.37
 48    x = rng.normal(size=f.n)
 49    z0, z1 = rng.normal(size=(2, f.n))
 50    grad = f.grad(x)
 51    r0 = residual_from_grad(x, z0, gamma, grad)
 52    direct = residual_from_grad(x, z1, gamma, grad)
 53    transported = r0 + z0 - z1
 54    transport_err = np.linalg.norm(direct - transported)
 55
 56    secant_errors = []
 57    secant_ratios = []
 58    for _ in range(500):
 59        u, v = rng.normal(size=(2, f.n))
 60        s = v - u
 61        q = gamma * (f.A @ s) + s
 62        expected = np.dot(s, s) + gamma * s @ f.A @ s
 63        secant_errors.append(abs(s @ q - expected))
 64        secant_ratios.append((s @ q) / (s @ s))
 65
 66    # For exact H=(I+gamma A)^-1 and undamped unit Newton steps, residual
 67    # contracts to zero in one step. With H=I, the factor is 1-gamma*lambda.
 68    boundary = []
 69    for lam in [0.5, 1.0, 1.8, 2.2]:
 70        for g in [0.2, 0.5, 0.9, 1.1]:
 71            factor = abs(1.0 - g * lam)
 72            boundary.append((g * lam, factor, factor < 1.0))
 73    pred_ok = all((prod < 2.0 and stable) or (prod >= 2.0 and not stable)
 74                  for prod, _, stable in boundary)
 75    # Sweep confirms the exact scalar stability interval gamma*lambda in (0,2).
 76    observed_boundary = max(prod for prod, factor, _ in boundary if factor < 1.0)
 77    return {
 78        "transport_l2_error": float(transport_err),
 79        "max_secant_identity_error": float(max(secant_errors)),
 80        "secant_ratio_min": float(min(secant_ratios)),
 81        "secant_ratio_predicted_min": 1.0,
 82        "stability_sweep": [{"gamma_lambda": p, "observed_factor": a, "contractive": c} for p, a, c in boundary],
 83        "stability_prediction_pass": bool(pred_ok),
 84        "stability_interval_prediction": "0 < gamma*lambda < 2 for H=I",
 85        "stability_sweep_max_contracting_product": float(observed_boundary),
 86    }
 87
 88
 89def restarted_pg(f, centers, gamma, inner_steps, l1, step):
 90    x = centers[0].copy()
 91    calls = 0
 92    residuals = []
 93    for z in centers:
 94        x = z.copy()  # independently restarted solve
 95        for _ in range(inner_steps):
 96            g = f.grad(x); calls += 1
 97            r = residual_from_grad(x, z, gamma, g)
 98            x -= step * r
 99        residuals.append(float(np.linalg.norm(residual_from_grad(x, z, gamma, f.A @ x - f.b))))
100        x = soft_threshold(x, gamma * l1)
101    return x, calls, residuals
102
103
104def recycled_bfgs(f, centers, gamma, inner_steps, l1, alpha=0.8):
105    # Frozen CR-DRS-style predictor: pseudo-steps use B=H^{-1} only;
106    # exactly one expensive gradient is evaluated at the final predictor.
107    x = centers[0].copy()
108    H = np.eye(f.n)
109    old_z = centers[0].copy()
110    g = f.grad(x); calls = 1
111    r = residual_from_grad(x, old_z, gamma, g)
112    residuals = []
113    for z in centers:
114        r = r + old_z - z  # exact center transport, no gradient call
115        x_anchor, r_anchor = x.copy(), r.copy()
116        B = np.linalg.inv(H)
117        for _ in range(inner_steps):
118            d = -H @ r
119            Bd = B @ d
120            denom = float(d @ Bd)
121            numer = float(-(r @ d))
122            eta = min(alpha, numer / denom) if denom > 1e-14 and numer > 0 else 0.0
123            if eta == 0.0:
124                break
125            dx = eta * d
126            x = x + dx
127            r = r + B @ dx  # modeled residual; no expensive evaluation
128        # One true gradient at the endpoint, as in the paper.
129        gt = f.grad(x); calls += 1
130        rt = residual_from_grad(x, z, gamma, gt)
131        s, q = x - x_anchor, rt - r_anchor
132        H = bfgs_inverse(H, s, q)
133        r = rt
134        residuals.append(float(np.linalg.norm(r)))
135        x = soft_threshold(x, gamma * l1)
136        old_z = z.copy()
137    return x, calls, residuals
138
139
140def run_experiment():
141    rng = np.random.default_rng(7)
142    f1 = Quadratic(np.linspace(0.5, 8.0, 24), seed=9)
143    f2 = Quadratic(np.linspace(0.5, 8.0, 24), seed=9)
144    centers = []
145    z = rng.normal(scale=1.5, size=24)
146    for k in range(30):
147        centers.append(z.copy())
148        z = 0.88 * z + 0.12 * rng.normal(size=24)
149    gamma, inner, l1 = 0.08, 3, 0.015
150    # conservative fixed residual-gradient step; same three expensive calls/outer update
151    step = 1.0 / (1.0 + gamma * 8.0)
152    t0 = time.perf_counter(); _, cb, rb = restarted_pg(f1, centers, gamma, inner, l1, step); tb = time.perf_counter()-t0
153    t0 = time.perf_counter(); _, ci, ri = recycled_bfgs(f2, centers, gamma, inner, l1); ti = time.perf_counter()-t0
154    return {
155        "outer_updates": len(centers), "inner_steps": inner,
156        "baseline_gradient_calls": cb, "idea_gradient_calls": ci,
157        "gradient_call_reduction_percent": 100.0*(cb-ci)/cb,
158        "baseline_final_residual": rb[-1], "idea_final_residual": ri[-1],
159        "baseline_mean_residual": float(np.mean(rb)), "idea_mean_residual": float(np.mean(ri)),
160        "baseline_seconds": tb, "idea_seconds": ti
161    }
162
163
164if __name__ == "__main__":
165    out = {"math_verification": verify_math(), "mini_experiment": run_experiment()}
166    with open("results.json", "w") as f:
167        json.dump(out, f, indent=2)
168    print(json.dumps(out, indent=2))