import sys, json, math, copy, random from pathlib import Path import numpy as np import torch sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) WORKERS = 4 EPOCHS = 24 CANDIDATES = [0.001, 0.003, 0.01] # Symmetric stochastic ring: each worker keeps half its value and averages neighbors. def ring_matrix(n=4): W = np.zeros((n, n), dtype=float) for i in range(n): W[i, i] = 0.5 W[i, (i-1) % n] += 0.25 W[i, (i+1) % n] += 0.25 return W W = ring_matrix() LAMS = np.linalg.eigvalsh(W)[0:-1] # exclude consensus eigenvalue 1 def roots(lam, q, alpha): return np.roots([1.0, q * alpha - 2.0 * lam, lam * lam - q * alpha]) def pole_radius(alpha, lams=LAMS, qs=np.linspace(0.1, 10.0, 17)): return float(max(abs(r) for lam in lams for q in qs for r in roots(lam, q, alpha))) def flat(tensors): return torch.cat([x.detach().reshape(-1) for x in tensors]) def grad_at(model, x, y): model.zero_grad(set_to_none=True) loss = ((model(x) - y) ** 2).mean() loss.backward() return [p.grad.detach().clone() for p in model.parameters()], float(loss.detach()) def curvature_interval(models, xs, ys): # Empirical gradient-difference Rayleigh quotients, as specified in the idea. vals = [] for m, x, y in zip(models, xs, ys): base = [p.detach().clone() for p in m.parameters()] g0, _ = grad_at(m, x, y) torch.manual_seed(9137) delta = [0.01 * torch.randn_like(p) for p in m.parameters()] with torch.no_grad(): for p, d in zip(m.parameters(), delta): p.add_(d) g1, _ = grad_at(m, x, y) with torch.no_grad(): for p, b in zip(m.parameters(), base): p.copy_(b) dx = flat(delta); dg = flat([a-b for a,b in zip(g1,g0)]) q = float(torch.dot(dg, dx) / (torch.dot(dx, dx) + 1e-12)) if np.isfinite(q) and q > 1e-6: vals.append(q) if not vals: return 0.1, 10.0 return max(0.05, float(np.quantile(vals, .2))), max(0.1, float(np.quantile(vals, .8))) def choose_alpha(models, xs, ys): mu, hi = curvature_interval(models, xs, ys) qs = np.linspace(mu, hi, 17) scores = {a: pole_radius(a, LAMS, qs) for a in CANDIDATES} alpha = min(CANDIDATES, key=lambda a: scores[a]) return alpha, mu, hi, scores def train_diging(seed, alpha, tune=False, return_signature=False): torch.manual_seed(seed); np.random.seed(seed); random.seed(seed) d = get_dataset('tabular', seed, n_train=400, n_test=200) # Same architecture and same partition for both systems. xparts = list(torch.chunk(d['xtr'], WORKERS)) yparts = list(torch.chunk(d['ytr'], WORKERS)) models = [make_model('mlp_tiny', d['input_shape'], d['out_dim']) for _ in range(WORKERS)] xs = xparts; ys = yparts grads = []; losses = [] for m,x,y in zip(models,xs,ys): g, loss = grad_at(m,x,y); grads.append(g); losses.append(loss) selected = alpha; mu=hi=None; scores=None if tune: selected, mu, hi, scores = choose_alpha(models, xs, ys) # y is the gradient tracker, represented as per-model parameter lists. tracker = [[g.clone() for g in gg] for gg in grads] obs_dis = [] for _ in range(EPOCHS): old_params = [[p.detach().clone() for p in m.parameters()] for m in models] old_tracker = [[z.detach().clone() for z in tt] for tt in tracker] with torch.no_grad(): for i,m in enumerate(models): for k,p in enumerate(m.parameters()): comm = sum(W[i,j] * old_params[j][k] for j in range(WORKERS)) p.copy_(comm - selected * old_tracker[i][k]) new_grads=[] for m,x,y in zip(models,xs,ys): g, loss = grad_at(m,x,y); new_grads.append(g); losses.append(loss) with torch.no_grad(): for i in range(WORKERS): for k in range(len(tracker[i])): comm = sum(W[i,j] * old_tracker[j][k] for j in range(WORKERS)) tracker[i][k].copy_(comm + new_grads[i][k] - grads[i][k]) grads = new_grads tv = torch.stack([flat(t) for t in tracker]) obs_dis.append(float(torch.linalg.norm(tv - tv.mean(0,keepdim=True)))) with torch.no_grad(): pred = torch.stack([m(d['xte']) for m in models]).mean(0) metric = float(((pred-d['yte'])**2).mean()) sig = None if return_signature: ratios = [obs_dis[i+1]/(obs_dis[i]+1e-12) for i in range(len(obs_dis)-1)] observed = float(np.median(ratios[-8:])) if ratios else float('nan') predicted = pole_radius(selected, LAMS, np.linspace(mu or .1, hi or 10., 17)) sig = {'alpha': selected, 'mu_est': mu, 'L_est': hi, 'predicted_pole_radius': predicted, 'observed_tracker_disagreement_ratio': observed, 'nondesignated_training_model': 'mlp_tiny_worker_ensemble', 'confirmed': bool(np.isfinite(observed) and abs(observed-predicted) <= max(.08, .25*predicted))} return metric, sig def make_baseline(cfg): return lambda seed: train_diging(seed, float(cfg['lr']), tune=False)[0] def make_idea(cfg): # cfg lr is the shared candidate grid; tuning chooses the minimax candidate return lambda seed: train_diging(seed, float(cfg['lr']), tune=True)[0] if __name__ == '__main__': grid = [{'lr': a} for a in CANDIDATES] base = sweep_baseline(make_baseline, grid, seeds=(0,1,2,3)) # Evaluate every idea candidate on all eight paired seeds; select best full mean. idea_runs=[] for cfg in grid: r=evaluate(make_idea(cfg), seeds=SEEDS) idea_runs.append((cfg,r)) idea_cfg, idea = min(idea_runs, key=lambda z:z[1]['mean']) # Signature is measured on trained systems, not on the analytic toy alone. _, signature = train_diging(0, idea_cfg['lr'], tune=True, return_signature=True) report = make_report('tabular', 'mlp_tiny', base, idea, {'track_structure':'optimizer/decentralized regression', **signature, 'pole_spectrum': LAMS.tolist(), 'idea_grid': grid, 'selected_idea_cfg': idea_cfg, 'candidate_full_means': {str(c['lr']): r['mean'] for c,r in idea_runs}}) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2))