import sys, json, random from pathlib import Path import numpy as np import torch from scipy.optimize import linprog sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report TRACK, MODEL = 'tabular', 'mlp_tiny' SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS, BATCH = 12, 128 # The union is shared by baseline and idea. Fixed ratios are the standard # alternatives to equal weighting and expose the method's replacement knob. LRS = [1e-3, 3e-3, 1e-2] RATIOS = [[1., 1., 1.], [10., 3., 1.], [100., 10., 1.]] GRID = [{'lr': lr, 'ratio': ratio} for lr in LRS for ratio in RATIOS] 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 certificate(q, S, epsilon=0.02, wmax=30.): q, S = np.asarray(q, float), np.asarray(S, float) d, tiers = len(q), S.shape[1] # variables: w_1..w_L, delta; maximize delta A, b = [], [] for i in range(tiers): row = np.zeros(tiers + 1); row[i] = -1; row[-1] = 1 A.append(row); b.append(0.) for k in range(d): row = np.zeros(tiers + 1); row[:tiers] = S[k] A.append(row); b.append(epsilon - q[k]) row = np.zeros(tiers + 1); row[:tiers] = -S[k] A.append(row); b.append(epsilon + q[k]) r = linprog(np.r_[np.zeros(tiers), -1.], A_ub=np.asarray(A), b_ub=np.asarray(b), bounds=[(0., wmax)] * tiers + [(0., wmax)], method='highs') if not r.success: return np.ones(tiers), 0., float('inf'), False w = r.x[:tiers] return w, float(r.x[-1]), float(np.max(np.abs(q + S @ w))), True def tier_values(pred, target, scale): # Three ordered accuracy requirements: tier 1 is the strictest. err = (pred - target).abs() / scale return [torch.relu(err - t).mean() for t in (0.25, 0.50, 0.80)] def train_idea(model, ds, lr, ratio, epsilon=0.02, refresh=4): device = 'cuda' if torch.cuda.is_available() else 'cpu' try: net = model.to(device); x, y = ds['xtr'].to(device), ds['ytr'].to(device) scale = y.std().clamp_min(1e-4) opt = torch.optim.Adam(net.parameters(), lr=lr) weights = torch.tensor(ratio, device=device); stats = [] for ep in range(EPOCHS): net.train(); perm = torch.randperm(len(x), device=device) for start in range(0, len(x), BATCH): idx = perm[start:start+BATCH]; pred = net(x[idx]) tiers = tier_values(pred, y[idx], scale) perf = ((pred-y[idx])**2).mean() if (ep * ((len(x)+BATCH-1)//BATCH) + start//BATCH) % refresh == 0: params = [p for p in net.parameters() if p.requires_grad] def flat_grads(loss): gs = torch.autograd.grad(loss, params, retain_graph=True, allow_unused=True) return torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1) for p,g in zip(params,gs)]) q = flat_grads(perf).detach().cpu().numpy() S = np.stack([flat_grads(v).detach().cpu().numpy() for v in tiers], axis=1) nw, delta, residual, ok = certificate(q, S, epsilon=epsilon, wmax=30.) if ok and delta > 1e-5: weights = torch.tensor(nw, dtype=torch.float32, device=device) stats.append((delta, residual, float(weights.min()))) loss = perf + sum(w*v for w,v in zip(weights, tiers)) opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): out = net(ds['xte'].to(device)); tv = tier_values(out, ds['yte'].to(device), scale) metric = float(((out-ds['yte'].to(device))**2).mean()) violations = [float(v) for v in tv] return net, metric, {'violations': violations, 'cert_stats': stats} except RuntimeError: if torch.cuda.is_available(): torch.cuda.empty_cache() return train_idea_cpu(model.to('cpu'), ds, lr, ratio, epsilon, refresh) def train_idea_cpu(model, ds, lr, ratio, epsilon, refresh): # Same intervention and hyperparameters, explicit CPU retry after CUDA errors. x,y=ds['xtr'],ds['ytr']; scale=y.std().clamp_min(1e-4); net=model opt=torch.optim.Adam(net.parameters(),lr=lr); weights=torch.tensor(ratio); stats=[] for ep in range(EPOCHS): perm=torch.randperm(len(x)) for st in range(0,len(x),BATCH): idx=perm[st:st+BATCH]; pred=net(x[idx]); tiers=tier_values(pred,y[idx],scale); perf=((pred-y[idx])**2).mean() if (ep*((len(x)+BATCH-1)//BATCH)+st//BATCH)%refresh==0: ps=[p for p in net.parameters() if p.requires_grad] def fg(z): gs=torch.autograd.grad(z,ps,retain_graph=True,allow_unused=True) return torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1) for p,g in zip(ps,gs)]) q=fg(perf).detach().numpy(); S=np.stack([fg(v).detach().numpy() for v in tiers],1) nw,de,re,ok=certificate(q,S,epsilon,30.) if ok and de>1e-5: weights=torch.tensor(nw,dtype=torch.float32) stats.append((de,re,float(weights.min()))) opt.zero_grad(); (perf+sum(w*v for w,v in zip(weights,tiers))).backward(); opt.step() with torch.no_grad(): out=net(ds['xte']); vals=tier_values(out,ds['yte'],scale) return net,float(((out-ds['yte'])**2).mean()),{'violations':[float(v) for v in vals],'cert_stats':stats} def baseline_fn(cfg): def run(seed): seed_all(seed); ds=get_dataset(TRACK, seed, n_train=400, n_test=200) # Standard weighted multi-objective loss, implemented locally so tier # objectives and architecture are identical to the certificate system. net=make_model(MODEL,ds['input_shape'],ds['out_dim']); device='cuda' if torch.cuda.is_available() else 'cpu' try: net=net.to(device); x,y=ds['xtr'].to(device),ds['ytr'].to(device); scale=y.std().clamp_min(1e-4) opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']); ws=torch.tensor(cfg['ratio'],device=device) for _ in range(EPOCHS): for st in range(0,len(x),BATCH): ix=torch.randperm(len(x),device=device)[st:st+BATCH]; p=net(x[ix]); tv=tier_values(p,y[ix],scale); loss=((p-y[ix])**2).mean()+sum(w*v for w,v in zip(ws,tv)); opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): return float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean()) except RuntimeError: return float('inf') return run def idea_fn(cfg): def run(seed): seed_all(seed); ds=get_dataset(TRACK,seed,n_train=400,n_test=200) _,m,_=train_idea(make_model(MODEL,ds['input_shape'],ds['out_dim']),ds,cfg['lr'],cfg['ratio']) return m return run def signature(cfg): seed=0; seed_all(seed); ds=get_dataset(TRACK,seed,n_train=400,n_test=200) net,m,info=train_idea(make_model(MODEL,ds['input_shape'],ds['out_dim']),ds,cfg['lr'],cfg['ratio']) vals=np.asarray(info['cert_stats'],float); finite=vals[np.isfinite(vals[:,1])] if len(vals) else np.empty((0,3)) observed=float(np.mean(finite[:,1])) if len(finite) else float('inf') predicted=float(0.02); return {'quantity':'parameter-gradient infinity stationarity residual at refreshes','predicted_epsilon':predicted,'observed_mean_residual':observed,'positive_margin_frequency':float(np.mean(vals[:,0]>1e-5)) if len(vals) else 0.,'confirmed':bool(observed <= predicted*1.5)} def main(): # Cheap analytical LP sanity check before NN training. w,d,r,ok=certificate(np.array([-3.,-1.]),np.eye(2),epsilon=1e-9,wmax=100.) math_check={'predicted_weights':[3.,1.],'observed_weights':w.tolist(),'residual':r,'passed':bool(np.max(np.abs(w-[3.,1.]))<1e-5)} base=sweep_baseline(baseline_fn,GRID,seeds=SWEEP_SEEDS) best_cfg=base['best_cfg']; idea_trials=[] for cfg in GRID: rr=evaluate(idea_fn(cfg),seeds=SEEDS); idea_trials.append({'cfg':cfg,'result':rr}) best=min(idea_trials,key=lambda z:z['result']['mean']) extra={'mechanism_signature':signature(best['cfg']),'math_check':math_check,'idea_sweep':idea_trials,'protocol_notes':'Matched tabular Friedman regression and mlp_tiny; baseline and controller share all lr and ratio configurations, Adam, epochs, batch, and data seeds.'} report=make_report(TRACK,MODEL,base,best['result'],extra) Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2)) if __name__=='__main__': main()