import sys, json, time 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, make_report SEEDS = tuple(range(8)) LRS = [0.001, 0.003, 0.006] EPOCHS, BATCH = 12, 128 def pdas(H, c, A, b, active_init=(), tol=1e-8, max_iter=30): n, m = len(c), len(b); active = set(active_init) for it in range(1, max_iter + 1): inds = sorted(active) if inds: K = np.block([[H, -A[inds].T], [A[inds], np.zeros((len(inds), len(inds)))]]) try: sol = np.linalg.solve(K, np.r_[-c, b[inds]]) except np.linalg.LinAlgError: active.remove(inds[-1]); continue v = sol[:n]; mu = np.zeros(m); mu[inds] = sol[n:] else: v = np.linalg.solve(H, -c); mu = np.zeros(m) r = A @ v - b bad = [j for j in range(m) if j not in active and r[j] < -tol] neg = [j for j in active if mu[j] < -tol] if not bad and not neg: return v, mu, tuple(sorted(active)), it if neg: active.remove(min(neg, key=lambda j: mu[j])) else: active.add(min(bad, key=lambda j: r[j])) return v, mu, tuple(sorted(active)), max_iter def seed_model(ds): return make_model('mlp_tiny', tuple(ds['xtr'].shape[1:]), int(ds['out_dim'])) def baseline(ds, lr): torch.manual_seed(2265) return train_model(seed_model(ds), ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None) def idea(ds, lr, alpha=0.20): torch.manual_seed(2265) net = seed_model(ds) # One scalar control per parameter tensor. The linearized barrier is a # parameter-norm trust barrier: ||theta + eta*s*g|| <= cap. caps = [1.02 * p.detach().norm().item() + 1e-6 for p in net.parameters()] opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = nn.MSELoss() x, y = ds['xtr'], ds['ytr'] if x.ndim == 2: x = x.float() if y.ndim == 1: y = y.float().reshape(-1, 1) active = (); its=[]; changes=0; violations=[]; predicted=[]; observed=[] prev_active = () for ep in range(EPOCHS): gen = torch.Generator().manual_seed(2265 + ep) order = torch.randperm(len(x), generator=gen) for start in range(0, len(x), BATCH): idx = order[start:start+BATCH] opt.zero_grad(set_to_none=True) out = net(x[idx]); loss = lossf(out, y[idx]); loss.backward() H = np.eye(len(list(net.parameters()))) * 1.01 c = -np.ones(len(H)); A=[]; b=[]; nominal=[] params = list(net.parameters()) for j,p in enumerate(params): g = p.grad.detach(); th=p.detach(); gn=float(g.norm()) tn=float(th.norm()) # scalar s multiplies Adam's current nominal direction; # linearized cap gives a_j*s >= b_j. direction = -g coeff = -lr * float((th * direction).sum()) / max(tn, 1e-8) h = caps[j] - tn # h(theta + eta*s*(-g)) >= (1-alpha)h(theta) # => coeff*s >= -alpha*h A.append([coeff if k == j else 0.0 for k in range(len(params))]) b.append(-alpha * h) nominal.append(1.0) A=np.asarray(A); b=np.asarray(b); nominal=np.asarray(nominal) # If a nearly-zero gradient makes a true barrier infeasible, use # a harmless broad lower bound; this keeps the small MVP robust. for j in range(len(b)): if abs(A[j,j]) < 1e-10: A[j,j]=1e-10; b[j]=min(b[j], 0.) s, mu, active, it = pdas(H, -nominal, A, b, active) its.append(it); changes += int(active != prev_active); prev_active=active predicted.append(int(np.sum(np.asarray(b) > np.asarray(A).diagonal()))); observed.append(len(active)) for j,p in enumerate(params): p.grad.mul_(float(s[j])) opt.step() # measured post-update barrier violations violations.extend(max(0., float(p.detach().norm())-caps[j]) for j,p in enumerate(params)) net.eval() with torch.no_grad(): metric=float(lossf(net(ds['xte']), ds['yte'].float().reshape(-1,1)).item()) return net, metric, {'iterations':its,'active_changes':changes,'violations':violations,'predicted_active':predicted,'observed_active':observed} def main(): # Correct structural match: optimizer/filter ideas belong to tabular. allres=[]; sweep=[] for lr in LRS: vals=[] for seed in SEEDS: ds=get_dataset('tabular', seed, n_train=400, n_test=200) vals.append(baseline(ds,lr)[1]) sweep.append({'config':{'lr':lr,'epochs':EPOCHS},'mean':float(np.mean(vals)),'per_seed':vals}) best=min(sweep,key=lambda z:z['mean']); idea_grid=[]; idea_diag=[] for lr in LRS: vals=[] for seed in SEEDS: ds=get_dataset('tabular', seed, n_train=400, n_test=200) net,v,d=idea(ds,lr); vals.append(v) if lr==best['config']['lr']: idea_diag.append(d) idea_grid.append({'config':{'lr':lr,'alpha':0.20,'epochs':EPOCHS},'mean':float(np.mean(vals)),'per_seed':vals}) best_i=min(idea_grid,key=lambda z:z['mean']) base_block={'best_cfg':best['config'],'sweep':[{'cfg':z['config'],'mean':z['mean']} for z in sweep], 'full':{'mean':float(np.mean(best['per_seed'])),'std':float(np.std(best['per_seed'])),'per_seed':best['per_seed'],'n':len(best['per_seed'])}} idea_res={'config':best_i['config'],'mean':best_i['mean'],'std':float(np.std(best_i['per_seed'])),'per_seed':best_i['per_seed'],'n':len(best_i['per_seed']), 'grid':idea_grid} pred=np.concatenate([d['predicted_active'] for d in idea_diag]); obs=np.concatenate([d['observed_active'] for d in idea_diag]) sig={'predicted_mean_active':float(pred.mean()),'observed_mean_active':float(obs.mean()),'mean_iterations':float(np.mean([x for d in idea_diag for x in d['iterations']])), 'one_iteration_fraction':float(np.mean([x==1 for d in idea_diag for x in d['iterations']])), 'active_set_changes':int(sum(d['active_changes'] for d in idea_diag)), 'max_violation':float(max([max(d['violations']) for d in idea_diag],default=0.0))} sig['confirmed']=bool(abs(sig['predicted_mean_active']-sig['observed_mean_active']) <= max(0.5,0.25*max(1,sig['observed_mean_active']))) report=make_report('tabular','mlp_tiny',base_block,idea_res,extra=sig) report['baseline_sweep']=sweep; report['idea_grid']=idea_grid; report['search_space_parity']=True with open('bench_report.json','w') as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__=='__main__': main()