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, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) EPOCHS = 8 BATCH = 128 ALPHA = 0.08 MARGIN = 0.01 PENALTY_WEIGHT = 0.08 # The intervention is a sampled discrete differential-contraction penalty on # the recurrent state transition. The task head and all non-intervention # architecture are exactly bench's rnn_small. def device_and_seed(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' try: if dev == 'cuda': torch.cuda.empty_cache() except Exception: dev = 'cpu' return dev def transition_jacobian(net, x, h): """Differentiable closed-form Jacobian of one GRU step wrt h.""" c = net.rnn; wi, wh = c.weight_ih_l0, c.weight_hh_l0 gi = x @ wi.T + c.bias_ih_l0; gh = h @ wh.T + c.bias_hh_l0 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) d = h.shape[-1] Dr = torch.diag_embed(r * (1-r)); Dz = torch.diag_embed(z * (1-z)) Dn = torch.diag_embed(1-n*n); D1mz = torch.diag_embed(1-z) Dh = torch.diag_embed(h.squeeze(1) - n.squeeze(1)) # PyTorch GRU reset-after-matrix convention: n=tanh(inn+r*(Wh_n h+b)). Jr = Dr @ wh[:d] Jz = Dz @ wh[d:2*d] Jn = Dn @ (wh[2*d:] + torch.diag_embed(hnn.squeeze(1)) @ Jr) return (D1mz @ Jn - Dh @ Jz + torch.diag_embed(z.squeeze(1)))[0] def recurrent_outputs(net, x): seq = x.view(x.shape[0], -1, 3) out, _ = net.rnn(seq) return seq, out def contraction_penalty(net, x, alpha=ALPHA): seq, hs = recurrent_outputs(net, x) vals = [] I = torch.eye(hs.shape[-1], device=x.device) for b in range(min(2, x.shape[0])): for t in (seq.shape[1] - 1,): M = transition_jacobian(net, seq[b:b+1, t], hs[b, t:t+1]) lam = torch.linalg.eigvalsh(M.T @ M - (1.0-alpha) * I)[-1] vals.append(torch.nn.functional.softplus(lam + MARGIN) ** 2) return torch.stack(vals).mean() def train_idea(net, d, epochs=EPOCHS, lr=3e-3, seed=0): dev = next(net.parameters()).device net.train(); opt = torch.optim.Adam(net.parameters(), lr=lr) x, y = d['xtr'].to(dev), d['ytr'].to(dev) gen = torch.Generator(device=dev).manual_seed(seed + 91) n = len(x) for _ in range(epochs): for ix in torch.randperm(n, generator=gen, device=dev).split(BATCH): xb, yb = x[ix], y[ix] opt.zero_grad(); pred = net(xb) task = ((pred - yb) ** 2).mean() pen = contraction_penalty(net, xb) (task + PENALTY_WEIGHT * pen).backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0) opt.step() net.eval() with torch.no_grad(): return float(((net(d['xte'].to(dev))-d['yte'].to(dev))**2).mean().cpu()) def run_idea(cfg): def fn(seed): dev = device_and_seed(seed) d = get_dataset('dynamics', seed, n_train=400, n_test=400) try: net = make_model('rnn_small', d['input_shape'], d['out_dim']).to(dev) return train_idea(net, d, lr=cfg['lr'], seed=seed) except Exception: if dev != 'cuda': raise net = make_model('rnn_small', d['input_shape'], d['out_dim']).cpu() return train_idea(net, {k:(v.cpu() if torch.is_tensor(v) else v) for k,v in d.items()}, lr=cfg['lr'], seed=seed) return fn def run_base(cfg): def fn(seed): device_and_seed(seed) d = get_dataset('dynamics', seed, n_train=400, n_test=400) # Required standard path; bench itself supplies GPU->CPU fallback. _, metric, _ = train_model(make_model('rnn_small', d['input_shape'], d['out_dim']), d, epochs=EPOCHS, lr=cfg['lr']) return float(metric) return fn def signature(seed, cfg, idea): dev = device_and_seed(seed); d = get_dataset('dynamics', seed, 400, 400) try: net = make_model('rnn_small', d['input_shape'], d['out_dim']).to(dev) if idea: train_idea(net, d, lr=cfg['lr'], seed=seed) else: train_model(net, d, epochs=EPOCHS, lr=cfg['lr']) x = d['xte'][:16].to(dev) with torch.no_grad(): seq, hs = recurrent_outputs(net, x) vals=[]; ratios=[]; I=torch.eye(hs.shape[-1], device=dev) for b in range(4): M=transition_jacobian(net,seq[b:b+1,0],hs[b,0:1]) vals.append(float(torch.linalg.eigvalsh(M.T@M-(1-ALPHA)*I)[-1].detach().cpu())) h1=hs[b,0:1].detach(); h2=h1+0.01*torch.ones_like(h1) with torch.no_grad(): a=transition_jacobian(net, seq[b:b+1,0], h1) @ (h1-h1).T # finite perturbation through the actual GRU recurrence def step(q): return recurrent_outputs(net, torch.cat([seq[b:b+1,0:1], seq[b:b+1,0:1]], 1))[1][:,0:1] c=net.rnn; xx=seq[b:b+1,0] def actual(q): gi=xx@c.weight_ih_l0.T+c.bias_ih_l0; gh=q@c.weight_hh_l0.T+c.bias_hh_l0 ir,iz,inn=gi.chunk(3,-1); hr,hz,hnn=gh.chunk(3,-1) rr=torch.sigmoid(ir+hr); zz=torch.sigmoid(iz+hz); nnx=torch.tanh(inn+rr*hnn) return (1-zz)*nnx+zz*q aa,zz=actual(h1),actual(h2) ratios.append(float((torch.linalg.vector_norm(zz-aa)/torch.linalg.vector_norm(h2-h1)).cpu())) return {'predicted_ratio_bound': math.sqrt(1-ALPHA), 'observed_mean_ratio': float(np.mean(ratios)), 'predicted_lmi_bound': 0.0, 'observed_max_lmi': max(vals), 'confirmed': float(np.mean(ratios)) <= math.sqrt(1-ALPHA)+0.03 and max(vals) <= 0.03} except Exception as e: return {'predicted_ratio_bound': math.sqrt(1-ALPHA), 'observed_mean_ratio': None, 'predicted_lmi_bound': 0.0, 'observed_max_lmi': None, 'confirmed': False, 'error': repr(e)} def main(): grid=[{'lr':1e-3},{'lr':3e-3},{'lr':1e-2}] base=sweep_baseline(run_base, grid, seeds=SWEEP_SEEDS) idea_runs=[] for cfg in grid: r=evaluate(run_idea(cfg), seeds=SEEDS) idea_runs.append((r,cfg)) best_idea,best_cfg=min(idea_runs, key=lambda z:z[0]['mean']) report=make_report('dynamics','rnn_small',base,best_idea,{'mechanism_signature': signature(0,best_cfg,True), 'idea_sweep':[{'cfg':c,'result':r} for r,c in idea_runs], 'structural_match':'Dynamics track directly tests recurrent trajectory sensitivity and long-horizon stability.'}) report['baseline']['union_grid']=grid report['selected_idea_cfg']=best_cfg Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()