import json import sys import time 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, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) # Search-space parity: every idea learning rate is also a baseline candidate. GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}] EPOCHS = 12 BATCH = 128 TAU = 1e-6 RHO = 1e-3 def seed_all(seed): np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def core_sanity(): s = torch.tensor([1.0, 1e-9]) r = torch.tensor([0.1, 0.1]) full = -r / s safe = -torch.where(s >= 1e-6 * s[0], r / s, torch.zeros_like(s)) return { 'singular_values': s.tolist(), 'full_step_norm': float(torch.linalg.norm(full)), 'truncated_step_norm': float(torch.linalg.norm(safe)), 'amplification_ratio': float(torch.linalg.norm(full) / torch.linalg.norm(safe)), 'kept_rank': int((s >= 1e-6 * s[0]).sum()), } def baseline_fn(cfg): def run(seed): seed_all(seed) ds = get_dataset('tabular', seed, n_train=400, n_test=400) model = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']) _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None) return float(metric) return run class VPMLP(nn.Module): """Same mlp_tiny architecture, exposing its penultimate features.""" def __init__(self, d=10, width=64): super().__init__() self.l1 = nn.Linear(d, width) self.l2 = nn.Linear(width, width) self.out = nn.Linear(width, 1) def features(self, x): return torch.relu(self.l2(torch.relu(self.l1(x)))) def forward(self, x): return self.out(self.features(x)) def project(model, x, y, tau=TAU): phi = model.features(x) # Exact least-squares output projection, with rank-safe SVD cutoff. u, s, vh = torch.linalg.svd(phi, full_matrices=False) keep = s >= (tau * s[0] if s.numel() else 0.0) inv = torch.where(keep, 1.0 / s, torch.zeros_like(s)) w = vh.transpose(0, 1) @ (inv[:, None] * (u.transpose(0, 1) @ y)) r = phi @ w - y return w, r, s def vectorized_model(z, x, d=10, width=64): n1 = width * d n2 = width * width w1 = z[:n1].reshape(width, d) b1 = z[n1:n1 + width] off = n1 + width w2 = z[off:off + n2].reshape(width, width) b2 = z[off + n2:off + n2 + width] return torch.relu(torch.relu(x @ w1.transpose(0, 1) + b1) @ w2.transpose(0, 1) + b2) def vp_run(cfg, seed, collect_signature=False): seed_all(seed) ds = get_dataset('tabular', seed, n_train=400, n_test=400) device = 'cuda' if torch.cuda.is_available() else 'cpu' try: torch.tensor(0.0, device=device).item() except Exception: device = 'cpu' x, y = ds['xtr'].to(device), ds['ytr'].to(device) xt, yt = ds['xte'].to(device), ds['yte'].to(device) model = VPMLP().to(device) # Hidden parameters are represented as one differentiable vector. z = torch.cat([p.detach().reshape(-1) for p in (model.l1.weight, model.l1.bias, model.l2.weight, model.l2.bias)]) z = z.requires_grad_() alphas = [] ranks = [] conds = [] t0 = time.perf_counter() for _ in range(EPOCHS): with torch.no_grad(): phi = vectorized_model(z, x) u, s, vh = torch.linalg.svd(phi, full_matrices=False) inv = torch.where(s >= TAU * s[0], 1.0 / s, torch.zeros_like(s)) w = vh.T @ (inv[:, None] * (u.T @ y)) r = phi @ w - y def fixed_res(zz): return (vectorized_model(zz, x) @ w - y).reshape(-1) try: jac = torch.autograd.functional.jacobian(fixed_res, z, create_graph=False) uj, sj, vhj = torch.linalg.svd(jac, full_matrices=False) ranks.append(int((sj >= RHO * sj[0]).sum())) conds.append(float(sj[0] / max(float(sj[-1]), 1e-30))) rhs = uj.T @ r.reshape(-1) invj = torch.where(sj >= RHO * sj[0], 1.0 / sj, torch.zeros_like(sj)) delta = -(vhj.T @ (invj * rhs)) except RuntimeError: # A safe fallback retains exact variable projection if dense GN fails. delta = torch.zeros_like(z) ranks.append(0); conds.append(float('inf')) base = float(0.5 * (r * r).sum()) accepted = 0.0 # cfg['lr'] is the shared step-size knob for the GN intervention. for k in range(9): a = min(1.0, cfg['lr'] / 3e-3) * (0.5 ** k) with torch.no_grad(): ztry = z + a * delta phit = vectorized_model(ztry, x) ut, st, vht = torch.linalg.svd(phit, full_matrices=False) it = torch.where(st >= TAU * st[0], 1.0 / st, torch.zeros_like(st)) wt = vht.T @ (it[:, None] * (ut.T @ y)) loss = float(0.5 * ((phit @ wt - y) ** 2).sum()) if loss < base: z = ztry.detach().requires_grad_() accepted = a break alphas.append(accepted) with torch.no_grad(): phi = vectorized_model(z, x) u, s, vh = torch.linalg.svd(phi, full_matrices=False) inv = torch.where(s >= TAU * s[0], 1.0 / s, torch.zeros_like(s)) w = vh.T @ (inv[:, None] * (u.T @ y)) test = vectorized_model(z, xt) @ w metric = float(((test - yt) ** 2).mean()) result = {'metric': metric, 'seconds': time.perf_counter() - t0, 'accepted_steps': int(sum(a > 0 for a in alphas)), 'median_alpha': float(np.median(alphas)), 'mean_rank': float(np.mean(ranks)), 'max_jacobian_condition': float(max(conds))} if collect_signature: # Trained-model signature: compare actual full and truncated GN steps. with torch.no_grad(): phi0 = vectorized_model(z, x) u0, s0, vh0 = torch.linalg.svd(phi0, full_matrices=False) i0 = torch.where(s0 >= TAU * s0[0], 1.0 / s0, torch.zeros_like(s0)) w0 = vh0.T @ (i0[:, None] * (u0.T @ y)) r0 = phi0 @ w0 - y def res_sig(zz): return (vectorized_model(zz, x) @ w0 - y).reshape(-1) jac = torch.autograd.functional.jacobian(res_sig, z) us, ss, vhs = torch.linalg.svd(jac, full_matrices=False) rhs = us.T @ r0.reshape(-1) full = -(vhs.T @ (rhs / torch.clamp(ss, min=1e-30))) trunc = -(vhs.T @ torch.where(ss >= RHO * ss[0], rhs / ss, torch.zeros_like(ss))) result['signature'] = {'predicted_full_to_truncated_amplification': core_sanity()['amplification_ratio'], 'observed_full_step_norm': float(torch.linalg.norm(full)), 'observed_truncated_step_norm': float(torch.linalg.norm(trunc)), 'observed_amplification': float(torch.linalg.norm(full) / max(float(torch.linalg.norm(trunc)), 1e-30)), 'confirmed': bool(float(torch.linalg.norm(full)) > 10.0 * max(float(torch.linalg.norm(trunc)), 1e-30))} return result def idea_eval(cfg, seeds=SEEDS): vals = [] details = [] for seed in seeds: z = vp_run(cfg, seed, collect_signature=(seed == 0)) vals.append(z['metric']); details.append(z) out = {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'n': len(vals), 'details': details} return out def main(): # Baseline sweep and idea sweep use exactly the same three learning-rate configs. base = sweep_baseline(baseline_fn, GRID, seeds=SWEEP_SEEDS) idea_sweep = [] for cfg in GRID: r = idea_eval(cfg, seeds=SWEEP_SEEDS) idea_sweep.append({'cfg': cfg, 'mean': r['mean']}) best_cfg = min(idea_sweep, key=lambda q: q['mean'])['cfg'] idea = idea_eval(best_cfg, seeds=SEEDS) sig = idea['details'][0].get('signature', {}) rep = make_report('tabular', 'mlp_tiny', base, idea, {'mechanism_signature': sig, 'method_notes': {'track_justification': 'Optimizer idea matched to tabular regression.', 'idea_sweep': idea_sweep, 'idea_best_cfg': best_cfg, 'sanity': core_sanity(), 'epochs': EPOCHS, 'batch': BATCH, 'tau': TAU, 'rho': RHO}}) Path('bench_report.json').write_text(json.dumps(rep, indent=2)) print(json.dumps(rep, indent=2)) if __name__ == '__main__': main()