Pole-radius tuning for gradient tracking / pole_radius_experiment.py
Failed on benchmark
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()