import os, sys, json, math, 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, make_model, train_model, evaluate, sweep_baseline, make_report SEEDS=tuple(range(8)) # n=8 tangent subspace keeps the compound penalty small enough for the bench. N_TAN=8 HIDDEN=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 pairs(n): return [(i,j) for i in range(n) for j in range(i+1,n)] def additive_compound(A): n=A.shape[-1]; ps=pairs(n); m=len(ps) B=A.new_zeros(A.shape[:-2]+(m,m)) for a,(i,j) in enumerate(ps): B[...,a,a]=A[...,i,i]+A[...,j,j] for b,(k,l) in enumerate(ps): if a != b: # coefficient of E_ij in A^[2] acting on e_k wedge e_l B[...,a,b] = (A[...,i,k]*(1.0 if j==l else 0.0) + A[...,j,l]*(1.0 if i==k else 0.0) - A[...,i,l]*(1.0 if j==k else 0.0) - A[...,j,k]*(1.0 if i==l else 0.0)) return B def wedge_math_check(): torch.manual_seed(7); A=torch.randn(5,5,dtype=torch.float64) # finite difference derivative of multiplicative compound, using minors ps=pairs(5) def wedge(M): W=torch.empty(10,10,dtype=M.dtype) for a,(i,j) in enumerate(ps): for b,(k,l) in enumerate(ps): W[a,b]=M[i,k]*M[j,l]-M[i,l]*M[j,k] return W hs=torch.tensor([.02,.01,.005,.0025],dtype=torch.float64); errs=[] for h in hs: errs.append(float(torch.linalg.norm(wedge(torch.matrix_exp(h*A))-torch.eye(10)-h*additive_compound(A)))) slope=float(np.polyfit(np.log(hs.numpy()),np.log(errs),1)[0]) return {'identity_error_slope':slope,'expected':2.0,'max_error':max(errs),'passed':abs(slope-2)<.12} def cell8(model, x3, h8): """First eight rows/columns of the exact GRU Jacobian, analytically.""" W=model.rnn.weight_ih_l0; U=model.rnn.weight_hh_l0 b1=model.rnn.bias_ih_l0; b2=model.rnn.bias_hh_l0 h=torch.cat([h8, torch.zeros(h8.shape[0], HIDDEN-N_TAN, device=h8.device, dtype=h8.dtype)], dim=1) gi=torch.nn.functional.linear(x3,W,b1); gh=torch.nn.functional.linear(h,U,b2) ar=gi[:,:HIDDEN]+gh[:,:HIDDEN]; aq=gi[:,HIDDEN:2*HIDDEN]+gh[:,HIDDEN:2*HIDDEN] r=torch.sigmoid(ar); q=torch.sigmoid(aq) an=gi[:,2*HIDDEN:3*HIDDEN]+r*gh[:,2*HIDDEN:3*HIDDEN] n=torch.tanh(an); d=gh[:,2*HIDDEN:3*HIDDEN] # Batch is one in tangent_penalty; retain differentiability wrt parameters. dr=r*(1-r) dq=q*(1-q) # Each derivative tensor has shape [1,HIDDEN,N_TAN]. Ur=U[:HIDDEN,:N_TAN].unsqueeze(0) Uq=U[HIDDEN:2*HIDDEN,:N_TAN].unsqueeze(0) Un=U[2*HIDDEN:3*HIDDEN,:N_TAN].unsqueeze(0) drdh=dr.unsqueeze(2)*Ur dqdh=dq.unsqueeze(2)*Uq dndh=(1-n*n).unsqueeze(2)*(r.unsqueeze(2)*Un+d.unsqueeze(2)*drdh) dy=(1-q).unsqueeze(2)*dndh + dqdh*(h[:,:HIDDEN].unsqueeze(2)-n.unsqueeze(2)) dy=dy + q.unsqueeze(2)*torch.eye(HIDDEN,device=h.device,dtype=h.dtype)[:,:N_TAN].unsqueeze(0) return ((1-q)*n+q*h[:,:HIDDEN])[:,:N_TAN], dy[:,:N_TAN,:] def tangent_penalty(model, x): x3=x[:, -3:][:1] h=torch.zeros(1,N_TAN,device=x.device) _, A=cell8(model,x3,h) A=A[0] C=additive_compound(A) off=C-torch.diag(torch.diagonal(C)) neg=torch.relu(-off) return (neg*neg).mean()+0.02*torch.relu(0.01-off).pow(2).mean(), A.detach(), C.detach() def train_idea(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']) device='cuda' if torch.cuda.is_available() else 'cpu' try: model.to(device); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['wd']) lossf=nn.MSELoss(); xtr,ytr=ds['xtr'].to(device),ds['ytr'].to(device) for ep in range(cfg['epochs']): model.train(); perm=torch.randperm(len(xtr),device=device) for i in range(0,len(xtr),64): ix=perm[i:i+64]; pred=model(xtr[ix]); task=lossf(pred,ytr[ix]) # Applying at every mini-batch makes the intervention part of training. pen,_,_=tangent_penalty(model,xtr[ix]) loss=task+cfg['lam']*pen opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step() model.eval() with torch.no_grad(): metric=float(((model(ds['xte'].to(device))-ds['yte'].to(device))**2).mean()) if return_sig: with torch.enable_grad(): pen,A,C=tangent_penalty(model,xtr[:1]) off=C-torch.diag(torch.diagonal(C)); return metric, {'negative_offdiag_fraction':float((off<0).float().mean()),'mean_negative_magnitude':float(torch.relu(-off).mean()),'penalty':float(pen),'jacobian_norm':float(torch.linalg.norm(A))} return metric except RuntimeError: # CPU fallback for tight CUDA/cuDNN environments. seed_all(seed); os.environ['CUDA_VISIBLE_DEVICES']='' return train_idea_cpu(seed,cfg,return_sig) def train_idea_cpu(seed,cfg,return_sig=False): # same implementation after disabling CUDA; recursive body is intentionally avoided. old=torch.cuda.is_available ds=get_dataset('dynamics',seed,n_train=200,n_test=100); model=make_model('rnn_small',ds['input_shape'],ds['out_dim']) opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['wd']); lossf=nn.MSELoss(); xtr,ytr=ds['xtr'],ds['ytr'] for _ in range(cfg['epochs']): for i in range(0,len(xtr),64): task=lossf(model(xtr[i:i+64]),ytr[i:i+64]); pen,_,_=tangent_penalty(model,xtr[i:i+64]); opt.zero_grad(); (task+cfg['lam']*pen).backward(); opt.step() with torch.no_grad(): metric=float(((model(ds['xte'])-ds['yte'])**2).mean()) if return_sig: pen,A,C=tangent_penalty(model,xtr[:1]); off=C-torch.diag(torch.diagonal(C)); return metric,{'negative_offdiag_fraction':float((off<0).float().mean()),'mean_negative_magnitude':float(torch.relu(-off).mean()),'penalty':float(pen),'jacobian_norm':float(torch.linalg.norm(A))} 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(): check=wedge_math_check(); lrs=[0.0015,0.003,0.006]; wds=[0.0] grid=[{'lr':lr,'wd':wd,'epochs':3} for lr in lrs for wd in [0.0]] base=sweep_baseline(base_fn,grid,seeds=(0,1,2,3)) # Same lr/weight-decay union; lambda sweep is the sole method difference. ideas=[] for lr in lrs: for lam in [0.001,0.005,0.02]: ideas.append({'lr':lr,'wd':base['best_cfg']['wd'],'epochs':3,'lam':lam}) best=None; rows=[] for cfg in ideas: r=evaluate(lambda s,cfg=cfg: train_idea(s,cfg),seeds=SEEDS); rows.append({'cfg':cfg,'mean':r['mean']}) if best is None or r['mean']