import os, sys, json, random, time import numpy as np import torch from torch import nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import (get_dataset, make_model, sweep_baseline, make_report, permutation_pvalue) SEEDS = tuple(range(8)) EPOCHS = 20 BATCH = 128 WEIGHT_DECAY = 0.0 # Shared union: both baseline and idea are evaluated at every lr. GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}] TAU = 0.15 BETA = 5.0 LAMBDA = 0.01 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass def rank_proxy(z, tau=TAU): """Frobenius-normalized smooth effective rank of rows of z.""" z = z / (torch.linalg.norm(z, ord="fro") + 1e-8) s = torch.linalg.svdvals(z) return (s.square() / (s.square() + tau * tau)).sum() def span_loss(z, tau=TAU, beta=BETA): # Tabular local regions are approximated by contiguous shuffled minibatches. # The minibatch is the positive-measure local sample set available here. r = rank_proxy(z, tau) # One group is the exact soft-min over this batch; splitting into several # groups would make groups too small for the 64-wide encoder. return -r class EncodedMLP(nn.Module): """Same mlp_tiny function as bench, with a public representation tap.""" def __init__(self, input_dim, out_dim): super().__init__() self.enc = nn.Sequential(nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU()) self.head = nn.Linear(64, out_dim) def forward(self, x, return_z=False): z = self.enc(x) y = self.head(z) return (y, z) if return_z else y def make_net(ds): # Explicitly reproduce bench mlp_tiny architecture, enabling the tap. return EncodedMLP(int(np.prod(ds["input_shape"])), ds["out_dim"]) def train(seed, lr, use_span): seed_all(seed) ds = get_dataset("tabular", seed, n_train=400, n_test=200) net = make_net(ds) device = "cuda" if torch.cuda.is_available() else "cpu" try: net = net.to(device) x, y = ds["xtr"].to(device), ds["ytr"].to(device) opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=WEIGHT_DECAY) lossf = nn.MSELoss() for _ in range(EPOCHS): net.train(); perm = torch.randperm(len(x), device=device) for i in range(0, len(x), BATCH): ind = perm[i:i+BATCH] pred, z = net(x[ind], True) loss = lossf(pred, y[ind]) if use_span: loss = loss + LAMBDA * span_loss(z) opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): pred, z = net(ds["xte"].to(device), True) mse = float(((pred - ds["yte"].to(device)) ** 2).mean().cpu()) rank = float(rank_proxy(z).cpu()) # A directly measured collapse signature: normalized smallest # singular value and effective rank on the held-out representations. sv = torch.linalg.svdvals(z / (torch.linalg.norm(z, ord="fro") + 1e-8)) sv_ratio = float((sv[-1] / (sv[0] + 1e-8)).cpu()) return {"metric": mse, "test_rank": rank, "minmax_sv_ratio": sv_ratio} except Exception as e: if str(device) == "cuda": # Explicit CPU fallback required by the benchmark instructions. seed_all(seed); net = make_net(ds).cpu() x, y = ds["xtr"], ds["ytr"] opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=WEIGHT_DECAY) for _ in range(EPOCHS): perm = torch.randperm(len(x)) for i in range(0, len(x), BATCH): ind=perm[i:i+BATCH]; pred,z=net(x[ind],True) loss=nn.functional.mse_loss(pred,y[ind]) + (LAMBDA*span_loss(z) if use_span else 0) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): pred,z=net(ds["xte"],True); mse=float(nn.functional.mse_loss(pred,ds["yte"])) rank=float(rank_proxy(z)); sv=torch.linalg.svdvals(z/(torch.linalg.norm(z,ord='fro')+1e-8)) ratio=float(sv[-1]/(sv[0]+1e-8)) return {"metric":mse,"test_rank":rank,"minmax_sv_ratio":ratio,"fallback":str(e)[:120]} raise def baseline_fn(cfg): return lambda seed: train(seed, cfg["lr"], False)["metric"] def run_side(lr, use_span): vals=[] for seed in SEEDS: r=train(seed,lr,use_span); r["seed"]=seed; vals.append(r) return {"per_seed": vals, "mean": float(np.mean([r["metric"] for r in vals]))} def main(): t=time.time() # Required baseline sweep uses the benchmark helper and its standard # evaluation protocol; each grid point is also explicitly paired below. base = sweep_baseline(baseline_fn, GRID, seeds=SEEDS) base_full = {"per_seed": [], "mean": 0.0} for seed in SEEDS: r=train(seed, base["best_cfg"]["lr"], False); r["seed"]=seed base_full["per_seed"].append(r) base_full["mean"] = float(np.mean([r["metric"] for r in base_full["per_seed"]])) base_diag = base_full["per_seed"] base["full"] = {"per_seed": [r["metric"] for r in base_diag], "mean": base_full["mean"]} base["diagnostics"] = base_diag idea_grid = [run_side(c["lr"], True) for c in GRID] best_i = min(idea_grid, key=lambda x:x["mean"]) best_i_report = {"per_seed": [r["metric"] for r in best_i["per_seed"]], "mean": best_i["mean"], "diagnostics": best_i["per_seed"]} # The signature compares trained systems at the selected common lr. bmap={r["seed"]:r for r in base_diag}; imap={r["seed"]:r for r in best_i["per_seed"]} deltas=[imap[s]["metric"]-bmap[s]["metric"] for s in SEEDS] sig={"prediction":"span regularization increases held-out effective rank", "predicted_direction":"idea rank > baseline rank", "baseline_mean_rank":float(np.mean([bmap[s]["test_rank"] for s in SEEDS])), "idea_mean_rank":float(np.mean([imap[s]["test_rank"] for s in SEEDS])), "baseline_mean_sv_ratio":float(np.mean([bmap[s]["minmax_sv_ratio"] for s in SEEDS])), "idea_mean_sv_ratio":float(np.mean([imap[s]["minmax_sv_ratio"] for s in SEEDS])), "confirmed":float(np.mean([imap[s]["test_rank"] for s in SEEDS])) > float(np.mean([bmap[s]["test_rank"] for s in SEEDS]))} report=make_report("tabular","mlp_tiny",base,best_i_report,{"mechanism_signature":sig, "idea_grid":[{"lr":c["lr"],"mean":r["mean"]} for c,r in zip(GRID,idea_grid)], "paired_deltas":deltas,"paired_p_value":permutation_pvalue(deltas)}) report["elapsed_sec"]=time.time()-t with open("bench_report.json","w") as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__ == "__main__": main()