1import json, random, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12LR_GRID = [1e-3, 3e-3, 1e-2]
 13EPOCHS = 6
 14BATCH = 128
 15M = 4
 16ALPHA = 0.10
 17LAMBDA = 0.002
 18BETA = 0.90
 19
 20
 21def seed_all(seed):
 22    random.seed(seed)
 23    np.random.seed(seed)
 24    torch.manual_seed(seed)
 25    if torch.cuda.is_available():
 26        torch.cuda.manual_seed_all(seed)
 27
 28
 29def train_one(seed, lr, regularized, collect=False):
 30    seed_all(seed)
 31    ds = get_dataset("dynamics", seed, n_train=800, n_test=300)
 32    net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 33    dev = "cuda" if torch.cuda.is_available() else "cpu"
 34    try:
 35        return _train(net, ds, lr, regularized, dev, collect)
 36    except Exception:
 37        if dev == "cuda":
 38            torch.cuda.empty_cache()
 39            return _train(net.cpu(), ds, lr, regularized, "cpu", collect)
 40        raise
 41
 42
 43def _train(net, ds, lr, regularized, dev, collect=False):
 44    net = net.to(dev)
 45    x, y = ds["xtr"].to(dev), ds["ytr"].to(dev)
 46    xt, yt = ds["xte"].to(dev), ds["yte"].to(dev)
 47    opt = torch.optim.Adam(net.parameters(), lr=lr)
 48    lossf = nn.MSELoss()
 49    queue = []
 50    eps = 0.0
 51    violations, ratios, grad_norms, losses = [], [], [], []
 52    for _ in range(EPOCHS):
 53        net.train()
 54        perm = torch.randperm(len(x), device=dev)
 55        for i in range(0, len(x), BATCH):
 56            idx = perm[i:i + BATCH]
 57            pred = net(x[idx])
 58            task = lossf(pred, y[idx])
 59            grads = torch.autograd.grad(task, tuple(net.parameters()),
 60                                        create_graph=regularized, retain_graph=True)
 61            flat = torch.cat([g.reshape(-1) for g in grads])
 62            # Gradient certificate q_k and V_k = ||q_k||^2/2.
 63            v_now = 0.5 * (flat * flat).sum()
 64            penalty = torch.zeros((), device=dev)
 65            if regularized and len(queue) >= M:
 66                v_old = queue[-M]
 67                delta = v_now - v_old + ALPHA * v_old
 68                eps = BETA * eps + (1.0 - BETA) * float(delta.detach().abs().cpu())
 69                residual = delta - float(eps)
 70                # Normalization and clipping prevent certificate scale domination.
 71                penalty = LAMBDA * torch.relu(residual).clamp(max=10.0).pow(2) / (1.0 + v_old)
 72                violations.append(float((residual.detach() > 0).cpu()))
 73                ratios.append(float((v_now.detach() / (v_old + 1e-12)).cpu()))
 74            total = task + penalty if regularized else task
 75            opt.zero_grad(set_to_none=True)
 76            total.backward()
 77            opt.step()
 78            queue.append(v_now.detach())
 79            grad_norms.append(float(flat.detach().norm().cpu()))
 80            losses.append(float(task.detach().cpu()))
 81    net.eval()
 82    with torch.no_grad():
 83        metric = float(((net(xt) - yt) ** 2).mean().cpu())
 84    if not collect:
 85        return metric
 86    return {
 87        "metric": metric,
 88        "model": net,
 89        "signature": {
 90            "violation_rate": float(np.mean(violations)) if violations else 0.0,
 91            "mean_v_ratio": float(np.mean(ratios)) if ratios else float("nan"),
 92            "tail_grad_std": float(np.std(grad_norms[-100:])) if grad_norms else float("nan"),
 93            "tail_loss_std": float(np.std(losses[-100:])) if losses else float("nan"),
 94        },
 95    }
 96
 97
 98def metric_fn(cfg, regularized):
 99    lr = float(cfg["lr"])
100    return lambda seed: train_one(seed, lr, regularized, False)
101
102
103def main():
104    grid = [{"lr": lr} for lr in LR_GRID]
105    base = sweep_baseline(lambda cfg: metric_fn(cfg, False), grid, seeds=SWEEP_SEEDS)
106    # The same union of lrs is used for the idea-side three-configuration sweep.
107    idea_trials = []
108    for cfg in grid:
109        r = evaluate(metric_fn(cfg, True), seeds=SWEEP_SEEDS)
110        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
111    best_idea_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"]
112    idea = evaluate(metric_fn(best_idea_cfg, True), seeds=SEEDS)
113    report = make_report(
114        "dynamics", "rnn_small", base, idea,
115        extra={
116            "prediction": "finite-horizon gradient energy should contract on average and unstable training should have positive residual violations",
117            "trained_model_measurements": {
118                "idea_cfg": best_idea_cfg,
119                "idea_signature_per_seed": [train_one(s, best_idea_cfg["lr"], True, True)["signature"] for s in SEEDS],
120            },
121            "confirmed": False,
122            "idea_sweep": idea_trials,
123            "certificate": {"M": M, "alpha": ALPHA, "lambda": LAMBDA, "beta": BETA},
124        },
125    )
126    # Add a compact quantitative signature summary, measured from trained models.
127    sigs = report["mechanism_signature"]["trained_model_measurements"]["idea_signature_per_seed"]
128    report["mechanism_signature"]["observed_mean_v_ratio"] = float(np.nanmean([s["mean_v_ratio"] for s in sigs]))
129    report["mechanism_signature"]["observed_violation_rate"] = float(np.mean([s["violation_rate"] for s in sigs]))
130    ratio = report["mechanism_signature"]["observed_mean_v_ratio"]
131    report["mechanism_signature"]["confirmed"] = bool(np.isfinite(ratio) and ratio < 1.0)
132    Path("bench_report.json").write_text(json.dumps(report, indent=2))
133    print(json.dumps(report, indent=2))
134
135
136if __name__ == "__main__":
137    main()