1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  8
  9TRACK, MODEL = "tabular", "mlp_tiny"
 10EPOCHS, BATCH = 8, 128
 11LRS = [0.001, 0.003, 0.006]  # identical union on baseline and idea sides
 12LAMBDAS = [0.02, 0.05, 0.10]
 13torch.set_num_threads(1)
 14
 15
 16def seed_all(s):
 17    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 18    if torch.cuda.is_available():
 19        try: torch.cuda.manual_seed_all(s)
 20        except Exception: pass
 21
 22
 23def mgs_numpy(Z, Y, d, eps=1e-6):
 24    qs, ys, res = [], [], []
 25    for z, y in zip(Z, Y):
 26        nz = np.linalg.norm(z)
 27        if nz < eps: continue
 28        u = z / nz; v = u.copy()
 29        for q in qs: v -= q * np.dot(q, v)
 30        nv = np.linalg.norm(v); res.append(nv)
 31        if nv > eps:
 32            qs.append(v / nv); ys.append(y / nz)
 33            if len(qs) == d: break
 34    if not qs: return None, None, res
 35    return np.asarray(qs, dtype=np.float32), np.asarray(ys, dtype=np.float32), res
 36
 37
 38def run(seed, lr, lam=0.0, signature=False):
 39    seed_all(seed)
 40    ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=200)
 41    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
 42    use_cuda = torch.cuda.is_available()
 43    device = "cuda" if use_cuda else "cpu"
 44    try:
 45        net.to(device)
 46        x, y = ds["xtr"].to(device), ds["ytr"].to(device)
 47        xt, yt = ds["xte"].to(device), ds["yte"].to(device)
 48        opt = torch.optim.Adam(net.parameters(), lr=lr)
 49        lossf = nn.MSELoss()
 50        head = [m for m in net.modules() if isinstance(m, nn.Linear)][-1]
 51        memz, memy, Q = [], [], None
 52        ranks, orth, eigs = [], [], []
 53        for ep in range(EPOCHS):
 54            net.train(); perm = torch.randperm(len(x), device=device)
 55            for st in range(0, len(x), BATCH):
 56                ix = perm[st:st+BATCH]; loss = lossf(net(x[ix]), y[ix])
 57                opt.zero_grad(); loss.backward(); opt.step()
 58                if lam:
 59                    with torch.no_grad():
 60                        h = x[ix]
 61                        for layer in list(net.children())[:-1]: h = layer(h)
 62                        memz.append(h.detach().cpu().numpy()); memy.append(y[ix].detach().cpu().numpy())
 63            if lam:
 64                # Refresh once per epoch, using at most 256 recent examples.
 65                Z = np.concatenate(memz)[-256:]; Y = np.concatenate(memy)[-256:]
 66                q, qy, _ = mgs_numpy(Z, Y, head.in_features)
 67                if q is not None:
 68                    Q = torch.from_numpy(q).to(device); Qy = torch.from_numpy(qy).to(device)
 69                    with torch.no_grad():
 70                        old = head.weight.detach().clone()
 71                        residual = Qy - Q @ head.weight.T
 72                        head.weight.sub_(lr * lam * (residual.T @ Q))
 73                        G = Q.T @ Q
 74                        ranks.append(int(Q.shape[1])); orth.append(float(torch.max(torch.abs(G-torch.eye(G.shape[0],device=device))).cpu()))
 75                        eigs.append(float(torch.linalg.eigvalsh(G).max().cpu()))
 76        with torch.no_grad(): metric = float(torch.mean((net(xt)-yt)**2).cpu())
 77        sig = None
 78        if signature and eigs:
 79            sig = {"feature_dim": head.in_features, "observed_rank_last": ranks[-1],
 80                   "predicted_identity_max_eigen_mean": float(np.mean(eigs)),
 81                   "observed_QtQ_error_max": max(orth),
 82                   "prediction": "MGS memory predicts identity Gramian after d independent directions",
 83                   "confirmed": bool(ranks[-1] == head.in_features and max(orth) < 1e-4 and abs(np.mean(eigs)-1)<1e-4)}
 84        return metric, sig
 85    except RuntimeError:
 86        if use_cuda:
 87            torch.cuda.empty_cache()
 88            # Explicit CPU retry, avoiding recursive CUDA detection.
 89            torch.cuda.is_available = lambda: False
 90            try: return run(seed, lr, lam, signature)
 91            finally: torch.cuda.is_available = lambda: True
 92        raise
 93
 94
 95def fn(cfg):
 96    return lambda s: run(s, cfg["lr"], cfg.get("lambda", 0.0))[0]
 97
 98
 99def main():
100    base = sweep_baseline(fn, [{"lr":lr, "weight_decay":0.0} for lr in LRS])
101    ideas = []
102    for lam in LAMBDAS:
103        cfg = {"lr": base["best_cfg"]["lr"], "lambda": lam}
104        ideas.append({"cfg":cfg, "result":evaluate(fn(cfg))})
105    best = min(ideas, key=lambda z:z["result"]["mean"])
106    _, sig = run(0, best["cfg"]["lr"], best["cfg"]["lambda"], True)
107    report = make_report(TRACK, MODEL, base, best["result"], {
108        "prediction": sig, "idea_configs": ideas,
109        "confirmed": bool(sig and sig["confirmed"])
110    })
111    report["track_justification"] = "Tabular is the prescribed optimizer track; Friedman regression with mlp_tiny has a trainable final linear head and fixed feature dimension, matching the memory correction."
112    Path("bench_report.json").write_text(json.dumps(report, indent=2))
113    print(json.dumps(report, indent=2))
114
115if __name__ == "__main__": main()