import sys, json, random from pathlib import Path import numpy as np sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import torch import torch.nn as nn from bench import get_dataset, make_model, train_model, sweep_baseline, make_report SEEDS = tuple(range(8)) LRS = [1e-3, 3e-3, 6e-3] EPOCHS = 8 BATCH = 128 PRIOR = 0.05 BETA = 0.01 TAU = 0.02 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 baseline_fn(cfg): def run(seed): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=300, n_test=150) net = make_model('rnn_small', ds['input_shape'], ds['out_dim']) _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=0.0, log=lambda *_: None) return float(metric) if metric is not None else float('inf') return run def softmin(vals, tau): return -tau * torch.logsumexp(-vals / tau, dim=0) def curvature_penalty(net, x): # Input Jacobian is measured on the actual recurrent model and actual benchmark # windows. For scalar output, J^T J is rank one per sample; prior makes H PD. x = x.detach().requires_grad_(True) out = net(x).reshape(-1) rows = [] for i in range(len(out)): # Samples are independent, so one gradient of the summed scalar # yields every per-sample input Jacobian row. g = torch.autograd.grad(out.sum(), x, retain_graph=True, create_graph=True)[0] rows = [g[j].reshape(-1) for j in range(len(g))] J = torch.stack(rows) H = PRIOR * torch.eye(J.shape[1], device=J.device, dtype=J.dtype) + J.T @ J / max(1, len(rows)) ev = torch.linalg.eigvalsh(H) return softmin(ev, TAU), H.detach(), ev.detach() def idea_run(seed, lr): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=300, n_test=150) device = 'cuda' if torch.cuda.is_available() else 'cpu' for dev in ([device, 'cpu'] if device == 'cuda' else ['cpu']): try: net = make_model('rnn_small', ds['input_shape'], ds['out_dim']).to(dev) xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev) opt = torch.optim.Adam(net.parameters(), lr=lr) n = len(xtr) for _ in range(EPOCHS): net.train() perm = torch.randperm(n, device=dev) for start in range(0, n, BATCH): idx = perm[start:start+BATCH] pred = net(xtr[idx]) task = ((pred - ytr[idx]) ** 2).mean() # Small subset keeps second-order autodiff affordable and fixed. take = idx[:min(16, len(idx))] margin, _, _ = curvature_penalty(net, xtr[take]) loss = task - BETA * margin opt.zero_grad(set_to_none=True) loss.backward() opt.step() net.eval() with torch.no_grad(): metric = float(((net(ds['xte'].to(dev)) - ds['yte'].to(dev)) ** 2).mean()) return metric, net, dev except RuntimeError: if dev == 'cpu': raise torch.cuda.empty_cache() raise RuntimeError('training failed') def idea_fn(cfg): return lambda seed: idea_run(seed, cfg['lr'])[0] def measured_signature(base_lr, idea_lr): # Re-train one paired seed and measure the trained systems' curvature and # perturbation response using the same validation inputs. seed = 0 seed_all(seed) ds = get_dataset('dynamics', seed, n_train=300, n_test=150) base = make_model('rnn_small', ds['input_shape'], ds['out_dim']) base, _, _ = train_model(base, ds, epochs=EPOCHS, lr=base_lr, batch=BATCH, log=lambda *_: None) imetric, idea, idev = idea_run(seed, idea_lr) bdev = next(base.parameters()).device xb = ds['xte'][:16].to(bdev) xi = ds['xte'][:16].to(idev) _, Hb, eb = curvature_penalty(base, xb) _, Hi, ei = curvature_penalty(idea, xi) # Parameter/input perturbation is observed model behaviour; compare output # displacement under a fixed small perturbation of the input window. eps = 1e-3 direction = torch.randn_like(xb) direction = direction / direction.norm() with torch.no_grad(): db = (base(xb + eps * direction) - base(xb)).norm().item() di = (idea(xi + eps * direction.to(idev)) - idea(xi)).norm().item() return { 'baseline_min_curvature': float(eb[0]), 'idea_min_curvature': float(ei[0]), 'baseline_trace_curvature': float(torch.trace(Hb)), 'idea_trace_curvature': float(torch.trace(Hi)), 'baseline_observed_perturbation': db, 'idea_observed_perturbation': di, 'predicted_inverse_margin_ratio': float(eb[0] / max(ei[0], 1e-12)), 'observed_perturbation_ratio': float(di / max(db, 1e-12)), 'confirmed': bool((ei[0] > eb[0]) and (di < db)) } def main(): grid = [{'lr': x} for x in LRS] base = sweep_baseline(baseline_fn, grid, seeds=(0, 1, 2, 3)) best_lr = float(base['best_cfg']['lr']) idea_grid = [{'lr': best_lr}, {'lr': LRS[max(0, LRS.index(best_lr)-1)]}, {'lr': LRS[min(len(LRS)-1, LRS.index(best_lr)+1)]}] idea_trials = [] for cfg in idea_grid: vals = [float(idea_fn(cfg)(s)) for s in SEEDS] idea_trials.append({'cfg': cfg, 'result': {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'n': len(vals)}}) chosen = min(idea_trials, key=lambda z: z['result']['mean']) report = make_report('dynamics', 'rnn_small', base, chosen['result'], {'mechanism_signature': measured_signature(best_lr, chosen['cfg']['lr']), 'idea_trials': idea_trials, 'structural_match': 'dynamics: recurrent actuated pendulum and stability/control task', 'intervention': 'soft-min of input Jacobian Gauss-Newton curvature plus fixed prior'}) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()