import sys, json, random, 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, sweep_baseline, evaluate, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 12 BATCH = 128 LR_GRID = [1.5e-3, 3e-3, 6e-3] DELTA_GRID = [0.03, 0.05, 0.10] 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 rho_hh(net): w = net.rnn.weight_hh_l0.detach().float().cpu().numpy() # GRU has three gate blocks; use the largest block as the feedback estimate. h = w.shape[1] vals = [] for block in np.array_split(w, 3, axis=0): vals.append(np.max(np.abs(np.linalg.eigvals(block)))) return float(max(vals)) def controller(net, delta): r = rho_hh(net) target = 1.0 - delta if r > target: scale = target / (r + 1e-8) with torch.no_grad(): net.rnn.weight_hh_l0.mul_(scale) return r def train(seed, lr, controlled=False, delta=0.05, clip=None, capture=False): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=100) net = make_model('rnn_small', 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) xt, yt = ds['xte'].to(device), ds['yte'].to(device) opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = nn.MSELoss(); hist=[]; rh=[] for ep in range(EPOCHS): net.train(); perm=torch.randperm(len(x), device=device); total=0. for i in range(0,len(x),BATCH): idx=perm[i:i+BATCH]; loss=lossf(net(x[idx]),y[idx]) opt.zero_grad(); loss.backward() if clip is not None: torch.nn.utils.clip_grad_norm_(net.parameters(), clip) opt.step() if controlled: rh.append(controller(net, delta)) total += float(loss.detach())*len(idx) hist.append(total/len(x)) net.eval() with torch.no_grad(): metric=float(((net(xt)-yt)**2).mean()) if not rh: rh=[rho_hh(net)] result=(metric, net, ds, {'rho_mean':float(np.mean(rh)), 'rho_max':float(np.max(rh)), 'rho_final':rho_hh(net), 'history':hist}) return result if capture else metric except RuntimeError: if device == 'cuda': torch.cuda.empty_cache(); return train_cpu(seed, lr, controlled, delta, clip, capture) raise def train_cpu(seed, lr, controlled=False, delta=0.05, clip=None, capture=False): old=torch.cuda.is_available # Re-execute on CPU explicitly, avoiding any CUDA allocation. seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=100) net=make_model('rnn_small',ds['input_shape'],ds['out_dim']).cpu(); x,y=ds['xtr'],ds['ytr']; xt,yt=ds['xte'],ds['yte'] opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss(); hist=[]; rh=[] for ep in range(EPOCHS): perm=torch.randperm(len(x)); total=0. for i in range(0,len(x),BATCH): idx=perm[i:i+BATCH]; loss=lossf(net(x[idx]),y[idx]); opt.zero_grad(); loss.backward() if clip is not None: torch.nn.utils.clip_grad_norm_(net.parameters(),clip) opt.step() if controlled: rh.append(controller(net,delta)) total+=float(loss.detach())*len(idx) hist.append(total/len(x)) net.eval(); metric=float(((net(xt)-yt)**2).mean()); info={'rho_mean':float(np.mean(rh or [rho_hh(net)])),'rho_max':float(np.max(rh or [rho_hh(net)])),'rho_final':rho_hh(net),'history':hist} return (metric,net,ds,info) if capture else metric def baseline_fn(cfg): return lambda s: train(s, cfg['lr'], False, 0.05, cfg['clip']) def idea_fn(cfg): return lambda s: train(s, cfg['lr'], True, cfg['delta'], None) def signature(base_cfg, idea_cfg): # Signature is measured on trained benchmark models: linear feedback prediction # versus local hidden-state Jacobian amplification on actual test sequences. pred=[]; obs=[] for s in (0,1,2,3): _, net, ds, _ = train(s, idea_cfg['lr'], True, idea_cfg['delta'], None, True) cell=nn.GRUCell(3,64); cell.load_state_dict({'weight_ih':net.rnn.weight_ih_l0.detach().cpu(), 'weight_hh':net.rnn.weight_hh_l0.detach().cpu(), 'bias_ih':net.rnn.bias_ih_l0.detach().cpu(), 'bias_hh':net.rnn.bias_hh_l0.detach().cpu()}) x=ds['xte'][0].view(-1,3); h=torch.zeros(64); ratios=[] for t in range(x.shape[0]): z=x[t]; h0=h.detach().requires_grad_(True) J=torch.autograd.functional.jacobian(lambda q: cell(z,q), h0) ratios.append(float(torch.linalg.svdvals(J).max())) h=cell(z,h).detach() pred.append(rho_hh(net)); obs.append(float(np.exp(np.mean(np.log(np.maximum(ratios,1e-9)))))) p=float(np.mean(pred)); o=float(np.mean(obs)) return {'predicted_feedback_radius':p,'observed_local_hidden_jacobian_gain':o,'relative_gap':abs(p-o)/(abs(p)+1e-9),'n_models':4,'confirmed':bool(abs(p-o)/(abs(p)+1e-9)<0.35)} def main(): # Baseline method knob parity: Adam gradient clipping is swept alongside all lrs. grid=[{'lr':lr,'clip':clip} for lr in LR_GRID for clip in (None,1.0)] base=sweep_baseline(baseline_fn,grid,seeds=SWEEP_SEEDS) idea_grid=[{'lr':lr,'delta':d} for lr,d in zip(LR_GRID,[0.03,0.05,0.10])] # choose idea by same 4-seed budget; all idea lrs are in baseline union. tried=[] for cfg in idea_grid: r=evaluate(idea_fn(cfg),SWEEP_SEEDS); tried.append({'cfg':cfg,'mean':r['mean']}) best=min(tried,key=lambda z:z['mean'])['cfg'] idea=evaluate(idea_fn(best),SEEDS); base['idea_grid']=tried; base['idea_best_cfg']=best # mechanism signature uses independently trained models, not toy arithmetic. sig=signature(base['best_cfg'],best) report=make_report('dynamics','rnn_small',base,idea,{'prediction':'feedback spectral radius should be held below 1-delta and correspond to local hidden sensitivity','measurement':sig}) report['protocol']={'epochs':EPOCHS,'batch':BATCH,'paired_seeds':list(SEEDS),'baseline_grid':grid,'idea_grid':idea_grid} Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()