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 SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) EPOCHS = 30 BATCH = 128 PEAK_LR = 0.003 # Union is used on both sides: lr and nominal decay are method knobs for AdamW. GRID = [ {"lr": 0.0015, "weight_decay": 0.001}, {"lr": 0.0015, "weight_decay": 0.01}, {"lr": 0.0030, "weight_decay": 0.001}, {"lr": 0.0030, "weight_decay": 0.01}, {"lr": 0.0060, "weight_decay": 0.001}, {"lr": 0.0060, "weight_decay": 0.01}, ] 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 lr_at(epoch, epochs, peak): # Same warmup/cooldown schedule for both systems. warm = max(1, epochs // 5) if epoch < warm: return peak * (epoch + 1) / warm z = (epoch - warm) / max(1, epochs - warm - 1) return peak * 0.5 * (1.0 + math.cos(math.pi * z)) def train_one(seed, cfg, scaled, collect=False): seed_all(seed) ds = get_dataset("tabular", seed, n_train=400, n_test=200) model = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"]) # Explicit device ladder mirrors bench.train_model's robust fallback. devices = ["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"] last_err = None for dev in devices: try: net = model.to(dev) x, y = ds["xtr"].to(dev), ds["ytr"].to(dev) opt = torch.optim.AdamW(net.parameters(), lr=float(cfg["lr"]), weight_decay=0.0) lossf = nn.MSELoss() decay_logs = [] for ep in range(EPOCHS): lr = lr_at(ep, EPOCHS, float(cfg["lr"])) perm = torch.randperm(len(x), device=dev) net.train() for i in range(0, len(x), BATCH): ix = perm[i:i+BATCH] loss = lossf(net(x[ix]), y[ix]) opt.zero_grad(set_to_none=True); loss.backward() # AdamW moments/update are computed by the optimizer; to keep # the intervention isolated, temporarily apply its raw step # with zero decay, then apply our decoupled factor ourselves. for group in opt.param_groups: group["lr"] = lr opt.step() # Correct the just-applied zero-decay update with decoupled decay. frac = min(max(lr / float(cfg["lr"]), 0.0), 1.0) lam_t = float(cfg["weight_decay"]) * (frac if scaled else 1.0) factor = 1.0 - lr * lam_t before = 0.0; after = 0.0 with torch.no_grad(): for p in net.parameters(): before += float((p.detach() ** 2).sum()) p.mul_(factor) after += float((p.detach() ** 2).sum()) if collect: # NN-scale observed norm multiplier immediately around # the decay operation, compared with the predicted factor. decay_logs.append((math.sqrt(after / max(before, 1e-30)), factor, frac)) net.eval() with torch.no_grad(): pred = net(ds["xte"].to(dev)) metric = float(((pred - ds["yte"].to(dev)) ** 2).mean()) norm = float(torch.sqrt(sum((p.detach() ** 2).sum() for p in net.parameters()))) result = {"metric": metric, "final_norm": norm} if collect: result["decay_logs"] = decay_logs return result except RuntimeError as e: last_err = str(e) model = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"]) raise RuntimeError(last_err or "training failed") def metric_factory(scaled, cfg): return lambda seed: train_one(seed, cfg, scaled)["metric"] def main(): # Baseline sweep uses the same six configurations that are available to idea. base = sweep_baseline(lambda c: metric_factory(False, c), GRID, seeds=SWEEP_SEEDS) # Equal-size idea-side sweep, then full paired evaluation at selected config. idea_sweep = [{"cfg": c, "mean": evaluate(metric_factory(True, c), SWEEP_SEEDS)["mean"]} for c in GRID] idea_cfg = min(idea_sweep, key=lambda z: z["mean"])["cfg"] idea_eval = evaluate(metric_factory(True, idea_cfg), SEEDS) # Collect behavior from the actual full paired trained systems. observed = [] for s in SEEDS: r = train_one(s, idea_cfg, True, collect=True) observed.extend(r["decay_logs"]) obs_mult = float(np.mean([a for a, _, _ in observed])) pred_mult = float(np.mean([b for _, b, _ in observed])) # Also test the stage-1 claim on the cooldown: log shrinkage ratio follows lr fraction. 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] # Float32 rounds the factor to one at the very end of cooldown; omit only # those unresolvable observations from the logarithmic ratio statistic. observed_ratio = float(np.mean([math.log(a) / math.log(b) for a,b,_ in cooldown])) if cooldown else float("nan") predicted_ratio = float(np.mean([f for _,_,f in cooldown])) if cooldown else float("nan") signature = { "source": "trained mlp_tiny models on Friedman#1; decay boundaries instrumented during full paired runs", "predicted_mean_decay_multiplier": pred_mult, "observed_mean_decay_multiplier": obs_mult, "multiplier_abs_error": abs(obs_mult - pred_mult), "cooldown_predicted_log_ratio": predicted_ratio, "cooldown_observed_log_ratio": observed_ratio, "cooldown_ratio_abs_error": abs(observed_ratio - predicted_ratio), "confirmed": abs(obs_mult-pred_mult) < 1e-6 and np.isfinite(observed_ratio) and abs(observed_ratio-predicted_ratio) < 1e-4, } # Include the idea sweep transparently while make_report supplies paired test. idea = dict(idea_eval); idea["selected_cfg"] = idea_cfg; idea["sweep"] = idea_sweep rep = make_report("tabular", "mlp_tiny", base, idea, signature) rep["protocol_notes"] = {"epochs": EPOCHS, "batch": BATCH, "grid_union": GRID, "selection_seeds": list(SWEEP_SEEDS), "paired_seeds": list(SEEDS), "structural_match": "tabular is the built-in optimizer/regularizer track"} Path("bench_report.json").write_text(json.dumps(rep, indent=2)) print(json.dumps(rep, indent=2)) if __name__ == "__main__": main()