Pole-radius tuning for gradient tracking / pole_radius_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6SEED = 1324
  7rng = np.random.default_rng(SEED)
  8
  9
 10def roots(lam, q, alpha):
 11    # (r-lambda)^2 + q alpha (r-1) = 0
 12    return np.roots([1.0, q * alpha - 2.0 * lam, lam * lam - q * alpha])
 13
 14
 15def pole_radius(lams, qs, alpha):
 16    return float(max(abs(r) for lam in lams for q in qs for r in roots(lam, q, alpha)))
 17
 18
 19def predicted_boundary(lam, q):
 20    # Jury boundary for 0 <= lambda < 1: q alpha = (1+lambda)^2/2.
 21    return (1.0 + lam) ** 2 / (2.0 * q)
 22
 23
 24def make_symmetric_W(eigs):
 25    n = len(eigs) + 1
 26    # Deterministic orthonormal basis with first vector equal to consensus.
 27    A = rng.normal(size=(n, n))
 28    A[:, 0] = 1.0
 29    Q, _ = np.linalg.qr(A)
 30    if np.sum(Q[:, 0]) < 0:
 31        Q[:, 0] *= -1
 32    vals = np.array([1.0] + list(eigs), dtype=float)
 33    W = Q @ np.diag(vals) @ Q.T
 34    # Numerical cleanup preserves symmetry and stochasticity closely.
 35    W = (W + W.T) / 2.0
 36    return W
 37
 38
 39def diging(W, qs_worker, centers, alpha, rounds=250):
 40    n, d = centers.shape
 41    x = np.zeros((n, d))
 42    g = qs_worker[:, None] * (x - centers)
 43    y = g.copy()
 44    losses = []
 45    tracker = []
 46    for _ in range(rounds):
 47        x = W @ x - alpha * y
 48        g_new = qs_worker[:, None] * (x - centers)
 49        y = W @ y + g_new - g
 50        g = g_new
 51        local_loss = 0.5 * qs_worker[:, None] * (x - centers) ** 2
 52        losses.append(float(np.mean(np.sum(local_loss, axis=1))))
 53        tracker.append(float(np.linalg.norm(y - np.mean(y, axis=0, keepdims=True))))
 54    return np.asarray(losses), np.asarray(tracker), x
 55
 56
 57def main():
 58    # Core numerical sanity check: formula roots match polynomial evaluation.
 59    checks = []
 60    for lam, q, a in [(0.2, 1.0, .1), (.75, 4.0, .03), (-.3, 2.0, .2)]:
 61        rr = roots(lam, q, a)
 62        residual = max(abs((r-lam)**2 + q*a*(r-1)) for r in rr)
 63        checks.append(float(residual))
 64
 65    # Prediction 1: stability transition at alpha_c=(1+lambda)^2/(2q).
 66    lam0, q0 = .6, 3.0
 67    ac = predicted_boundary(lam0, q0)
 68    ratios = np.array([.98, .995, 1.005, 1.02])
 69    boundary_rows = []
 70    for ratio in ratios:
 71        rho = pole_radius([lam0], [q0], ac * ratio)
 72        boundary_rows.append({"ratio": float(ratio), "alpha": float(ac*ratio),
 73                              "rho": rho, "stable_observed": bool(rho < 1.0)})
 74
 75    # Prediction 2: boundary scales inversely with curvature q.
 76    lam1 = .6
 77    qvals = np.array([1., 2., 4., 8.])
 78    measured = []
 79    for q in qvals:
 80        grid = np.linspace(.0001, 1.2 * predicted_boundary(lam1, q), 3000)
 81        rhos = np.array([pole_radius([lam1], [q], a) for a in grid])
 82        # first grid point at/above one is a numerical estimate of the transition
 83        idx = np.flatnonzero(rhos >= 1.0)
 84        observed = float(grid[idx[0]]) if len(idx) else float("nan")
 85        measured.append({"q": float(q), "predicted_alpha_c": predicted_boundary(lam1, q),
 86                         "observed_alpha_c": observed,
 87                         "relative_error": abs(observed-predicted_boundary(lam1,q))/predicted_boundary(lam1,q)})
 88
 89    # Prediction 3: minimax tuning picks a compromise over the lambda/q rectangle.
 90    lams = np.array([.2, .6, .85])
 91    qs = np.linspace(1., 5., 17)
 92    candidates = np.linspace(.0005, .30, 1200)
 93    candidate_rho = np.array([pole_radius(lams, qs, a) for a in candidates])
 94    best_i = int(np.argmin(candidate_rho))
 95    best_a = float(candidates[best_i])
 96    best_rho = float(candidate_rho[best_i])
 97    # A deliberately generic conservative choice, and its predicted contraction.
 98    fixed_a = .05
 99    fixed_rho = pole_radius(lams, qs, fixed_a)
100
101    # Small decentralized quadratic experiment, same setup and rounds.
102    W = make_symmetric_W(lams)
103    worker_q = np.array([1., 2., 3., 5.])
104    centers = rng.normal(scale=1.0, size=(4, 2))
105    tuned_loss, tuned_tracker, tuned_x = diging(W, worker_q, centers, best_a)
106    fixed_loss, fixed_tracker, fixed_x = diging(W, worker_q, centers, fixed_a)
107    # Centralized optimum for reporting a meaningful parameter error.
108    x_star = np.sum(worker_q[:, None] * centers, axis=0) / np.sum(worker_q)
109    tuned_err = float(np.linalg.norm(np.mean(tuned_x, axis=0) - x_star))
110    fixed_err = float(np.linalg.norm(np.mean(fixed_x, axis=0) - x_star))
111
112    report = {
113        "seed": SEED,
114        "root_residual_max": max(checks),
115        "prediction_1_boundary": {"lambda": lam0, "q": q0, "predicted_alpha_c": ac,
116                                   "rows": boundary_rows},
117        "prediction_2_inverse_curvature": measured,
118        "prediction_3_minimax": {"lambdas": lams.tolist(), "q_interval": [1., 5.],
119                                  "tuned_alpha": best_a, "tuned_predicted_rho": best_rho,
120                                  "fixed_alpha": fixed_a, "fixed_predicted_rho": fixed_rho,
121                                  "predicted_improvement": fixed_rho-best_rho},
122        "toy_diging": {"rounds": 250, "tuned_alpha": best_a, "fixed_alpha": fixed_a,
123                        "initial_loss": float(fixed_loss[0]),
124                        "final_loss_fixed": float(fixed_loss[-1]),
125                        "final_loss_tuned": float(tuned_loss[-1]),
126                        "final_tracker_fixed": float(fixed_tracker[-1]),
127                        "final_tracker_tuned": float(tuned_tracker[-1]),
128                        "final_mean_error_fixed": fixed_err,
129                        "final_mean_error_tuned": tuned_err},
130    }
131    Path("results.json").write_text(json.dumps(report, indent=2))
132    print(json.dumps(report, indent=2))
133
134
135if __name__ == "__main__":
136    main()