import sys, json, random 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, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 12 BATCH = 128 # Union of learning rates is used on both sides. Adam weight decay is also # swept for the baseline's central optimization knob. LRS = (1e-3, 3e-3, 1e-2) WDS = (0.0, 1e-4) IDEA_LAMBDA = 0.02 TARGET = 0.90 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 ds_for(seed): return get_dataset("dynamics", seed=seed, n_train=400, n_test=400) def baseline_fn(cfg): def run(seed): seed_all(seed) ds = ds_for(seed) model = make_model("rnn_small", tuple(ds["xtr"].shape[1:]), 1) _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, weight_decay=cfg["weight_decay"], log=lambda *_: None) return float(metric) if metric is not None else float("inf") return run def recurrent_contraction_penalty(model, x, target=TARGET): """Penalty on the local hidden transition expansion of the trained GRU. The hidden state before the final observed control input is computed, then two deterministic one-step GRU transitions are compared under a small hidden perturbation. Only expansion above target is penalized, preserving useful contracting dynamics instead of forcing every transition to zero. """ seq = x.view(x.shape[0], -1, 3) with torch.no_grad(): _, h0 = model.rnn(seq[:, :-1]) h0 = h0.detach().requires_grad_(True) last = seq[:, -1:, :] out, _ = model.rnn(last, h0) eps = torch.randn_like(h0) * 1e-3 outp, _ = model.rnn(last, h0 + eps) # Batch-averaged local gain estimate; denominator avoids scale artifacts. gain = (outp - out).pow(2).mean(dim=(0, 2)).sqrt() / (eps.pow(2).mean(dim=(0, 2)).sqrt() + 1e-8) return torch.relu(gain - target).pow(2).mean() def idea_train(model, ds, epochs, lr, weight_decay, lam): errs = [] ladder = [("cuda", False)] if torch.cuda.is_available() else [] ladder += [("cpu", False)] for device, _ in ladder: try: net = model.to(device) xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device) opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay) lossf = nn.MSELoss() for _ in range(epochs): net.train(); perm = torch.randperm(len(xtr), device=device) for i in range(0, len(xtr), BATCH): idx = perm[i:i+BATCH]; xb, yb = xtr[idx], ytr[idx] pred = net(xb) loss = lossf(pred, yb) + lam * recurrent_contraction_penalty(net, xb) opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0); opt.step() net.eval() with torch.no_grad(): metric = float(((net(ds["xte"].to(device)) - ds["yte"].to(device))**2).mean()) return net, metric except RuntimeError as e: errs.append(str(e)); if device == "cuda": torch.cuda.empty_cache() return None, float("inf") def idea_fn(cfg): def run(seed): seed_all(seed); ds = ds_for(seed) model = make_model("rnn_small", tuple(ds["xtr"].shape[1:]), 1) _, metric = idea_train(model, ds, EPOCHS, cfg["lr"], cfg["weight_decay"], cfg["lambda"]) return metric return run def trained_signature(seed, cfg, idea): seed_all(seed); ds = ds_for(seed) model = make_model("rnn_small", tuple(ds["xtr"].shape[1:]), 1) if idea: model, _ = idea_train(model, ds, EPOCHS, cfg["lr"], cfg["weight_decay"], cfg["lambda"]) else: model, _, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, weight_decay=cfg["weight_decay"], log=lambda *_: None) # Signature probing is deliberately CPU-only: it is a small diagnostic and # avoids consuming the shared GPU allocator after the paired runs. model = model.cpu() device = torch.device("cpu") x = ds["xte"][:128].to(device) seq = x.view(x.shape[0], -1, 3) with torch.no_grad(): _, h = model.rnn(seq[:, :-1]) eps = torch.randn_like(h) * 1e-3 with torch.no_grad(): a, _ = model.rnn(seq[:, -1:], h); b, _ = model.rnn(seq[:, -1:], h + eps) gain = float(((b-a).pow(2).mean().sqrt() / (eps.pow(2).mean().sqrt()+1e-8)).cpu()) return gain def main(): baseline_grid = [{"lr": lr, "weight_decay": wd} for lr in LRS for wd in WDS] base = sweep_baseline(baseline_fn, baseline_grid, seeds=SWEEP_SEEDS) best = base["best_cfg"] # Three idea settings: best baseline lr and both adjacent rates, all also # present in baseline_grid (same shared architecture and optimizer family). idea_grid = [ {"lr": 1e-3, "weight_decay": best["weight_decay"], "lambda": IDEA_LAMBDA}, {"lr": 3e-3, "weight_decay": best["weight_decay"], "lambda": IDEA_LAMBDA}, {"lr": 1e-2, "weight_decay": best["weight_decay"], "lambda": IDEA_LAMBDA}, ] idea_runs = [] for cfg in idea_grid: r = __import__('bench').evaluate(idea_fn(cfg), seeds=SEEDS) idea_runs.append({"cfg": cfg, "result": r}) chosen = min(idea_runs, key=lambda z: z["result"]["mean"]) rep = make_report("dynamics", "rnn_small", base, chosen["result"], extra={}) sig_b = trained_signature(0, best, False); sig_i = trained_signature(0, chosen["cfg"], True) # The quantitative stage-1 prediction is contraction/stability; this is # measured on trained systems, not analytically asserted. rep["mechanism_signature"] = { "prediction": "cycle-stabilizing intervention lowers local recurrent hidden-state gain", "baseline_hidden_gain": sig_b, "idea_hidden_gain": sig_i, "relative_gain_change": (sig_i-sig_b)/(abs(sig_b)+1e-8), "confirmed": bool(sig_i < sig_b), "measurement": "finite 1e-3 hidden perturbation through trained GRU final-step transition" } rep["idea_sweep"] = idea_runs rep["protocol_notes"] = {"epochs": EPOCHS, "n_train": 400, "n_test": 400, "baseline_grid": baseline_grid, "idea_grid": idea_grid} Path("bench_report.json").write_text(json.dumps(rep, indent=2)) print(json.dumps(rep, indent=2)) if __name__ == "__main__": main()