import os, sys, json, math, random import numpy as np import torch import torch.nn as nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEED0 = 1505 EPOCHS = 15 BATCH = 128 # This is the complete shared search-space union. The idea is evaluated at all # three learning rates; baseline is also evaluated at all three rates. LR_GRID = [0.0015, 0.003, 0.006] WD_GRID = [0.0, 1e-4] # Fixed a priori stable, moderate asymmetric coupling; nearby settings vary lr. K1, K2 = 0.20, 0.05 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def device_for(): return "cuda" if torch.cuda.is_available() else "cpu" def batches(n, batch, rng): order = rng.permutation(n) for i in range(0, n, batch): yield order[i:i+batch] def baseline_run(cfg, seed, collect=False): seed_all(seed) ds = get_dataset("tabular", seed, n_train=4000, n_test=1000) dev = device_for() try: net = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"]).to(dev) opt = torch.optim.Adam(net.parameters(), lr=float(cfg["lr"]), weight_decay=float(cfg["weight_decay"])) lossf = nn.MSELoss() rng = np.random.default_rng(seed + 991) net.train() for ep in range(EPOCHS): for ix in batches(len(ds["xtr"]), BATCH, rng): x = ds["xtr"][ix].to(dev); y = ds["ytr"][ix].to(dev) opt.zero_grad(set_to_none=True) lossf(net(x), y).backward(); opt.step() net.eval() with torch.no_grad(): val = float(lossf(net(ds["xte"].to(dev)), ds["yte"].to(dev)).cpu()) return val except Exception: # Robust CPU retry, including CUDA/cuDNN allocation failures. if dev != "cuda": raise torch.cuda.empty_cache() os.environ["CUDA_VISIBLE_DEVICES"] = "" return baseline_run(cfg, seed, collect) def idea_run(cfg, seed, collect=False): seed_all(seed) ds = get_dataset("tabular", seed, n_train=4000, n_test=1000) dev = device_for() try: a = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"]).to(dev) b = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"]).to(dev) b.load_state_dict(a.state_dict()) # SGD is the displayed Brownian update; baseline uses standard Adam. pa, pb = list(a.parameters()), list(b.parameters()) lr = float(cfg["lr"]); wd = float(cfg["weight_decay"]) rng1 = np.random.default_rng(seed + 1771); rng2 = np.random.default_rng(seed + 2771) lossf = nn.MSELoss(); history = [] a.train(); b.train() for ep in range(EPOCHS): # Anneal directed coupling and effective temperature to zero. frac = 1.0 - ep / max(1, EPOCHS - 1) k1, k2 = K1 * frac, K2 * frac for ix1, ix2 in zip(batches(len(ds["xtr"]), BATCH, rng1), batches(len(ds["xtr"]), BATCH, rng2)): x1=ds["xtr"][ix1].to(dev); y1=ds["ytr"][ix1].to(dev) x2=ds["xtr"][ix2].to(dev); y2=ds["ytr"][ix2].to(dev) la = lossf(a(x1), y1); lb = lossf(b(x2), y2) ga = torch.autograd.grad(la, pa, create_graph=False) gb = torch.autograd.grad(lb, pb, create_graph=False) with torch.no_grad(): for u,v,gu,gv in zip(pa,pb,ga,gb): # Couplings are applied as parameter forces; a small # shared weight decay is the only common regularizer. du = gu + k1*(u-v) + wd*u dv = gv + k2*(v-u) + wd*v u.add_(du, alpha=-lr); v.add_(dv, alpha=-lr) if collect: with torch.no_grad(): va = torch.cat([p.detach().flatten().cpu() for p in pa]) vb = torch.cat([p.detach().flatten().cpu() for p in pb]) history.append((va, vb)) a.eval(); b.eval() with torch.no_grad(): avg = [(u+v)*0.5 for u,v in zip(pa,pb)] # Evaluate the trained averaged system identically on the task. pred = ds["xte"].to(dev) out = torch.zeros((len(pred),1), device=dev) for xpart, opart in [(None,None)]: # Functional evaluation avoids changing either trained replica. from torch.func import functional_call sd = {n:(p+q)*0.5 for (n,p),(n2,q) in zip(a.named_parameters(), b.named_parameters())} out = functional_call(a, sd, (pred,)) val = float(lossf(out, ds["yte"].to(dev)).cpu()) if collect: return val, history return val except Exception: if dev != "cuda": raise torch.cuda.empty_cache(); os.environ["CUDA_VISIBLE_DEVICES"] = "" return idea_run(cfg, seed, collect) def main(): # Baseline's decisive standard knobs (Adam lr and weight decay) are swept. grid = [{"lr": lr, "weight_decay": wd} for lr in LR_GRID for wd in WD_GRID] base = sweep_baseline(lambda c: (lambda s: baseline_run(c, s)), grid) # Full paired idea sweep at the same learning-rate union; weight decay is # held at the baseline winner's value and coupling is fixed a priori. wd = float(base["best_cfg"]["weight_decay"]) idea_grid = [{"lr": lr, "weight_decay": wd} for lr in LR_GRID] idea_scores = [] for c in idea_grid: r = evaluate(lambda s, c=c: idea_run(c, s), seeds=tuple(range(8))) idea_scores.append({"cfg": c, "result": r}) best = min(idea_scores, key=lambda z: z["result"]["mean"]) # Model-behaviour signature: area/circulation in the two-replica state, # measured from every epoch's trained parameter vectors on seed 0. sig_val, traj = idea_run(best["cfg"], 0, collect=True) if len(traj) > 2: x1=np.array([float(z[0][0]) for z in traj]); x2=np.array([float(z[1][0]) for z in traj]) area=float(np.mean(x1[:-1]*x2[1:]-x2[:-1]*x1[1:])) sep=float(np.mean([torch.linalg.vector_norm(z[0]-z[1]).item() for z in traj])) else: area=0.0; sep=0.0 # The predicted direction is nonzero circulation for k1 != k2. Confirmed # only if it is distinguishable from numerical zero; no claimed task win. signature={"prediction":"nonzero replica circulation when k1!=k2", "k1":K1,"k2":K2,"observed_epoch_area_seed0":area, "observed_mean_replica_separation":sep, "confirmed": bool(abs(area)>1e-12 and sep>1e-8)} report=make_report("tabular","mlp_tiny",base,best["result"],extra=signature) report["idea_sweep"]=idea_scores report["protocol_notes"]={"epochs":EPOCHS,"batch":BATCH,"paired_seeds":list(range(8)), "track_rationale":"tabular is the prescribed structural track for optimizer modifications", "equal_budget_note":"same epochs and minibatch budget per replica; the idea uses two replicas and thus approximately 2x parameter-update compute"} with open("bench_report.json","w") as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__ == "__main__": main()