import sys, json, math, random from pathlib import Path import numpy as np import torch from torch import nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report TRACK = "tabular" MODEL = "mlp_tiny" EPOCHS = 18 BATCH = 128 WEIGHT_DECAY = 0.0 # Shared union: baseline and idea both run every learning rate below. LRS = [1e-3, 3e-3, 1e-2] IDEA_SETTINGS = [ {"lr": 1e-3, "r0": 0.02}, {"lr": 3e-3, "r0": 0.02}, {"lr": 1e-2, "r0": 0.02}, ] SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) 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 global_norm(xs): return math.sqrt(sum(float(x.detach().pow(2).sum().cpu()) for x in xs)) def train_baseline(cfg, seed, return_model=False): seed_all(seed) ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=400) net = make_model(MODEL, ds["input_shape"], ds["out_dim"]) device = "cuda" if torch.cuda.is_available() else "cpu" try: net = net.to(device) opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"], weight_decay=WEIGHT_DECAY) lossf = nn.MSELoss() x, y = ds["xtr"].to(device), ds["ytr"].to(device) for _ in range(EPOCHS): net.train(); perm = torch.randperm(len(x), device=device) for i in range(0, len(x), BATCH): ix = perm[i:i+BATCH]; loss = lossf(net(x[ix]), y[ix]) opt.zero_grad(set_to_none=True); loss.backward(); opt.step() net.eval() with torch.no_grad(): metric = float(lossf(net(ds["xte"].to(device)), ds["yte"].to(device)).cpu()) return (metric, net, ds, device) if return_model else metric except (RuntimeError, torch.cuda.CudaError): if device == "cuda": seed_all(seed); net = make_model(MODEL, ds["input_shape"], ds["out_dim"]) return train_baseline(cfg, seed, return_model) if False else _train_cpu(net, ds, cfg, return_model) raise def _train_cpu(net, ds, cfg, return_model=False): net = net.cpu(); opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"], weight_decay=WEIGHT_DECAY) lossf = nn.MSELoss(); x, y = ds["xtr"], ds["ytr"] for _ in range(EPOCHS): perm = torch.randperm(len(x)) for i in range(0, len(x), BATCH): ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]), y[ix]); opt.zero_grad(set_to_none=True); loss.backward(); opt.step() with torch.no_grad(): metric=float(lossf(net(ds["xte"]), ds["yte"])) return (metric, net, ds, "cpu") if return_model else metric def train_idea(cfg, seed, return_model=False): seed_all(seed) ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=400) net = make_model(MODEL, ds["input_shape"], ds["out_dim"]) device = "cuda" if torch.cuda.is_available() else "cpu" try: net=net.to(device); opt=torch.optim.Adam(net.parameters(), lr=cfg["lr"], weight_decay=WEIGHT_DECAY) lossf=nn.MSELoss(); x,y=ds["xtr"].to(device),ds["ytr"].to(device) # Fixed cap is deliberately identical to initial adaptive radius. r=float(cfg["r0"]); rmin=.1*cfg["r0"]; rmax=10*cfg["r0"] qs=[]; radii=[]; scales=[]; proposal_norms=[] for _ in range(EPOCHS): net.train(); perm=torch.randperm(len(x),device=device) for i in range(0,len(x),BATCH): ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(set_to_none=True); loss.backward() # Adam's actual proposal is reconstructed from its state after a harmless step. # We instead use the standard Adam preconditioned proposal explicitly. props=[] beta1,beta2=opt.param_groups[0]["betas"]; eps=opt.param_groups[0]["eps"]; step=int(opt.state[next(net.parameters())].get("step",0))+1 for p in net.parameters(): if p.grad is None: continue st=opt.state[p] if not st: st["step"]=torch.tensor(0.,device=p.device); st["exp_avg"]=torch.zeros_like(p); st["exp_avg_sq"]=torch.zeros_like(p) m=st["exp_avg"]; v=st["exp_avg_sq"]; m.mul_(beta1).add_(p.grad,alpha=1-beta1); v.mul_(beta2).addcmul_(p.grad,p.grad,value=1-beta2) bc1=1-beta1**step; bc2=1-beta2**step props.append(-cfg["lr"]*(m/bc1)/(v.sqrt()/math.sqrt(bc2)+eps)) pn=global_norm(props); q=min(1.,pn/(r+1e-12)); scale=min(1.,r/(pn+1e-12)) with torch.no_grad(): for p,prop in zip([p for p in net.parameters() if p.grad is not None],props): p.add_(prop,alpha=scale) r=float(np.clip(r*math.exp(.02*q-.01*(1-q)),rmin,rmax)); qs.append(q); radii.append(r); scales.append(scale); proposal_norms.append(pn) net.eval(); with torch.no_grad(): metric=float(lossf(net(ds["xte"].to(device)),ds["yte"].to(device)).cpu()) if return_model: return metric,net,ds,device,{"mean_q":float(np.mean(qs)),"clip_fraction":float(np.mean(np.array(qs)>=1-1e-9)),"radius_final":r,"radius_initial":cfg["r0"],"mean_scale":float(np.mean(scales)),"mean_proposal_norm":float(np.mean(proposal_norms))} return metric except (RuntimeError, torch.cuda.CudaError): raise RuntimeError("CUDA failure; rerun with CUDA_VISIBLE_DEVICES='' for CPU fallback") def baseline_factory(cfg): return lambda seed: train_baseline(cfg, seed) def idea_factory(cfg): return lambda seed: train_idea(cfg, seed) def main(): # Official baseline sweep, then full-seed baseline at selected cfg. base=sweep_baseline(baseline_factory, [{"lr":lr} for lr in LRS], seeds=SWEEP_SEEDS) idea_cfgs=IDEA_SETTINGS idea_sweep=[] for cfg in idea_cfgs: idea_sweep.append({"cfg":cfg,"mean":evaluate(idea_factory(cfg),SWEEP_SEEDS)["mean"]}) best_idea_cfg=min(idea_sweep,key=lambda z:z["mean"])["cfg"] base_full=base["full"] idea_full=evaluate(idea_factory(best_idea_cfg),SEEDS) # Trained-model signature: inspect one paired seed's actual proposals/updates. _,_,_,_,sig=train_idea(best_idea_cfg,0,True) sig.update({"prediction":"sustained saturation expands radius and later unsaturation contracts it","observed_radius_increase":sig["radius_final"]>sig["radius_initial"],"observed_clipping_fraction":sig["clip_fraction"],"confirmed":bool(sig["radius_final"]>sig["radius_initial"] and sig["clip_fraction"]>0)}) report=make_report(TRACK,MODEL,{"best_cfg":base["best_cfg"],"sweep":base["sweep"],"full":base_full},idea_full,{"idea_sweep":idea_sweep,"mechanism_signature":sig}) Path("bench_report.json").write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=="__main__": main()