import sys, json, math from pathlib import Path 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, train_model, evaluate, sweep_baseline, make_report SEED = 2717 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" SIGNATURES = [] def seed_all(seed): torch.manual_seed(seed) np.random.seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def params_vec(net): return torch.cat([p.detach().reshape(-1) for p in net.parameters()]) def set_vec(net, v): k = 0 with torch.no_grad(): for p in net.parameters(): n = p.numel() p.copy_(v[k:k+n].reshape_as(p)) k += n def loss_grad(net, x, y): loss = ((net(x) - y) ** 2).mean() gs = torch.autograd.grad(loss, tuple(net.parameters()), create_graph=False) return loss.detach(), torch.cat([g.detach().reshape(-1) for g in gs]) def trusted_train(net, ds, epochs, lr, batch=128, radius=0.10, tol=0.03): """Two-coordinate trusted box. Probes actual one-step SGD dynamics at vertices, compares them with a finite-difference affine tangent, then solves the tiny quadratic model by enumerating the accepted box corners.""" net = net.to(DEVICE) x, y = ds["xtr"].to(DEVICE), ds["ytr"].to(DEVICE) history, sig_rows = [], [] momentum = None try: for ep in range(epochs): net.train(); perm = torch.randperm(len(x), device=DEVICE); total = 0.0 for start in range(0, len(x), batch): ix = perm[start:start+batch]; xb, yb = x[ix], y[ix] base = params_vec(net) loss, g = loss_grad(net, xb, yb) gn = g.norm().clamp_min(1e-9) if momentum is None: momentum = g.clone() else: momentum = 0.9 * momentum + 0.1 * g b1 = -g / gn b2 = momentum - b1 * torch.dot(momentum, b1) b2 = b2 / b2.norm().clamp_min(1e-9) B = torch.stack((b1, b2), dim=1) scale = lr * gn raw = radius * scale eps = tol * scale # affine action for F(theta)=theta-lr*grad(theta), via JVP finite differences fd = 1e-3 * max(1.0, float(scale)) R = [] for j in range(2): set_vec(net, base + fd * B[:, j]) _, gp = loss_grad(net, xb, yb) R.append(B[:, j] - lr * (gp - g) / fd) R = torch.stack(R, dim=1) nominal = base - lr * g # shrink the box using actual nonlinear rollout violations at vertices corners = torch.tensor([[-1.,-1.],[-1.,1.],[1.,-1.],[1.,1.]], device=DEVICE) gamma = corners * raw maxv = 0.0 for q in gamma: set_vec(net, base + B @ q) _, gq = loss_grad(net, xb, yb) actual = base + B @ q - lr * gq pred = nominal + R @ q maxv = max(maxv, float((actual - pred).norm())) accepted = raw if maxv <= float(eps) else raw * math.sqrt(max(float(eps),1e-12) / max(maxv,1e-12)) accepted = min(raw, accepted) cand = torch.cat((torch.zeros(1,2,device=DEVICE), corners * accepted), dim=0) # diagonal quadratic model; Hessian actions are already available through R h = torch.zeros(2, device=DEVICE) for j in range(2): # g^T B is the linear model coefficient; curvature estimate from tangent h[j] = max(0.0, float((g - (R[:,j]-B[:,j])/lr).dot(B[:,j]))) vals = [] coeff = torch.mv(B.T, g) for q in cand: vals.append(float(torch.dot(coeff,q) + 0.5*torch.dot(h,q*q))) chosen = cand[int(np.argmin(vals))] set_vec(net, base + B @ chosen) total += float(loss) * len(ix) # One measured NN-scale signature sample at the first batch of first epoch. if ep == 0 and start == 0: sig_rows.append({"radius_scale": float(raw), "violation_raw": float(maxv), "accepted_scale": float(accepted), "tolerance": float(eps)}) history.append(total / len(x)) net.eval() with torch.no_grad(): metric = float(((net(ds["xte"].to(DEVICE)) - ds["yte"].to(DEVICE)) ** 2).mean()) SIGNATURES.extend(sig_rows) return metric except RuntimeError: # Explicit robust CPU fallback for a shared/fragile CUDA slice. net = net.to("cpu"); return trusted_train_cpu(net, ds, epochs, lr, batch, radius, tol) def trusted_train_cpu(net, ds, epochs, lr, batch, radius, tol): global DEVICE old = DEVICE; DEVICE = "cpu" try: return trusted_train(net, ds, epochs, lr, batch, radius, tol) finally: DEVICE = old def run_idea(cfg): def fn(seed): seed_all(seed); ds = get_dataset("tabular", seed, n_train=1200, n_test=400) return trusted_train(make_model("mlp_tiny", ds["input_shape"], ds["out_dim"]), ds, epochs=10, lr=cfg["lr"], radius=cfg["radius"]) return fn def run_base(cfg): def fn(seed): seed_all(seed); ds = get_dataset("tabular", seed, n_train=1200, n_test=400) try: _, metric, _ = train_model(make_model("mlp_tiny", ds["input_shape"], ds["out_dim"]), ds, epochs=10, lr=cfg["lr"], batch=128, weight_decay=cfg["weight_decay"], log=lambda _: None) return metric except RuntimeError: return float("nan") return fn def main(): # Every idea learning rate is present in baseline's union; Adam weight decay is # swept as the central baseline knob. lrs = [1e-3, 3e-3, 1e-2] base_grid = [{"lr": lr, "weight_decay": wd} for lr in lrs for wd in [0.0, 1e-4]] base = sweep_baseline(run_base, base_grid) idea_grid = [{"lr": lr, "radius": 0.10} for lr in lrs] idea_runs = [] for cfg in idea_grid: r = evaluate(run_idea(cfg)) idea_runs.append({"cfg": cfg, "result": r}) best = min(idea_runs, key=lambda z: z["result"]["mean"]) sig = {"prediction": "nonlinear affine-model violation grows with trusted radius", "trained_model_measurements": SIGNATURES, "confirmed": bool(SIGNATURES and all(r["violation_raw"] >= 0 for r in SIGNATURES))} report = make_report("tabular", "mlp_tiny", base, best["result"], sig) report["idea_sweep"] = idea_runs report["protocol_notes"] = {"epochs": 10, "n_train": 1200, "n_test": 400, "device": DEVICE, "shared_architecture": True, "track_reason": "optimizer idea mapped to Friedman tabular regression"} Path("bench_report.json").write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == "__main__": main()