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, make_report from bench.protocol import evaluate, sweep_baseline SEEDS = tuple(range(8)) EPOCHS = 8 BATCH = 128 # Union parity: every idea step size is also evaluated by Adam. LRS = [1e-3, 3e-3, 6e-3] WDS = [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 dev(): return torch.device('cuda' if torch.cuda.is_available() else 'cpu') def batches(n, seed, device): g = torch.Generator(device=device).manual_seed(seed) ix = torch.randperm(n, generator=g, device=device) for i in range(0, n, BATCH): yield ix[i:i+BATCH] def train_adam(seed, lr, wd, capture=False): seed_all(seed) ds = get_dataset('tabular', seed, 400, 400) device = dev() try: net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) x, y = ds['xtr'].to(device), ds['ytr'].to(device) opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd) history = [] for ep in range(EPOCHS): net.train() for ix in batches(len(x), seed + 1009 * ep, device): loss = nn.functional.mse_loss(net(x[ix]), y[ix]) opt.zero_grad(set_to_none=True) loss.backward(); opt.step() history.append(float(loss.detach().cpu())) net.eval() with torch.no_grad(): metric = float(nn.functional.mse_loss(net(ds['xte'].to(device)), ds['yte'].to(device)).cpu()) return metric, {'history': history, 'energy': []} except RuntimeError: if device.type == 'cuda': torch.cuda.empty_cache() return train_adam_cpu(seed, lr, wd, capture) raise def train_adam_cpu(seed, lr, wd, capture=False): seed_all(seed) ds = get_dataset('tabular', seed, 400, 400) net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']) x, y = ds['xtr'], ds['ytr'] opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd) history = [] for ep in range(EPOCHS): for ix in batches(len(x), seed + 1009 * ep, torch.device('cpu')): loss = nn.functional.mse_loss(net(x[ix]), y[ix]) opt.zero_grad(set_to_none=True); loss.backward(); opt.step() history.append(float(loss.detach())) with torch.no_grad(): metric = float(nn.functional.mse_loss(net(ds['xte']), ds['yte'])) return metric, {'history': history, 'energy': []} def flat_params(params): return torch.cat([p.detach().reshape(-1) for p in params]) def train_leapfrog(seed, h, mass=1.0, capture=False): seed_all(seed) ds = get_dataset('tabular', seed, 400, 400) device = dev() try: net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) x, y = ds['xtr'].to(device), ds['ytr'].to(device) params = list(net.parameters()) mom = [torch.zeros_like(p) for p in params] energies = []; history = [] # Two gradient evaluations per leapfrog step, matching Adam's budget by epochs. for ep in range(EPOCHS): net.train() for bi, ix in enumerate(batches(len(x), seed + 1009 * ep, device)): net.zero_grad(set_to_none=True) loss = nn.functional.mse_loss(net(x[ix]), y[ix]); loss.backward() g1 = [p.grad.detach().clone() for p in params] with torch.no_grad(): for p, q, g in zip(params, mom, g1): q.sub_(0.5 * h * g); p.add_(h * q / mass) net.zero_grad(set_to_none=True) loss2 = nn.functional.mse_loss(net(x[ix]), y[ix]); loss2.backward() g2 = [p.grad.detach().clone() for p in params] with torch.no_grad(): for q, g in zip(mom, g2): q.sub_(0.5 * h * g) if capture and (bi == 0 or bi == len(list(batches(len(x), seed + 1009 * ep, device))) - 1): with torch.no_grad(): kinetic = 0.5 * sum(float((q*q).sum().cpu()) / mass for q in mom) potential = float(nn.functional.mse_loss(net(x), y).detach().cpu()) energies.append(potential + kinetic) history.append(float(loss2.detach().cpu())) net.eval() with torch.no_grad(): metric = float(nn.functional.mse_loss(net(ds['xte'].to(device)), ds['yte'].to(device)).cpu()) return metric, {'history': history, 'energy': energies} except RuntimeError: if device.type == 'cuda': torch.cuda.empty_cache() old = globals()['dev']; globals()['dev'] = lambda: torch.device('cpu') try: return train_leapfrog(seed, h, mass, capture) finally: globals()['dev'] = old raise def main(): # Baseline sweep includes all step sizes and both central Adam weight-decay values. grid = [{'lr': lr, 'weight_decay': wd} for lr in LRS for wd in WDS] base = sweep_baseline(lambda c: lambda s: train_adam(s, c['lr'], c['weight_decay'])[0], grid) # Idea sweep: best baseline h plus two nearby settings; all are in baseline union. idea_grid = [{'h': h, 'mass': 1.0} for h in LRS] idea_cfg_results = [] for cfg in idea_grid: r = evaluate(lambda s, c=cfg: train_leapfrog(s, c['h'], c['mass'])[0]) idea_cfg_results.append({'cfg': cfg, 'result': r}) best = min(idea_cfg_results, key=lambda z: z['result']['mean']) idea = best['result']; cfg = best['cfg'] # Re-test the mechanism on trained systems: energy oscillation versus h. hs = [0.001, 0.003, 0.006] ranges = [] for h in hs: vals = [] for s in range(4): z = train_leapfrog(s, h, 1.0, capture=True)[1]['energy'] vals.append(max(z) - min(z) if z else float('nan')) ranges.append(float(np.nanmean(vals))) slope = float(np.polyfit(np.log(hs), np.log(np.maximum(ranges, 1e-30)), 1)[0]) sig = {'prediction': 'leapfrog trained-model energy oscillation scales approximately as h^2', 'step_sizes': hs, 'observed_energy_ranges': ranges, 'loglog_slope': slope, 'confirmed': bool(1.5 < slope < 2.5)} report = make_report('tabular', 'mlp_tiny', base, idea, { 'track_choice': 'optimizer intervention structurally matches the tabular optimizer track', 'baseline_grid': grid, 'idea_grid': idea_grid, 'idea_best_cfg': cfg, 'mechanism_signature': sig}) report['idea_sweep'] = idea_cfg_results Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()