import os, sys, json, random import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, sweep_baseline, evaluate, make_report SEEDS = tuple(range(8)) LRS = [1e-3, 3e-3, 1e-2] # union is used for both methods EPOCHS = 12 BATCH = 128 SIGMA = 0.20 TARGET = 0.92 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 device(): # Probe allocation and cuDNN with the exact recurrent primitive; shared GPU # failures are handled by the required CPU fallback. if torch.cuda.is_available(): try: torch.zeros(1, device='cuda') probe = nn.GRU(3, 4, batch_first=True).cuda() probe(torch.zeros(2, 2, 3, device='cuda')) return 'cuda' except Exception: try: torch.cuda.empty_cache() except Exception: pass return 'cpu' class RNNSmall(nn.Module): # Same 32-unit GRU-style recurrent architecture for both systems. def __init__(self, out_dim=1): super().__init__() self.rnn = nn.GRU(3, 32, batch_first=True) self.head = nn.Linear(32, out_dim) def forward(self, x, return_hidden=False): seq = x.view(x.shape[0], -1, 3) z, h = self.rnn(seq) if return_hidden: return self.head(h[-1]), h[-1], z return self.head(h[-1]) def rho_proxy(net): # For scalar multiplicative uncertainty M_t=(1+sigma xi)M0, # lifted spectral radius is approximately (1+sigma^2)*rho(M0)^2. # Use differentiable power iteration on recurrent hidden-hidden blocks. W = net.rnn.weight_hh_l0 v = torch.ones(W.shape[1], device=W.device) / np.sqrt(W.shape[1]) for _ in range(12): v = W.T @ (W @ v) v = v / (v.norm() + 1e-8) s2 = (W @ v).pow(2).sum() return (1.0 + SIGMA**2) * s2 def train_one(seed, lr, penalty): seed_all(seed) ds = get_dataset('dynamics', seed=seed, n_train=400, n_test=400) dev = device() net = RNNSmall().to(dev) x, y = ds['xtr'].to(dev), ds['ytr'].to(dev) opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = nn.MSELoss() for _ in range(EPOCHS): net.train(); perm = torch.randperm(len(x), device=dev) for i in range(0, len(x), BATCH): q = perm[i:i+BATCH] loss = lossf(net(x[q]), y[q]) if penalty: loss = loss + penalty * torch.relu(rho_proxy(net) - TARGET).pow(2) opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): return float(lossf(net(ds['xte'].to(dev)), ds['yte'].to(dev)).cpu()) def train_and_signature(seed, lr, penalty): seed_all(seed) ds = get_dataset('dynamics', seed=seed, n_train=400, n_test=400) dev = device(); net = RNNSmall().to(dev) x, y = ds['xtr'].to(dev), ds['ytr'].to(dev) opt = torch.optim.Adam(net.parameters(), lr=lr); lossf = nn.MSELoss() for _ in range(EPOCHS): perm = torch.randperm(len(x), device=dev) for i in range(0, len(x), BATCH): q=perm[i:i+BATCH]; loss=lossf(net(x[q]),y[q]) if penalty: loss = loss + penalty*torch.relu(rho_proxy(net)-TARGET).pow(2) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric=float(lossf(net(ds['xte'].to(dev)),ds['yte'].to(dev)).cpu()) W=net.rnn.weight_hh_l0.detach(); v=torch.ones(32,device=dev)/np.sqrt(32) for _ in range(30): v=W.T@(W@v); v=v/(v.norm()+1e-8) pred=float((1+SIGMA**2)*(W@v).pow(2).sum().cpu()) # Re-test the trained recurrent system under multiplicative recurrent noise. h=torch.randn(512,32,device=dev); zero=torch.zeros(512,1,3,device=dev) vals=[float(h.pow(2).mean().cpu())] for _ in range(20): # GRU zero-input rollout, with recurrent weights perturbed as specified. old=net.rnn.weight_hh_l0.data.clone() net.rnn.weight_hh_l0.data = old*(1+SIGMA*torch.randn_like(old)) _,hh=net.rnn(zero,h.unsqueeze(0)); h=hh[-1] net.rnn.weight_hh_l0.data = old vals.append(float(h.pow(2).mean().cpu())) ratios=np.asarray(vals[1:])/np.maximum(np.asarray(vals[:-1]),1e-12) obs=float(np.median(ratios[-8:])) return metric, pred, obs def main(): # Cheap exact math sanity check: scalar K=(a^2+sigma^2), boundary at 1. a=.8; sig=.6; math_check=float(a*a+sig*sig) grid=[{'lr':lr, 'penalty':0.0} for lr in LRS] base=sweep_baseline(lambda c: lambda s: train_one(s,c['lr'],False), grid, seeds=(0,1,2,3)) # Mandatory parity: baseline was evaluated at every idea lr; final baseline is best sweep config. best_lr=float(base['best_cfg']['lr']) idea_grid=[best_lr, 1e-3 if best_lr != 1e-3 else 3e-3, 1e-2 if best_lr != 1e-2 else 3e-3] idea_cfgs=[{'lr':float(lr),'penalty':p} for lr,p in zip(idea_grid,[1.0,3.0,10.0])] # Keep idea sweep size 3; all its lrs are in baseline union. best_idea_cfg=min(idea_cfgs, key=lambda c: np.mean([train_one(s,c['lr'],True) for s in (0,1,2,3)])) base_full=evaluate(lambda s: train_one(s,best_lr,False), SEEDS) idea_full=evaluate(lambda s: train_one(s,best_idea_cfg['lr'],True), SEEDS) sig=[train_and_signature(s,best_idea_cfg['lr'],True) for s in SEEDS] pred=float(np.mean([r[1] for r in sig])); obs=float(np.mean([r[2] for r in sig])) signature={'noise_sigma':SIGMA,'predicted_lifted_rho_mean':pred, 'observed_hidden_second_moment_ratio_mean':obs, 'relative_error':abs(obs-pred)/max(abs(pred),1e-8), 'confirmed': bool(abs(obs-pred)/max(abs(pred),1e-8) < .20), 'source':'trained idea GRU models on dynamics test rollout'} rep=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base_full},idea_full, {'mechanism_signature':signature,'math_sanity':{'a':a,'sigma':sig,'K_scalar':math_check}, 'idea_sweep':{'configs':idea_cfgs,'best_cfg':best_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()