import os, sys, json, math from pathlib import Path 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 from bench.protocol import evaluate, sweep_baseline, make_report OUT = Path('bench_report.json') DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' def schur(H, G, lam=1e-2, beta=2.0): A = H + lam*np.eye(H.shape[0]) R = G.T @ np.linalg.solve(A, G) R = (R + R.T)/2 ev, V = np.linalg.eigh(R) ev = np.maximum(ev, 0) R = (V*ev) @ V.T P = np.linalg.solve(np.eye(R.shape[0]) + beta*R, np.eye(R.shape[0])) return R, P def verify_math(seed=491): rng=np.random.default_rng(seed); k,m=4,3 A=rng.normal(size=(k,k)); H=A.T@A+.7*np.eye(k); G=rng.normal(size=(k,m)); a=rng.normal(size=m) def f(u): s=-np.linalg.solve(H,G@u) return .5*s@H@s+s@G@u+a@u e=2e-4; num=np.zeros((m,m)); eye=np.eye(m); u=rng.normal(size=m) for i in range(m): for j in range(m): ei,ej=eye[i],eye[j] num[i,j]=(f(u+e*ei+e*ej)-f(u+e*ei-e*ej)-f(u-e*ei+e*ej)+f(u-e*ei-e*ej))/(4*e*e) target=-G.T@np.linalg.solve(H,G) R,_=schur(H,G,0,1) return {'max_abs_hessian_error':float(np.max(abs(num-target))), 'min_R_eigenvalue':float(np.linalg.eigvalsh(R).min()), 'identity_pass':bool(np.max(abs(num-target))<2e-6 and np.linalg.eigvalsh(R).min()>-1e-9)} def baseline_one(seed, lr, wd, epochs=18): np.random.seed(seed); torch.manual_seed(seed) ds=get_dataset('tabular', seed, n_train=1600, n_test=500) model=make_model('mlp_tiny', ds['input_shape'], ds['out_dim']) _, metric, _=train_model(model, ds, epochs=epochs, lr=lr, batch=128, weight_decay=wd, log=lambda *_:None) return float(metric) def idea_one(seed, lr, wd, epochs=18, capture=None): np.random.seed(seed); torch.manual_seed(seed) ds=get_dataset('tabular', seed, n_train=1600, n_test=500) # Same mlp_tiny architecture; intervention is optimizer mechanism only. model=make_model('mlp_tiny', ds['input_shape'], ds['out_dim']) device=DEVICE try: model=model.to(device); x=ds['xtr'].to(device); y=ds['ytr'].to(device) groups=[list(model[:2].parameters()), list(model[2:4].parameters()), [model[4].weight], [model[4].bias]] # Deliberately heterogeneous initial amplitudes, as the proposal targets. amps=np.array([2.5,.4,2.0,.5], dtype=float) params=[p for g in groups for p in g] opt=torch.optim.Adam(params, lr=lr, weight_decay=wd) lossf=nn.MSELoss(); rng=np.random.default_rng(seed) # sigma is a relaxed, four-component controller (EMA gradient norm). sigma=np.zeros(4); prev=np.zeros(4); records=[]; maxeig=0. for ep in range(epochs): model.train(); perm=torch.randperm(len(x),device=device) for ii in range(0,len(x),128): ix=perm[ii:ii+128]; loss=lossf(model(x[ix]),y[ix]) opt.zero_grad(); loss.backward() gn=np.array([math.sqrt(sum(float((p.grad.detach()**2).sum()) for p in g if p.grad is not None)+1e-12) for g in groups]) # Controller relaxation and sensitivity G: finite difference in each u sigma=.9*sigma+.1*gn if (ep*math.ceil(len(x)/128)+ii//128)%8==0: # local controller energy: 1/2 sigma^T H sigma + sigma^T G u; # G is measured sensitivity of controller state to log-rate amplitude. H=np.diag(0.1+sigma/(np.mean(sigma)+1e-6)) G=np.zeros((4,4)); base=sigma.copy() for j in range(4): # one-step virtual Adam displacement estimates response of hidden state eps=.03; G[j,j]=-(abs(base[j])+0.05)*amps[j]*eps if j>0: G[j-1,j]=-.12*abs(base[j])*eps R,P=schur(H,G,lam=.05,beta=3.0); maxeig=max(maxeig,float(np.linalg.eigvalsh(R).max())) factors=np.diag(P) for j,g in enumerate(groups): # trust correction applies to mechanism amplitudes, not model/readout for p in g: p.grad.mul_(float(amps[j]*factors[j])) observed=float(np.linalg.norm(amps*factors)/(np.linalg.norm(amps)+1e-12)) records.append((float(np.linalg.norm(np.diag(P))),observed)) else: for j,g in enumerate(groups): for p in g: p.grad.mul_(float(amps[j])) opt.step() model.eval() with torch.no_grad(): metric=float(((model(ds['xte'].to(device))-ds['yte'].to(device))**2).mean()) if capture is not None: capture.append({'seed':seed,'max_R_eigenvalue':maxeig, 'predicted_mean_shrink':float(np.mean([r[0] for r in records])) if records else 1., 'observed_mean_shrink':float(np.mean([r[1] for r in records])) if records else 1.}) return metric except RuntimeError: if torch.cuda.is_available(): torch.cuda.empty_cache() # robust CPU retry old=DEVICE globals()['DEVICE']='cpu' try: return idea_one(seed,lr,wd,epochs,capture) finally: globals()['DEVICE']=old def main(): math_check=verify_math() # Union of all learning rates tried by either side; baseline also sweeps Adam WD. lrs=[1e-3,3e-3,6e-3]; wds=[0.0,1e-4] grid=[{'lr':lr,'weight_decay':wd} for lr in lrs for wd in wds] def mk(cfg): return lambda seed: baseline_one(seed,cfg['lr'],cfg['weight_decay']) base=sweep_baseline(mk,grid) # Three idea settings at the baseline-selected WD and shared LR union. sig=[] idea_cfgs=[{'lr':lr,'weight_decay':base['best_cfg']['weight_decay']} for lr in lrs] idea_runs=[] for cfg in idea_cfgs: cap=[] r=evaluate(lambda seed,cfg=cfg: idea_one(seed,cfg['lr'],cfg['weight_decay'],capture=cap)) idea_runs.append({'cfg':cfg,'result':r,'signature_records':cap}) best=min(idea_runs,key=lambda z:z['result']['mean']) sig=best['signature_records'] pred=float(np.mean([x['predicted_mean_shrink'] for x in sig])); obs=float(np.mean([x['observed_mean_shrink'] for x in sig])) signature={'quantity':'Schur trust shrinkage measured during trained MLP updates', 'predicted_vs_observed':{'predicted_mean_shrink':pred,'observed_mean_shrink':obs, 'relative_error':abs(pred-obs)/(abs(obs)+1e-12)}, 'max_R_eigenvalue_mean':float(np.mean([x['max_R_eigenvalue'] for x in sig])), 'confirmed':bool(abs(pred-obs)/(abs(obs)+1e-12)<0.25)} rep=make_report('tabular','mlp_tiny',base,best['result'],{ 'mechanism_signature':signature,'idea_sweep':idea_runs, 'math_verification':math_check,'device':DEVICE, 'protocol_note':'Baseline sweep uses all idea learning rates and both Adam weight decays; 8 paired seeds.'}) OUT.write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2)) if __name__=='__main__': main()