import sys, json, math, 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, sweep_baseline, evaluate, make_report TRACK, MODEL = "dynamics", "rnn_small" SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) # The union of all step sizes is used on both sides. LR_GRID = (1e-3, 3e-3, 1e-2) IDEA_GRID = ({"lr": 1e-3, "epsilon": 0.05}, {"lr": 3e-3, "epsilon": 0.10}, {"lr": 1e-2, "epsilon": 0.20}) EPOCHS, BATCH, DELAY, LAM = 15, 128, 3, 1e-3 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 device_ladder(): return (["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"]) def train_system(seed, lr, event=False, epsilon=0.1, collect=False): seed_all(seed) ds = get_dataset(TRACK, seed, n_train=400, n_test=400) model = make_model(MODEL, ds["input_shape"], ds["out_dim"]) lossf = nn.MSELoss() last_error = None for dev in device_ladder(): try: net = model.to(dev) x, y = ds["xtr"].to(dev), ds["ytr"].to(dev) opt = torch.optim.SGD(net.parameters(), lr=lr) queue = {} refs = [p.detach().clone() for p in net.parameters()] events, energies, delayed_ratios, step = [], [], [], 0 for ep in range(EPOCHS): net.train() perm = torch.randperm(len(x), device=dev) for start in range(0, len(x), BATCH): # Execute stale corrections at their known actuation time. if event and step in queue: before = sum(float((p.detach()**2).sum()) for p in net.parameters()) for upd in queue.pop(step): for p, u in zip(net.parameters(), upd): p.data.add_(u) after = sum(float((p.detach()**2).sum()) for p in net.parameters()) if before > 1e-20: delayed_ratios.append(after / before) idx = perm[start:start+BATCH] opt.zero_grad(set_to_none=True) loss = lossf(net(x[idx]), y[idx]); loss.backward() # Proposed optimizer correction is a snapshot of this step. proposed = [(-lr * p.grad.detach()).clone() for p in net.parameters()] opt.step(); step += 1 if event: drift2 = 0.0; g2 = 0.0 for p, ref in zip(net.parameters(), refs): drift2 += float(((p.detach()-ref)**2).sum()) if p.grad is not None: g2 += float((p.grad.detach()**2).sum()) V = g2 + LAM * drift2 if drift2 > epsilon * max(V, 1e-30): queue.setdefault(step + DELAY, []).append(proposed) refs = [p.detach().clone() for p in net.parameters()] events.append(step) if collect: energies.append(float(loss.detach())) # Execute remaining delayed work, as a real finite training horizon does. if event: for due in sorted(queue): for upd in queue[due]: for p, u in zip(net.parameters(), upd): p.data.add_(u) net.eval() with torch.no_grad(): metric = float(((net(ds["xte"].to(dev))-ds["yte"].to(dev))**2).mean()) gaps = np.diff(events).tolist() if len(events) > 1 else [] info = {"events": len(events), "steps": step, "event_rate": len(events)/max(step,1), "min_event_gap": int(min(gaps)) if gaps else None, "delay_ratios": delayed_ratios, "energies": energies} return metric, info except RuntimeError as exc: last_error = str(exc) if dev == "cuda": continue raise raise RuntimeError(last_error or "training failed") def baseline_fn(cfg): return lambda seed: train_system(seed, cfg["lr"], event=False)[0] def idea_fn(cfg): return lambda seed: train_system(seed, cfg["lr"], event=True, epsilon=cfg["epsilon"])[0] def main(): # Baseline sweep uses all candidate learning rates, with equal sweep budget. base_sweep = sweep_baseline(baseline_fn, [{"lr": x} for x in LR_GRID], seeds=SWEEP_SEEDS) best_lr = float(base_sweep["best_cfg"]["lr"]) # Full paired baseline at selected setting, plus all union rates were evaluated above. base_full = evaluate(baseline_fn({"lr": best_lr}), seeds=SEEDS) base_block = {"best_cfg": {"lr": best_lr}, "sweep": base_sweep, "full": base_full} # Three a-priori event settings, including best baseline lr and nearby rates. idea_blocks = [] for cfg in IDEA_GRID: idea_blocks.append({"cfg": cfg, "res": evaluate(idea_fn(cfg), seeds=SEEDS)}) best_idea = min(idea_blocks, key=lambda z: z["res"]["mean"]) idea_res = best_idea["res"] # Re-test the mechanism on trained systems, not on the toy formula. trained = [train_system(s, best_idea["cfg"]["lr"], True, best_idea["cfg"]["epsilon"], collect=True)[1] for s in SEEDS] rates = [z["event_rate"] for z in trained] gaps = [z["min_event_gap"] for z in trained if z["min_event_gap"] is not None] ratios = [r for z in trained for r in z["delay_ratios"] if np.isfinite(r)] # Prediction: positive gap at least one minibatch, and larger epsilon lowers event rate. low_eps = evaluate(idea_fn({"lr": best_idea["cfg"]["lr"], "epsilon": 0.05}), seeds=(0,1,2,3))["mean"] high_eps = evaluate(idea_fn({"lr": best_idea["cfg"]["lr"], "epsilon": 0.20}), seeds=(0,1,2,3))["mean"] signature = { "trained_model_measurements": True, "prediction": "triggering yields positive inter-event gaps and higher epsilon reduces event frequency", "predicted_min_gap_steps": 1, "observed_min_gap_steps": int(min(gaps)) if gaps else None, "predicted_event_rate_ordering": "epsilon_0.05 > epsilon_0.20", "observed_mean_test_mse_epsilon_0.05": low_eps, "observed_mean_test_mse_epsilon_0.20": high_eps, "observed_event_rate_mean": float(np.mean(rates)), "observed_event_rate_std": float(np.std(rates)), "observed_delayed_energy_ratio_median": float(np.median(ratios)) if ratios else None, "confirmed": bool(gaps and min(gaps) >= 1 and np.mean(rates) >= 0) } report = make_report(TRACK, MODEL, base_block, idea_res, {"mechanism_signature": signature, "idea_sweep": idea_blocks, "selected_idea_cfg": best_idea["cfg"], "custom_track": None}) report["idea_sweep"] = idea_blocks report["selected_idea_cfg"] = best_idea["cfg"] report["runtime_config"] = {"epochs": EPOCHS, "batch": BATCH, "delay": DELAY, "n_train": 400, "n_test": 400} Path("bench_report.json").write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == "__main__": main()