import sys, os, json, math, 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, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) EPOCHS = 15 NTRAIN, NTEST = 1200, 400 BATCH = 128 ALPHA = 1.5 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 stable_noise(shape, alpha=ALPHA, device='cpu'): """Unit-scale symmetric alpha-stable samples (Chambers--Mallows--Stuck).""" u = torch.rand(shape, device=device) * math.pi - math.pi / 2 w = torch.empty(shape, device=device).exponential_(1.0) # The symmetric CMS formula is well behaved for alpha in (1,2). return (torch.sin(alpha * u) / torch.cos(u).pow(1.0 / alpha) * (torch.cos((1.0 - alpha) * u) / w).pow((1.0-alpha)/alpha)) def train_one(seed, cfg, levy, collect=False): seed_all(seed) ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST) device = 'cuda' if torch.cuda.is_available() else 'cpu' try: net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) x, y = ds['xtr'].to(device), ds['ytr'].to(device) xt, yt = ds['xte'].to(device), ds['yte'].to(device) opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg.get('wd', 0.0)) lossf = nn.MSELoss() anchors = [p.detach().clone() for p in net.parameters()] drift_ratios, jump_samples = [], [] beta = cfg.get('beta', 0.95) lam, noise = cfg.get('lambda', 0.0), cfg.get('noise', 0.0) for _ in range(EPOCHS): net.train() perm = torch.randperm(len(x), device=device) for start in range(0, len(x), BATCH): idx = perm[start:start+BATCH] opt.zero_grad(set_to_none=True) loss = lossf(net(x[idx]), y[idx]) loss.backward() opt.step() with torch.no_grad(): for j, p in enumerate(net.parameters()): # EMA anchor is maintained per parameter tensor. old = p.detach().clone() anchors[j].mul_(beta).add_(old, alpha=1-beta) if levy: z = stable_noise(p.shape, device=device) jump = noise * z before = p.detach().clone() p.add_( -cfg['lr'] * lam * (p - anchors[j]) + jump ) if collect: d = (before - anchors[j]).flatten() det = (before - cfg['lr']*lam*(before-anchors[j]) - anchors[j]).flatten() good = d.abs() > 1e-12 if bool(good.any()): drift_ratios.append(float(det[good].norm() / d[good].norm())) jump_samples.append(jump.detach().flatten().cpu().numpy()) net.eval() with torch.no_grad(): metric = float(((net(xt) - yt) ** 2).mean().cpu()) stats = {'drift_ratios': drift_ratios, 'jumps': np.concatenate(jump_samples) if jump_samples else np.array([])} return metric, stats, net except RuntimeError: # Explicit CPU retry, matching the harness's GPU-fallback requirement. if device == 'cuda': torch.cuda.empty_cache() old = torch.cuda.is_available torch.cuda.is_available = lambda: False try: return train_one(seed, cfg, levy, collect) finally: torch.cuda.is_available = old raise def train_fn(cfg, levy, collect=False): def f(seed): v, st, _ = train_one(int(seed), cfg, levy, collect=collect) if collect: f.stats.append(st) return v f.stats = [] return f def mechanism_signature(cfg): # Measured from the actual trained networks, not the toy system. f = train_fn(cfg, True, collect=True) vals = evaluate(f, seeds=SEEDS) ratios = np.concatenate([s['drift_ratios'] for s in f.stats if s['drift_ratios']]) jumps = np.concatenate([s['jumps'] for s in f.stats if len(s['jumps'])]) # Robust empirical characteristic-function power fit on NN optimizer jumps. us = np.array([0.25, 0.4, 0.6, 0.8]) phi = np.array([np.mean(np.cos(u*jumps)) for u in us]) cf_y = -np.log(np.maximum(phi, 1e-8)) cf_power = float(np.polyfit(np.log(us), np.log(cf_y), 1)[0]) q = cfg['lr'] * cfg['lambda'] predicted = abs(1-q) observed = float(np.median(ratios)) if len(ratios) else float('nan') confirmed = bool(np.isfinite(observed) and abs(observed-predicted) <= 0.05 and abs(cf_power-ALPHA) <= 0.20) return {'alpha': ALPHA, 'lambda': cfg['lambda'], 'lr': cfg['lr'], 'q': q, 'predicted_drift_multiplier': predicted, 'observed_median_drift_multiplier': observed, 'observed_nn_jump_cf_power': cf_power, 'jump_mad': float(np.median(np.abs(jumps-np.median(jumps)))), 'confirmed': confirmed, 'n_observations': int(len(jumps)), 'training_metric_for_signature': vals} def main(): # The union of baseline and idea learning rates is identical (parity). lrs = [1e-3, 3e-3, 1e-2] base_grid = [{'lr': lr, 'wd': 0.0} for lr in lrs] base = sweep_baseline(lambda c: train_fn(c, False), base_grid, seeds=SWEEP_SEEDS) best_lr = base['best_cfg']['lr'] idea_grid = [ {'lr': best_lr, 'wd': 0.0, 'lambda': 1.0, 'beta': 0.95, 'noise': 0.001}, {'lr': 1e-3 if best_lr != 1e-3 else 3e-3, 'wd': 0.0, 'lambda': 1.0, 'beta': 0.95, 'noise': 0.001}, {'lr': 1e-2 if best_lr != 1e-2 else 3e-3, 'wd': 0.0, 'lambda': 1.0, 'beta': 0.95, 'noise': 0.001}, ] idea_trials = [] for cfg in idea_grid: r = evaluate(train_fn(cfg, True), seeds=SWEEP_SEEDS) idea_trials.append({'cfg': cfg, 'mean': r['mean']}) best_idea_cfg = min(idea_trials, key=lambda z: z['mean'])['cfg'] idea_full = evaluate(train_fn(best_idea_cfg, True), seeds=SEEDS) sig = mechanism_signature(best_idea_cfg) report = make_report('tabular', 'mlp_tiny', {'best_cfg': base['best_cfg'], 'sweep': base['sweep'], 'full': base['full']}, idea_full, {'signature_type': 'trained_nn_optimizer_dynamics', **sig}) report['idea_sweep'] = idea_trials report['protocol'] = {'paired_seeds': list(SEEDS), 'epochs': EPOCHS, 'n_train': NTRAIN, 'n_test': NTEST, 'structural_match': 'optimizer intervention on tabular regression'} 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()