1import sys, json, math, random
  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, sweep_baseline, evaluate, make_report
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12EPOCHS = 30
 13BATCH = 128
 14PEAK_LR = 0.003
 15# Union is used on both sides: lr and nominal decay are method knobs for AdamW.
 16GRID = [
 17    {"lr": 0.0015, "weight_decay": 0.001},
 18    {"lr": 0.0015, "weight_decay": 0.01},
 19    {"lr": 0.0030, "weight_decay": 0.001},
 20    {"lr": 0.0030, "weight_decay": 0.01},
 21    {"lr": 0.0060, "weight_decay": 0.001},
 22    {"lr": 0.0060, "weight_decay": 0.01},
 23]
 24
 25
 26def seed_all(seed):
 27    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 28    if torch.cuda.is_available():
 29        torch.cuda.manual_seed_all(seed)
 30
 31
 32def lr_at(epoch, epochs, peak):
 33    # Same warmup/cooldown schedule for both systems.
 34    warm = max(1, epochs // 5)
 35    if epoch < warm:
 36        return peak * (epoch + 1) / warm
 37    z = (epoch - warm) / max(1, epochs - warm - 1)
 38    return peak * 0.5 * (1.0 + math.cos(math.pi * z))
 39
 40
 41def train_one(seed, cfg, scaled, collect=False):
 42    seed_all(seed)
 43    ds = get_dataset("tabular", seed, n_train=400, n_test=200)
 44    model = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"])
 45    # Explicit device ladder mirrors bench.train_model's robust fallback.
 46    devices = ["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"]
 47    last_err = None
 48    for dev in devices:
 49        try:
 50            net = model.to(dev)
 51            x, y = ds["xtr"].to(dev), ds["ytr"].to(dev)
 52            opt = torch.optim.AdamW(net.parameters(), lr=float(cfg["lr"]), weight_decay=0.0)
 53            lossf = nn.MSELoss()
 54            decay_logs = []
 55            for ep in range(EPOCHS):
 56                lr = lr_at(ep, EPOCHS, float(cfg["lr"]))
 57                perm = torch.randperm(len(x), device=dev)
 58                net.train()
 59                for i in range(0, len(x), BATCH):
 60                    ix = perm[i:i+BATCH]
 61                    loss = lossf(net(x[ix]), y[ix])
 62                    opt.zero_grad(set_to_none=True); loss.backward()
 63                    # AdamW moments/update are computed by the optimizer; to keep
 64                    # the intervention isolated, temporarily apply its raw step
 65                    # with zero decay, then apply our decoupled factor ourselves.
 66                    for group in opt.param_groups: group["lr"] = lr
 67                    opt.step()
 68                    # Correct the just-applied zero-decay update with decoupled decay.
 69                    frac = min(max(lr / float(cfg["lr"]), 0.0), 1.0)
 70                    lam_t = float(cfg["weight_decay"]) * (frac if scaled else 1.0)
 71                    factor = 1.0 - lr * lam_t
 72                    before = 0.0; after = 0.0
 73                    with torch.no_grad():
 74                        for p in net.parameters():
 75                            before += float((p.detach() ** 2).sum())
 76                            p.mul_(factor)
 77                            after += float((p.detach() ** 2).sum())
 78                    if collect:
 79                        # NN-scale observed norm multiplier immediately around
 80                        # the decay operation, compared with the predicted factor.
 81                        decay_logs.append((math.sqrt(after / max(before, 1e-30)), factor, frac))
 82            net.eval()
 83            with torch.no_grad():
 84                pred = net(ds["xte"].to(dev))
 85                metric = float(((pred - ds["yte"].to(dev)) ** 2).mean())
 86                norm = float(torch.sqrt(sum((p.detach() ** 2).sum() for p in net.parameters())))
 87            result = {"metric": metric, "final_norm": norm}
 88            if collect:
 89                result["decay_logs"] = decay_logs
 90            return result
 91        except RuntimeError as e:
 92            last_err = str(e)
 93            model = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"])
 94    raise RuntimeError(last_err or "training failed")
 95
 96
 97def metric_factory(scaled, cfg):
 98    return lambda seed: train_one(seed, cfg, scaled)["metric"]
 99
100
101def main():
102    # Baseline sweep uses the same six configurations that are available to idea.
103    base = sweep_baseline(lambda c: metric_factory(False, c), GRID, seeds=SWEEP_SEEDS)
104    # Equal-size idea-side sweep, then full paired evaluation at selected config.
105    idea_sweep = [{"cfg": c, "mean": evaluate(metric_factory(True, c), SWEEP_SEEDS)["mean"]} for c in GRID]
106    idea_cfg = min(idea_sweep, key=lambda z: z["mean"])["cfg"]
107    idea_eval = evaluate(metric_factory(True, idea_cfg), SEEDS)
108    # Collect behavior from the actual full paired trained systems.
109    observed = []
110    for s in SEEDS:
111        r = train_one(s, idea_cfg, True, collect=True)
112        observed.extend(r["decay_logs"])
113    obs_mult = float(np.mean([a for a, _, _ in observed]))
114    pred_mult = float(np.mean([b for _, b, _ in observed]))
115    # Also test the stage-1 claim on the cooldown: log shrinkage ratio follows lr fraction.
116    cooldown = [(a, b, f) for a, b, f in observed if f < 0.5 and abs(math.log(b)) > 1e-12 and abs(math.log(a)) > 1e-12]
117    # Float32 rounds the factor to one at the very end of cooldown; omit only
118    # those unresolvable observations from the logarithmic ratio statistic.
119    observed_ratio = float(np.mean([math.log(a) / math.log(b) for a,b,_ in cooldown])) if cooldown else float("nan")
120    predicted_ratio = float(np.mean([f for _,_,f in cooldown])) if cooldown else float("nan")
121    signature = {
122        "source": "trained mlp_tiny models on Friedman#1; decay boundaries instrumented during full paired runs",
123        "predicted_mean_decay_multiplier": pred_mult,
124        "observed_mean_decay_multiplier": obs_mult,
125        "multiplier_abs_error": abs(obs_mult - pred_mult),
126        "cooldown_predicted_log_ratio": predicted_ratio,
127        "cooldown_observed_log_ratio": observed_ratio,
128        "cooldown_ratio_abs_error": abs(observed_ratio - predicted_ratio),
129        "confirmed": abs(obs_mult-pred_mult) < 1e-6 and np.isfinite(observed_ratio) and abs(observed_ratio-predicted_ratio) < 1e-4,
130    }
131    # Include the idea sweep transparently while make_report supplies paired test.
132    idea = dict(idea_eval); idea["selected_cfg"] = idea_cfg; idea["sweep"] = idea_sweep
133    rep = make_report("tabular", "mlp_tiny", base, idea, signature)
134    rep["protocol_notes"] = {"epochs": EPOCHS, "batch": BATCH, "grid_union": GRID,
135                             "selection_seeds": list(SWEEP_SEEDS), "paired_seeds": list(SEEDS),
136                             "structural_match": "tabular is the built-in optimizer/regularizer track"}
137    Path("bench_report.json").write_text(json.dumps(rep, indent=2))
138    print(json.dumps(rep, indent=2))
139
140if __name__ == "__main__": main()