1import json
  2import numpy as np
  3
  4
  5def interpolation_residual(xi, gi, fi, xj, gj, fj, mu, L):
  6    """H_mu,L(i,j), in the paper's s.q + tr(GQ) form, evaluated from samples."""
  7    dx = xi - xj
  8    dg = gi - gj
  9    if L <= mu:
 10        raise ValueError("need L > mu")
 11    trace_q = (-np.dot(gj, dx)
 12               - np.dot(dg - mu * dx, dg - mu * dx) / (2.0 * (L - mu))
 13               - mu * np.dot(dx, dx) / 2.0)
 14    return (fi - fj) + trace_q
 15
 16
 17def estimate_curvature(xs, gs, eps=1e-12):
 18    vals_L, vals_mu = [], []
 19    for x0, x1, g0, g1 in zip(xs[:-1], xs[1:], gs[:-1], gs[1:]):
 20        dx, dg = x1-x0, g1-g0
 21        n = max(np.dot(dx, dx), eps)
 22        vals_L.append(np.linalg.norm(dg) / np.sqrt(n))
 23        vals_mu.append(max(0.0, np.dot(dg, dx) / n))
 24    return max(vals_L), min(vals_mu)
 25
 26
 27def rho_scalar(eta, beta, lam):
 28    # z_{k+1}=(1+beta-eta*lam)z_k-beta*z_{k-1}
 29    roots = np.roots([1.0, -(1.0 + beta - eta*lam), beta])
 30    return float(np.max(np.abs(roots)))
 31
 32
 33def search_certificate(mu, L, beta_grid=None, eta_grid=None):
 34    """Small robust Lyapunov proxy: minimize worst endpoint spectral radius.
 35    The endpoint check is exact for scalar quadratics and is conservative for an
 36    interval. The returned support is the active curvature endpoints."""
 37    mu = max(float(mu), 1e-5)
 38    L = max(float(L), mu * 1.001)
 39    if beta_grid is None: beta_grid = np.linspace(0, .95, 40)
 40    if eta_grid is None: eta_grid = np.linspace(.001, 1.99 / L, 160)
 41    best = None
 42    for beta in beta_grid:
 43        for eta in eta_grid:
 44            r = max(rho_scalar(eta, beta, mu), rho_scalar(eta, beta, L))
 45            if best is None or r < best[0]:
 46                best = (r, eta, beta)
 47    return {"rho": best[0], "eta": best[1], "beta": best[2],
 48            "support": ["mu_endpoint", "L_endpoint"]}
 49
 50
 51def verify_math(seed=0):
 52    rng = np.random.default_rng(seed)
 53    mu, L = .7, 4.0
 54    # Random diagonal quadratics have exact [mu,L] smooth/strong convex bounds.
 55    A = np.diag([mu, 1.5, L])
 56    b = rng.normal(size=3)
 57    f = lambda x: .5*x@A@x + b@x
 58    g = lambda x: A@x+b
 59    xs = [rng.normal(size=3) for _ in range(10)]
 60    gs = [g(x) for x in xs]
 61    fs = [f(x) for x in xs]
 62    residuals = [interpolation_residual(xs[i],gs[i],fs[i],xs[j],gs[j],fs[j],mu,L)
 63                 for i in range(len(xs)) for j in range(len(xs)) if i != j]
 64    # The same expression must be nonnegative; deliberately underspecifying L breaks it.
 65    bad = [interpolation_residual(xs[0],gs[0],fs[0],xs[1],gs[1],fs[1],mu,1.0)
 66           if 1.0 > mu else 0.0]
 67    return {"min_valid_residual": float(min(residuals)),
 68            "max_abs_pair_asymmetry": float(max(abs(residuals[k] - residuals[k]) for k in range(len(residuals)))),
 69            "underspecified_L_residual": float(bad[0])}
 70
 71
 72def run_experiment(seed=7, steps=180):
 73    rng = np.random.default_rng(seed)
 74    # Two parameter blocks: a global step is dictated by the sharp block.
 75    curvatures = [np.array([1., 3., 10.]), np.array([.08, .2, .5])]
 76    xs0 = [rng.normal(size=len(c)) for c in curvatures]
 77    baseline_x = [x.copy() for x in xs0]
 78    idea_x = [x.copy() for x in xs0]
 79    base_eta = 1.0 / max(c.max() for c in curvatures)
 80    base_beta = .9
 81    certs = []
 82    # Warmup secants are obtained at a small stable step, as in the proposal.
 83    for c, x in zip(curvatures, xs0):
 84        histx, histg = [x.copy()], [c*x]
 85        for _ in range(8):
 86            x = x - .05 * (c*x) + rng.normal(0, .002, size=x.shape)
 87            histx.append(x.copy()); histg.append(c*x)
 88        Lhat, mhat = estimate_curvature(histx, histg)
 89        certs.append(search_certificate(mhat, Lhat))
 90    base_losses, idea_losses = [], []
 91    def loss(blocks): return float(sum(.5*np.sum(c*x*x) for c,x in zip(curvatures,blocks)))
 92    spikes_b = spikes_i = 0
 93    for t in range(steps):
 94        oldb = loss(baseline_x); oldi = loss(idea_x)
 95        # standard global momentum SGD
 96        if t == 0: prevb = [x.copy() for x in baseline_x]
 97        newb=[]
 98        for c,x,p in zip(curvatures,baseline_x,prevb):
 99            newb.append(x - base_eta*(c*x) + base_beta*(x-p))
100        prevb, baseline_x = baseline_x, newb
101        # certified per-block momentum update
102        if t == 0: previ = [x.copy() for x in idea_x]
103        newi=[]
104        for c,x,p,cert in zip(curvatures,idea_x,previ,certs):
105            newi.append(x - cert['eta']*(c*x) + cert['beta']*(x-p))
106        previ, idea_x = idea_x, newi
107        lb, li = loss(baseline_x), loss(idea_x)
108        spikes_b += int(lb > oldb * 1.05); spikes_i += int(li > oldi * 1.05)
109        base_losses.append(lb); idea_losses.append(li)
110    return {"baseline": {"initial_loss": loss(xs0), "final_loss": base_losses[-1],
111                          "loss_20pct": base_losses[int(steps*.2)], "eta": base_eta,
112                          "beta": base_beta, "spikes": spikes_b},
113            "idea": {"initial_loss": loss(xs0), "final_loss": idea_losses[-1],
114                     "loss_20pct": idea_losses[int(steps*.2)], "spikes": spikes_i,
115                     "certificates": certs},
116            "curvatures": [c.tolist() for c in curvatures]}
117
118
119if __name__ == '__main__':
120    out = {"math": verify_math(), "experiment": run_experiment()}
121    print(json.dumps(out, indent=2))