import os, sys, json, math import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) EPOCHS = 12 NTRAIN, NTEST = 800, 300 BATCH = 128 # Union is shared: all idea learning rates are also baseline-tested. LRS = (1e-3, 3e-3, 6e-3) WDS = (0.0, 1e-4) IDEA_CFGS = tuple((lr, c) for lr, c in zip(LRS, (0.005, 0.01, 0.02))) def pk_nll(h, xhat, eps=1e-6): """Poisson angular NLL after projecting recurrent states to C.""" z = h[..., :2] u = z / (torch.linalg.vector_norm(z, dim=-1, keepdim=True) + eps) r2 = (xhat * xhat).sum().clamp(max=1.0 - eps) den = 1.0 + r2 - 2.0 * (u * xhat).sum(-1) return (-torch.log(((1.0-r2) / den.clamp_min(eps)).clamp_min(eps))).mean() def attractor(device, depth=8): # Recent random disk-affine maps T(z)=q z + a; probes are fixed and gradients stop. # q/a are deterministic functions of the seed supplied by the caller. g = torch.Generator(device=device); g.manual_seed(1729) q = 0.70 + 0.08 * torch.rand(depth, generator=g, device=device) a = 0.10 * torch.randn(depth, 2, generator=g, device=device) probes = torch.tensor([[-.55,.10],[.15,-.45],[.40,.20],[-.20,.50]], device=device) z = probes.clone() for k in range(depth): z = q[k] * z + a[k] return z.mean(0).clamp(-.85, .85).detach() def forward_with_hidden(net, x): seq = x.view(x.shape[0], -1, 3) try: _, h = net.rnn(seq) except RuntimeError: old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False try: _, h = net.rnn(seq) finally: torch.backends.cudnn.enabled = old return net.head(h[-1]), h[-1] def run(seed, lr, coef, idea, collect=False, weight_decay=0.0): torch.manual_seed(seed); np.random.seed(seed) ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") try: net = make_model("rnn_small", ds["input_shape"], ds["out_dim"]).to(device) opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay) x, y = ds["xtr"].to(device), ds["ytr"].to(device) xhat = attractor(device) for _ in range(EPOCHS): net.train(); perm = torch.randperm(len(x), device=device) for j in range(0, len(x), BATCH): ix = perm[j:j+BATCH] pred, h = forward_with_hidden(net, x[ix]) loss = F.mse_loss(pred, y[ix]) if idea: loss = loss + coef * pk_nll(h, xhat) opt.zero_grad(set_to_none=True); loss.backward(); opt.step() net.eval() with torch.no_grad(): pred, h = forward_with_hidden(net, ds["xte"].to(device)) metric = float(F.mse_loss(pred, ds["yte"].to(device))) if collect: z = h[:, :2]; u = z / (torch.linalg.vector_norm(z, dim=-1, keepdim=True)+1e-6) phases = torch.atan2(u[:,1], u[:,0]).cpu().numpy() # Histogram KL is measured from the trained model, not an identity. bins=32; counts,_=np.histogram(phases, bins=bins, range=(-np.pi,np.pi)) p=(counts+1e-5)/(counts.sum()+bins*1e-5) centers=-np.pi+(np.arange(bins)+.5)*2*np.pi/bins xx=xhat.cpu().numpy(); uu=np.stack([np.cos(centers),np.sin(centers)],1) r2=min(float((xx*xx).sum()),1-1e-6) q=(1-r2)/(1+r2-2*uu.dot(xx)); q=q/q.sum() kl_pk=float(np.sum(p*np.log(p/q))); kl_uni=float(np.sum(p*np.log(p/(1/bins)))) # A direct recurrent stability observable: pair distances at final state. pair=float(torch.pdist(h).mean().cpu()) return metric, {"pairwise_hidden_distance":pair,"phase_kl_to_Px":kl_pk,"phase_kl_to_uniform":kl_uni,"pk_nll":float(pk_nll(h,xhat).cpu())} return metric except Exception: if device.type == "cuda": torch.cuda.empty_cache() torch.cuda.is_available = lambda: False try: return run(seed, lr, coef, idea, collect, weight_decay) finally: torch.cuda.is_available = lambda: True raise def main(): def baseline_fn(cfg): lr, wd = cfg return lambda seed: run(seed, lr, 0.0, False, False, wd) base_grid = [(lr, wd) for lr in LRS for wd in WDS] base = sweep_baseline(baseline_fn, base_grid, seeds=(0,1,2,3)) best_lr = float(base["best_cfg"][0]) # Evaluate idea at best baseline lr and two nearby predeclared settings. idea_grid = [(best_lr, .005), (best_lr, .01), (best_lr, .02)] idea_cfg_results=[] for cfg in idea_grid: r=evaluate(lambda s, cfg=cfg: run(s,cfg[0],cfg[1],True), seeds=SEEDS) idea_cfg_results.append({"cfg":cfg,"result":r}) best_idea=min(idea_cfg_results,key=lambda a:a["result"]["mean"]) idea=best_idea["result"] # Signature uses fresh trained models on all paired seeds. b_sig=[run(s,best_lr,0,False,True,base["best_cfg"][1])[1] for s in SEEDS] i_sig=[run(s,best_idea["cfg"][0],best_idea["cfg"][1],True,True)[1] for s in SEEDS] def mean(k, arr): return float(np.mean([a[k] for a in arr])) signature={ "prediction":"PK training should reduce recurrent ensemble spread and phase KL to P_x relative to unregularized rnn_small", "observed_baseline":{"pairwise_hidden_distance":mean("pairwise_hidden_distance",b_sig),"phase_kl_to_Px":mean("phase_kl_to_Px",b_sig),"phase_kl_to_uniform":mean("phase_kl_to_uniform",b_sig)}, "observed_idea":{"pairwise_hidden_distance":mean("pairwise_hidden_distance",i_sig),"phase_kl_to_Px":mean("phase_kl_to_Px",i_sig),"phase_kl_to_uniform":mean("phase_kl_to_uniform",i_sig)}, "confirmed": mean("pairwise_hidden_distance",i_sig) < mean("pairwise_hidden_distance",b_sig) and mean("phase_kl_to_Px",i_sig) < mean("phase_kl_to_Px",b_sig) } report=make_report("dynamics","rnn_small",base,idea,{"signature":signature,"idea_cfg":best_idea["cfg"],"baseline_grid":base_grid,"idea_grid":idea_grid}) report["comparison"]["idea_cfg_sweep"]=idea_cfg_results with open("bench_report.json","w") as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__ == "__main__": main()