import os, sys, json, random, 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, make_model, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0,1,2,3) H = 64 def seed_all(s): random.seed(s); np.random.seed(s); torch.manual_seed(s) if torch.cuda.is_available(): torch.cuda.manual_seed_all(s) def device_for(): try: d = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if d.type == 'cuda': torch.zeros(1, device=d) return d except Exception: return torch.device('cpu') def transition(model, x, h): """PyTorch GRUCell equations, preserving parameter graph.""" wih, whh = model.rnn.weight_ih_l0, model.rnn.weight_hh_l0 bih = model.rnn.bias_ih_l0 bhh = model.rnn.bias_hh_l0 gi = torch.nn.functional.linear(x, wih, bih) gh = torch.nn.functional.linear(h, whh, bhh) ir, iz, inn = gi.chunk(3, -1); hr, hz, hnn = gh.chunk(3, -1) r = torch.sigmoid(ir + hr) z = torch.sigmoid(iz + hz) n = torch.tanh(inn + r * hnn) return (1-z)*n + z*h def block_jacobian(model, x, h): # Two contiguous hidden blocks; exact Jacobian is only 64x64 on one sample. def f(v): return transition(model, x, v) J = torch.autograd.functional.jacobian(f, h, create_graph=True, vectorize=True) # J[output,input], aggregate absolute derivative magnitudes by blocks. a = J[:H//2,:H//2].abs().mean(); b = J[:H//2,H//2:].abs().mean() c = J[H//2:,:H//2].abs().mean(); d = J[H//2:,H//2:].abs().mean() K = torch.stack((torch.stack((a,b)), torch.stack((c,d)))) return K, J def cycle_penalty(model, xb, tau=1.0): seq = xb[:1].view(1,-1,3) h = torch.zeros(1,H,device=xb.device) # Build a representative hidden state using the same recurrent equations. for t in range(seq.shape[1]-1): h = transition(model, seq[:,t], h) K, J = block_jacobian(model, seq[0,-1], h[0]) tr2 = torch.trace(K @ K) return tr2/(tau*tau*2.0), float(tr2.detach().cpu()), float(torch.linalg.matrix_norm(J).detach().cpu()) def train_cycle(seed, cfg, return_sig=False): seed_all(seed); ds=get_dataset('dynamics', seed, n_train=200, n_test=100) model=make_model('rnn_small', ds['input_shape'], ds['out_dim']) dev=device_for() try: model.to(dev); xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev) opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['wd']) lossf=nn.MSELoss() for ep in range(cfg['epochs']): perm=torch.randperm(len(xtr),device=dev) for i in range(0,len(xtr),64): ix=perm[i:i+64]; pred=model(xtr[ix]); task=lossf(pred,ytr[ix]) pen,_,_=cycle_penalty(model,xtr[ix],cfg['tau']) opt.zero_grad(); (task+cfg['lam']*pen).backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step() with torch.no_grad(): metric=float(((model(ds['xte'].to(dev))-ds['yte'].to(dev))**2).mean()) if return_sig: pen,tr2,jn=cycle_penalty(model,xtr[:1],cfg['tau']) return metric, {'closed_walk_trace2':tr2,'jacobian_frobenius':jn,'penalty':float(pen.detach().cpu())} return metric except RuntimeError: # Explicit CPU fallback for shared/limited CUDA environments. seed_all(seed); os.environ['CUDA_VISIBLE_DEVICES']='' return train_cycle_cpu(seed,cfg,return_sig) def train_cycle_cpu(seed,cfg,return_sig=False): seed_all(seed); ds=get_dataset('dynamics',seed,n_train=200,n_test=100); model=make_model('rnn_small',ds['input_shape'],ds['out_dim']) xtr,ytr=ds['xtr'],ds['ytr']; opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['wd']); lossf=nn.MSELoss() for ep in range(cfg['epochs']): perm=torch.randperm(len(xtr)) for i in range(0,len(xtr),64): ix=perm[i:i+64]; task=lossf(model(xtr[ix]),ytr[ix]); pen,_,_=cycle_penalty(model,xtr[ix],cfg['tau']); opt.zero_grad(); (task+cfg['lam']*pen).backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step() metric=float(((model(ds['xte'])-ds['yte'])**2).mean()) if return_sig: pen,tr2,jn=cycle_penalty(model,xtr[:1],cfg['tau']); return metric, {'closed_walk_trace2':tr2,'jacobian_frobenius':jn,'penalty':float(pen.detach())} return metric def base_fn(cfg): def run(seed): seed_all(seed); ds=get_dataset('dynamics',seed,n_train=200,n_test=100); m=make_model('rnn_small',ds['input_shape'],ds['out_dim']) _,metric,_=train_model(m,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=64,weight_decay=cfg['wd'],log=lambda *_:None) return metric return run def main(): # tr(K^2)=2ab is the core closed-walk identity; threshold is tested numerically. a,b=1.8,0.8; K=np.array([[0,a],[b,0]],float); identity=float(np.trace(K@K)); rho=np.sqrt(a*b) sanity={'trace2':identity,'expected_2ab':2*a*b,'abs_error':abs(identity-2*a*b),'rho':float(rho),'critical_gain':float(1/rho)} lrs=[0.0015,0.003,0.006]; wds=[0.0,1e-4]; epochs=3 grid=[{'lr':lr,'wd':wd,'epochs':epochs} for lr in lrs for wd in wds] base=sweep_baseline(base_fn,grid,seeds=SWEEP_SEEDS) # Same lr/wd union; lambda is the sole extra method knob. ideas=[{'lr':lr,'wd':wd,'epochs':epochs,'lam':lam,'tau':tau} for lr in lrs for wd in wds for lam in [0.001,0.005,0.02] for tau in [1.0]] rows=[]; best=None for cfg in ideas: r=evaluate(lambda s,cfg=cfg: train_cycle(s,cfg),seeds=SEEDS); rows.append({'cfg':cfg,'mean':r['mean']}) if best is None or r['mean']