import sys, os, json, math, random, copy, time 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, train_model, evaluate, sweep_baseline, make_report) SEEDS = tuple(range(8)) # Shared union of learning rates: every idea lr is also a baseline lr. GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}] EPOCHS = 8 BATCH = 128 NTRAIN, NTEST = 400, 200 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) torch.set_num_threads(4) def dataset(seed): return get_dataset("tabular", int(seed), n_train=NTRAIN, n_test=NTEST) def baseline_fn(cfg): def run(seed): seed_all(seed) d = dataset(seed) seed_all(seed + 10000) net = make_model("mlp_tiny", d["input_shape"], d["out_dim"]) _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, weight_decay=0.0, log=lambda *_: None) return float(metric) return run def loss_fn(net, x, y): return ((net(x) - y) ** 2).mean() def apply_step(net, grads, eta): with torch.no_grad(): for p, g in zip(net.parameters(), grads): p.add_(g, alpha=-float(eta)) def controller_train(net, d, initial_lr, alpha=1.0, rho=0.9, c=0.1, beta=0.5, eta_min=1e-5, eta_max=1.0, max_bt=5): # This is deliberately a training loop intervention: probe and candidate # losses are evaluated on the same minibatch as F0. device = next(net.parameters()).device xtr, ytr = d["xtr"].to(device), d["ytr"].to(device) eta_probe = float(initial_lr) accepted = 0; total = 0; observed_lh = []; chosen = [] for ep in range(EPOCHS): net.train() perm = torch.randperm(len(xtr), device=device) for i in range(0, len(xtr), BATCH): idx = perm[i:i+BATCH]; x, y = xtr[idx], ytr[idx] net.zero_grad(set_to_none=True) f0 = loss_fn(net, x, y) grads = torch.autograd.grad(f0, tuple(net.parameters())) g2 = sum(float((g.detach() ** 2).sum()) for g in grads) G = math.sqrt(max(g2, 0.0)) total += 1 if G < 1e-12: continue # Probe at maintained rate, then restore exactly. apply_step(net, grads, eta_probe) with torch.no_grad(): fp = float(loss_fn(net, x, y)) apply_step(net, grads, -eta_probe) snorm = eta_probe * G rem = fp - float(f0.detach()) + eta_probe * g2 lhat = (1.0 + alpha) * max(rem, 0.0) / (snorm ** (1.0 + alpha) + 1e-12) observed_lh.append(lhat) eta_star = (G ** (1.0-alpha) / (lhat + 1e-12)) ** (1.0/alpha) eta = min(eta_max, max(eta_min, rho * eta_star)) ok = False fchosen = None for _ in range(max_bt + 1): apply_step(net, grads, eta) with torch.no_grad(): fc = float(loss_fn(net, x, y)) if fc <= float(f0.detach()) - c * eta * g2 + 1e-10: ok = True; fchosen = fc; break apply_step(net, grads, -eta) eta *= beta if not ok: # candidate is currently applied on final failed attempt apply_step(net, grads, -eta) eta_probe = max(eta_min, eta_probe * beta) else: accepted += 1 chosen.append(eta) # next probe gently grows from accepted scale eta_probe = min(eta_max, eta / 0.9) net.eval() with torch.no_grad(): xte, yte = d["xte"].to(device), d["yte"].to(device) metric = float(loss_fn(net, xte, yte)) stats = {"accepted_fraction": accepted / max(total, 1), "mean_lhat": float(np.mean(observed_lh)) if observed_lh else 0.0, "mean_accepted_eta": float(np.mean(chosen)) if chosen else 0.0, "n_updates": total} return metric, stats def idea_fn(cfg, collect=False): def run(seed): seed_all(seed) d = dataset(seed) seed_all(seed + 10000) net = make_model("mlp_tiny", d["input_shape"], d["out_dim"]) # Use official fallback placement convention by trying CUDA, then CPU. device = "cuda" if torch.cuda.is_available() else "cpu" try: net = net.to(device) metric, stats = controller_train(net, d, cfg["lr"]) except RuntimeError: net = make_model("mlp_tiny", d["input_shape"], d["out_dim"]).to("cpu") metric, stats = controller_train(net, d, cfg["lr"]) if collect: RUN_STATS.setdefault(str(cfg["lr"]), {})[str(seed)] = stats return float(metric) return run RUN_STATS = {} if __name__ == "__main__": t0 = time.time() base = sweep_baseline(baseline_fn, GRID) # Idea sweep uses exactly the same three learning rates; select on sweep seeds. idea_sweep = [] for cfg in GRID: r = evaluate(idea_fn(cfg), seeds=(0,1,2,3)) idea_sweep.append({"cfg": cfg, "mean": r["mean"]}) best_cfg = min(GRID, key=lambda c: next(x["mean"] for x in idea_sweep if x["cfg"] == c)) idea_full = evaluate(idea_fn(best_cfg, collect=True), seeds=SEEDS) # Re-test the NN-scale quantitative prediction: observed Lhat should predict # the accepted scale eta approximately as rho/Lhat for alpha=1. vals = [] for st in RUN_STATS.get(str(best_cfg["lr"]), {}).values(): if st["mean_lhat"] > 0 and st["mean_accepted_eta"] > 0: pred = 0.9 / st["mean_lhat"] obs = st["mean_accepted_eta"] vals.append({"predicted_eta": pred, "observed_eta": obs, "ratio_observed_to_predicted": obs/pred}) ratios = [v["ratio_observed_to_predicted"] for v in vals] signature = {"alpha": 1.0, "rho": 0.9, "prediction": "accepted_eta ~= rho / observed_directional_Lhat", "per_seed": vals, "mean_ratio": float(np.mean(ratios)) if ratios else None, "confirmed": bool(ratios and 0.5 <= float(np.mean(ratios)) <= 1.5)} base["idea_parity_sweep"] = idea_sweep report = make_report("tabular", "mlp_tiny", base, idea_full, signature) report["idea"]["selected_cfg"] = best_cfg report["runtime_seconds"] = time.time() - t0 report["protocol_notes"] = {"epochs": EPOCHS, "batch": BATCH, "dataset_sizes": [NTRAIN, NTEST], "structural_match": "optimizer -> tabular"} with open("bench_report.json", "w") as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2))