import sys, json, 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, train_model, evaluate, sweep_baseline, make_report EPOCHS = 18 BATCH = 128 SEEDS = tuple(range(8)) LRS = [1e-3, 3e-3, 1e-2] WEIGHT_DECAYS = [0.0, 1e-4] 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 baseline_run(cfg, seed, keep=False): seed_all(seed) ds = get_dataset('tabular', seed, n_train=400, n_test=400) net = make_model('mlp_tiny', tuple(ds['input_shape']), ds['out_dim']) net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None) out = {'metric': float(metric), 'history': hist} return out if keep else float(metric) def params(model): return [p for p in model.parameters() if p.requires_grad] def flatten(xs): return torch.cat([x.detach().reshape(-1) for x in xs]) def assign(model, vec): pos = 0 with torch.no_grad(): for p in params(model): n = p.numel(); p.copy_(vec[pos:pos+n].view_as(p)); pos += n def trust_run(cfg, seed, keep=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: net = make_model('mlp_tiny', tuple(ds['input_shape']), ds['out_dim']).to(device) x = ds['xtr'].to(device); y = ds['ytr'].to(device) lossf = nn.MSELoss() delta = float(cfg['delta0']); delta_max = 10.0 hist, rhos, rejects, radii, pred_obs = [], [], 0, [], [] for _ in range(EPOCHS): net.train(); ps = params(net) net.zero_grad(set_to_none=True) old = lossf(net(x), y); old.backward() g = flatten([p.grad for p in ps]) # Diagonal empirical-Fisher curvature plus damping; this is the local B model. b = flatten([p.grad * p.grad for p in ps]).clamp_min(1e-6) + cfg['damping'] gn = float(g.norm()) if gn < 1e-12: break # Exact trust-region solution for positive diagonal B via bisection on lambda. def step_for(lam): return -g / (b + lam) s = step_for(0.0) if float(s.norm()) > delta: lo, hi = 0.0, 1.0 while float(step_for(hi).norm()) > delta: hi *= 2.0 for _ in range(25): mid = (lo + hi) / 2 if float(step_for(mid).norm()) > delta: lo = mid else: hi = mid s = step_for(hi) cauchy_len = min(delta, gn / float(b.max())) sc = -cauchy_len * g / (gn + 1e-12) pred = -(torch.dot(g, s) + 0.5 * torch.dot(b * s, s)) cpred = -(torch.dot(g, sc) + 0.5 * torch.dot(b * sc, sc)) if float(pred) < 0.1 * float(cpred): s, pred = sc, cpred old_vec = flatten([p for p in ps]) assign(net, old_vec + s) with torch.no_grad(): new = lossf(net(x), y) ared = old.detach() - new rho = float(ared / (pred + 1e-12)) boundary = float(s.norm()) >= 0.99 * delta accepted = rho >= 0.1 and float(pred) > 0 if not accepted: assign(net, old_vec); delta *= 0.25; rejects += 1; value = float(old) else: value = float(new) if rho > 0.75 and boundary: delta = min(2.0 * delta, delta_max) hist.append(value); rhos.append(rho); radii.append(delta) pred_obs.append({'predicted': float(pred), 'observed': float(ared), 'rho': rho}) net.eval() with torch.no_grad(): metric = float(((net(ds['xte'].to(device)) - ds['yte'].to(device)) ** 2).mean()) out = {'metric': metric, 'history': hist, 'reject_fraction': rejects / max(1, EPOCHS), 'median_rho': float(np.median(rhos)) if rhos else float('nan'), 'final_radius': delta, 'radii': radii, 'pred_obs': pred_obs} return out if keep else metric except RuntimeError: # Robust CPU fallback for shared/unsupported CUDA environments. torch.cuda.empty_cache() old = torch.cuda.is_available torch.cuda.is_available = lambda: False try: return trust_run(cfg, seed, keep) finally: torch.cuda.is_available = old def main(): base_grid = [{'lr': lr, 'weight_decay': wd} for lr in LRS for wd in WEIGHT_DECAYS] base = sweep_baseline(lambda cfg: lambda seed: baseline_run(cfg, seed), base_grid) best = base['best_cfg'] idea_grid = [{'lr': lr, 'weight_decay': best['weight_decay'], 'delta0': d, 'damping': 0.01} for lr in LRS for d in [0.05, 0.2, 0.8]] # Equal-budget idea sweep on seeds 0..3, then full paired run for its best config. tried = [] for cfg in idea_grid: vals = [trust_run(cfg, s) for s in (0, 1, 2, 3)] tried.append({'cfg': cfg, 'mean': float(np.mean(vals)), 'per_seed': vals}) ibest = min(tried, key=lambda z: z['mean'])['cfg'] idea = evaluate(lambda seed: trust_run(ibest, seed), seeds=SEEDS) extra_vals = [trust_run(ibest, s, keep=True) for s in SEEDS] allro = [q for r in extra_vals for q in r['pred_obs']] ratios = np.array([q['observed'] / q['predicted'] for q in allro if q['predicted'] > 1e-10 and np.isfinite(q['observed'])]) signature = {'predicted_decrease_mean': float(np.mean([q['predicted'] for q in allro])), 'observed_decrease_mean': float(np.mean([q['observed'] for q in allro])), 'rho_median': float(np.median(ratios)) if len(ratios) else float('nan'), 'rho_iqr': [float(np.quantile(ratios, .25)), float(np.quantile(ratios, .75))] if len(ratios) else [], 'reject_fraction_mean': float(np.mean([r['reject_fraction'] for r in extra_vals])), 'confirmed': bool(len(ratios) > 0 and 0.5 <= float(np.median(ratios)) <= 1.5)} report = make_report('tabular', 'mlp_tiny', base, idea, signature) report['idea']['sweep'] = tried report['protocol_notes'] = 'Tabular is the built-in optimizer track; both systems use identical mlp_tiny, data, epochs, and paired seeds. Baseline is Adam via train_model; idea changes only the training optimizer.' Path('bench_report.json').write_text(json.dumps(report, indent=2, allow_nan=False)) print(json.dumps(report, indent=2, allow_nan=False)) if __name__ == '__main__': main()