import os, 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, evaluate, sweep_baseline, make_report from bench.protocol import DEFAULT_SEEDS, SWEEP_SEEDS E = 16 def seed_all(seed): np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def active_set_simplex(Z, b, lam=1e-3, tol=1e-9, max_pivots=100): """Exact-ish active-set solution of min .5*x'Ax-b'x, x>=0, 1'x=1.""" Z = np.asarray(Z, dtype=np.float64); b = np.asarray(b, dtype=np.float64) e = Z.shape[1]; A = Z.T @ Z + lam*np.eye(e); F = list(range(e)); pivots = 0 while pivots <= max_pivots: ii = np.asarray(F, dtype=int); Af = A[np.ix_(ii, ii)] K = np.zeros((len(ii)+1, len(ii)+1)); K[:-1,:-1] = Af K[:-1,-1] = 1.; K[-1,:-1] = 1. rhs = np.r_[b[ii], 1.] sol = np.linalg.solve(K, rhs); xf, nu = sol[:-1], sol[-1] x = np.zeros(e); x[ii] = xf bad = [(v, i) for v,i in zip(xf,ii) if v < -tol] if bad: m = min(v for v,_ in bad); i = min(i for v,i in bad if v <= m+1e-12) F.remove(i); pivots += 1; continue g = A @ x - b + nu blocked = [i for i in range(e) if i not in F and g[i] < -tol] if blocked: m = min(g[i] for i in blocked); i = min(i for i in blocked if g[i] <= m+1e-12) F.append(i); F.sort(); pivots += 1; continue x[np.abs(x) < tol] = 0. return x, {'pivots': pivots, 'support': int((x > tol).sum()), 'simplex_error': float(abs(x.sum()-1)), 'min_x': float(x.min()), 'kkt': float(max(np.max(np.abs(g[x>tol])) if np.any(x>tol) else 0., max(0., -np.min(g[x<=tol])) if np.any(x<=tol) else 0.))} raise RuntimeError('active-set pivot limit') class RouterMLP(nn.Module): def __init__(self, d): super().__init__() self.trunk = nn.Sequential(nn.Linear(d,64), nn.ReLU(), nn.Linear(64,64), nn.ReLU()) self.experts = nn.Linear(64,E) self.logits = nn.Parameter(torch.zeros(E)) def expert_values(self,x): return self.experts(self.trunk(x)) def forward(self,x,weights=None): z=self.expert_values(x) if weights is None: weights=torch.softmax(self.logits,dim=0) return z @ weights, z def fit(seed, cfg, idea=False, collect=False): seed_all(seed) d=get_dataset('tabular', seed, n_train=400, n_test=400) xtr,ytr=d['xtr'],d['ytr']; xte,yte=d['xte'],d['yte'] dev='cuda' if torch.cuda.is_available() else 'cpu' try: net=RouterMLP(int(np.prod(d['input_shape']))).to(dev) xtr=xtr.to(dev).float().reshape(len(xtr),-1); ytr=ytr.to(dev).float().reshape(-1) xte=xte.to(dev).float().reshape(len(xte),-1); yte=yte.to(dev).float().reshape(-1) opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'], weight_decay=cfg.get('weight_decay', 0.0)) weights=torch.full((E,),1./E,device=dev) epochs=cfg['epochs']; bs=cfg['batch']; stats=[] for ep in range(epochs): order=torch.randperm(len(xtr),device=dev) for start in range(0,len(xtr),bs): ix=order[start:start+bs]; pred,z=net(xtr[ix], None if not idea else weights) loss=((pred-ytr[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step() if idea: with torch.no_grad(): z=net.expert_values(xtr).detach().cpu().numpy(); yy=ytr.detach().cpu().numpy() # Ridge least-squares router calibration, then simplex projection via active set. weights_np, met=active_set_simplex(z, z.T@yy, lam=cfg['lam']) weights=torch.tensor(weights_np,dtype=torch.float32,device=dev) stats.append(met) with torch.no_grad(): pred,_=net(xte, weights if idea else None) metric=float(((pred-yte)**2).mean().cpu()) if collect: last=stats[-1] if stats else {'pivots':0,'support':E,'simplex_error':0.,'min_x':0.,'kkt':0.} return metric,last return metric except RuntimeError: # Robust CPU fallback for constrained shared GPU slots. if dev=='cuda': torch.cuda.empty_cache() torch.set_default_device('cpu') return fit(seed,cfg,idea,collect) raise def main(): # Union parity: all lrs occur in both grids; baseline also sweeps its central # router knob (Adam optimizer) through weight decay values. grid=[{'lr':lr,'weight_decay':wd,'epochs':12,'batch':64,'lam':lam} for lr in (1e-3,3e-3,1e-2) for wd in (0.,1e-4) for lam in (1e-3,)] def base_fn(c): return lambda s: fit(s,c,False) base=sweep_baseline(base_fn,grid,seeds=SWEEP_SEEDS) best=base['best_cfg']; idea_grid=[dict(best,lam=lam) for lam in (3e-4,1e-3,3e-3)] # Evaluate every idea setting on all paired seeds; report the best by sweep-seed mean. tried=[] for c in idea_grid: r=evaluate(lambda s,c=c: fit(s,c,True), seeds=SWEEP_SEEDS) tried.append((r['mean'],c)) idea_cfg=min(tried,key=lambda x:x[0])[1] idea=evaluate(lambda s: fit(s,idea_cfg,True),seeds=DEFAULT_SEEDS) base['idea_grid']= [{'cfg':c,'mean':m} for m,c in tried] sigs=[fit(s,idea_cfg,True,True)[1] for s in DEFAULT_SEEDS] sig={'track_match':'tabular optimizer/training-procedure structure', 'prediction':'simplex router remains feasible and active-set identifies sparse support', 'predicted':{'simplex_error':0.0,'min_x':0.0,'sparse_support':True}, 'observed':{'mean_simplex_error':float(np.mean([q['simplex_error'] for q in sigs])), 'min_x':float(min(q['min_x'] for q in sigs)), 'mean_support':float(np.mean([q['support'] for q in sigs])), 'mean_pivots':float(np.mean([q['pivots'] for q in sigs]))}, 'confirmed':bool(max(q['simplex_error'] for q in sigs)<1e-6 and min(q['min_x'] for q in sigs)>=-1e-7)} rep=make_report('tabular','mlp_tiny',base,idea,sig) rep['idea_cfg']=idea_cfg with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()