from __future__ import annotations import json, random, sys 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 SEEDS = tuple(range(8)) EPOCHS = 15 BATCH = 128 def kron_np(xs): z = xs[0] for x in xs[1:]: z = np.kron(z, x) return z def verify_majorization(trials=40, n=3, d=3): rng = np.random.default_rng(123) failures, worst, ratios = 0, -float('inf'), [] for _ in range(trials): aa, bb = [], [] for _ in range(n): q = rng.normal(size=(d,d)); aa.append(q @ q.T + .05*np.eye(d)) q = rng.normal(size=(d,d)); bb.append(q @ q.T + .05*np.eye(d)) lm = np.linalg.eigvalsh(kron_np(aa)+kron_np(bb))[::-1] av = [np.linalg.eigvalsh(x)[::-1] for x in aa] bv = [np.linalg.eigvalsh(x)[::-1] for x in bb] c = np.sort(kron_np(av)+kron_np(bv))[::-1] gap = float(np.max(np.cumsum(lm)-np.cumsum(c))) worst = max(worst, gap); failures += int(gap > 2e-8) ratios.append(float(lm[0]/c[0])) return {'trials':trials, 'failures':failures, 'worst_cumulative_gap':worst, 'max_eigenvalue_ratio_mean':float(np.mean(ratios))} def kron_t(xs): z = xs[0] for x in xs[1:]: z = torch.kron(z, x) return z class PSDKron64(nn.Module): """Shared MLP with its first 64x64 linear map represented as two PSD terms.""" def __init__(self, input_dim, out_dim, eps=1e-3): super().__init__(); self.eps = eps self.inp = nn.Linear(input_dim, 64) self.LA = nn.ParameterList([nn.Parameter(.08*torch.randn(8,8)) for _ in range(2)]) self.LB = nn.ParameterList([nn.Parameter(.08*torch.randn(8,8)) for _ in range(2)]) self.bias = nn.Parameter(torch.zeros(64)) self.out = nn.Linear(64, out_dim) def factors(self): dev=self.LA[0].device; I=torch.eye(8,device=dev) return ([x@x.T+self.eps*I for x in self.LA], [x@x.T+self.eps*I for x in self.LB]) def bound(self,k=1): A,B=self.factors() av=[torch.linalg.eigvalsh(x).flip(0) for x in A] bv=[torch.linalg.eigvalsh(x).flip(0) for x in B] c=torch.sort(kron_t(av)+kron_t(bv),descending=True).values return c[:k].mean() def exact_norm(self): A,B=self.factors(); return torch.linalg.eigvalsh(kron_t(A)+kron_t(B))[-1] def forward(self,x): h=torch.relu(self.inp(x)) A,B=self.factors(); M=kron_t(A)+kron_t(B) h=torch.relu(h@M.T+self.bias) return self.out(h) def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) def baseline_one(cfg, seed, collect=False): seed_all(seed); ds=get_dataset('tabular',seed,n_train=400,n_test=200) net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']) net,metric,hist=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['wd'],log=lambda _:None) return float(metric) def idea_one(cfg, seed, collect=False): seed_all(seed); ds=get_dataset('tabular',seed,n_train=400,n_test=200) dev='cuda' if torch.cuda.is_available() else 'cpu' try: net=PSDKron64(int(np.prod(ds['input_shape'])),ds['out_dim']).to(dev) xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['wd']) vals=[]; bounds=[]; norms=[] for ep in range(EPOCHS): net.train(); perm=torch.randperm(len(xtr),device=dev) for i in range(0,len(xtr),BATCH): ix=perm[i:i+BATCH]; pred=net(xtr[ix]); loss=((pred-ytr[ix])**2).mean()+cfg['lam']*net.bound(cfg['k']) opt.zero_grad(); loss.backward(); opt.step(); vals.append(float(loss.detach().cpu())) with torch.no_grad(): bounds.append(float(net.bound(cfg['k']).cpu())); norms.append(float(net.exact_norm().cpu())) net.eval() with torch.no_grad(): metric=float(((net(ds['xte'].to(dev))-ds['yte'].to(dev))**2).mean().cpu()) if collect: b1=[] with torch.no_grad(): for _ in norms: b1.append(float(net.bound(1).cpu())) return metric, {'bound_R4':float(np.mean(bounds)),'bound_R1':float(np.mean(b1)),'exact_norm':float(np.mean(norms)),'ratio_R1':float(np.mean(np.array(norms)/np.array(b1)))} return metric except Exception: if dev=='cuda': torch.cuda.empty_cache(); torch.set_default_device('cpu'); return idea_one(cfg,seed,collect) raise def main(): verification=verify_majorization() # Union parity: every idea lr/wd is also evaluated by baseline. grid=[{'lr':lr,'wd':wd,'lam':lam,'k':k} for lr in (0.001,0.003,0.006) for wd in (0.0,1e-4) for lam,k in ((0.0,1),)] basegrid=[{'lr':x['lr'],'wd':x['wd']} for x in grid] base=sweep_baseline(lambda c: (lambda s: baseline_one(c,s)),basegrid,seeds=tuple(range(4))) # Baseline sweep includes all idea learning rates; final best is evaluated on 8 seeds. idea_grid=[{'lr':lr,'wd':wd,'lam':lam,'k':k} for lr in (0.001,0.003,0.006) for wd in (0.0,1e-4) for lam,k in ((0.001,1),(0.003,1),(0.001,4))] tried=[] for c in idea_grid: r=evaluate(lambda s: idea_one(c,s),seeds=tuple(range(4))); tried.append({'cfg':c,'mean':r['mean']}) best=min(tried,key=lambda x:x['mean'])['cfg'] idea=evaluate(lambda s: idea_one(best,s),seeds=SEEDS) # Trained-model signature, independently measured on two paired trained models. sig=[] for s in (0,1): _,z=idea_one(best,s,True); sig.append(z) signature={'quantity':'exact operator norm / separable R1 bound on trained models','predicted':'ratio <= 1 by the Ky-Fan upper bound','observed_mean_ratio':float(np.mean([z['ratio_R1'] for z in sig])),'observed':sig,'confirmed':all(z['ratio_R1']<=1.0001 for z in sig)} base['idea_union_sweep']=tried rep=make_report('tabular','mlp_tiny',base,idea,{'mechanism_signature':signature,'verification':verification,'idea_best_cfg':best}) rep['stage2_note']='Baseline and idea use same tabular data, epochs, batch size, Adam, learning-rate/weight-decay union; idea changes only the first hidden operator and adds Ky-Fan penalty.' with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()