import sys, json, math, random import numpy as np sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import torch from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) LR_GRID = [0.003, 0.006, 0.009] EPOCHS = 18 BATCH = 128 CLIP = 1.0 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 run(seed, lr, scheduled): seed_all(seed) requested = 'cuda' if torch.cuda.is_available() else 'cpu' try: return _run_device(seed, lr, scheduled, requested) except Exception: if requested == 'cuda': try: torch.cuda.empty_cache() except Exception: pass return _run_device(seed, lr, scheduled, 'cpu') raise def _run_device(seed, lr, scheduled, device_name): device = torch.device(device_name) d = get_dataset('tabular', seed, n_train=400, n_test=200) xtr, ytr = d['xtr'].to(device), d['ytr'].to(device) xte, yte = d['xte'].to(device), d['yte'].to(device) model = make_model('mlp_tiny', d['input_shape'], d['out_dim']).to(device) opt = torch.optim.SGD(model.parameters(), lr=lr) fhat = None; qhist = []; gaphist = []; train_hist = [] local_streak = 0; switched = False; switch_step = None; clip_steps = 0 step = 0 n = len(xtr) for epoch in range(EPOCHS): # deterministic per-seed permutation, equivalent data budget in both arms gen = torch.Generator(device='cpu').manual_seed(seed * 1000 + epoch) perm = torch.randperm(n, generator=gen, device='cpu') for start in range(0, n, BATCH): idx = perm[start:start+BATCH].to(device) xb, yb = xtr[idx], ytr[idx] opt.zero_grad(set_to_none=True) pred = model(xb) loss = torch.nn.functional.mse_loss(pred, yb) loss.backward() g2 = 0.0 for p in model.parameters(): if p.grad is not None: g2 += float((p.grad.detach() ** 2).sum().cpu()) g = math.sqrt(max(g2, 0.0)); lv = float(loss.detach().cpu()) if fhat is None: fhat = lv else: fhat = min(fhat, 0.98 * fhat + 0.02 * lv) gap = max(lv - fhat, 1e-8) q = g / math.sqrt(gap) qhist.append(q); gaphist.append(gap); train_hist.append(lv) if scheduled and len(qhist) >= 20: qwin = np.asarray(qhist[-100:]); gapwin = np.asarray(gaphist[-100:]) floor = max(float(np.percentile(qwin, 10)) / 2.0, 1e-5) candidate = gap <= float(np.percentile(gapwin, 35)) and q >= floor local_streak = local_streak + 1 if candidate else 0 if local_streak >= 5 and not switched: for group in opt.param_groups: group['lr'] *= 1.5 switched = True; switch_step = step pre = torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP) clip_steps += int(float(pre) > CLIP) opt.step(); step += 1 with torch.no_grad(): test_loss = float(torch.nn.functional.mse_loss(model(xte), yte).cpu()) # Re-test the claimed local prediction on the trained model's observed trajectory. qarr, garr = np.asarray(qhist), np.asarray(gaphist) tail = min(40, len(qarr)) observed_slope = float(np.polyfit(np.arange(tail), np.log(np.maximum(garr[-tail:], 1e-12)), 1)[0]) if tail >= 3 else float('nan') predicted_slope = float(-np.mean(qarr[-tail:] ** 2)) if tail else float('nan') return {'metric': test_loss, 'switched': switched, 'switch_step': switch_step, 'lr_final': float(opt.param_groups[0]['lr']), 'clip_steps': clip_steps, 'observed_log_gap_slope': observed_slope, 'predicted_minus_q2': predicted_slope, 'q_tail_mean': float(np.mean(qarr[-tail:])) if tail else float('nan'), 'device': device_name} def main(): base = sweep_baseline(lambda cfg: (lambda seed: run(seed, cfg['lr'], False)['metric']), [{'lr': x} for x in LR_GRID], seeds=SWEEP_SEEDS) best_lr = float(base['best_cfg']['lr']) idea_cfgs = [{'lr': x} for x in LR_GRID] idea_sweep = [] for cfg in idea_cfgs: r = evaluate(lambda seed, x=cfg['lr']: run(seed, x, True)['metric'], seeds=SWEEP_SEEDS) idea_sweep.append({'cfg': cfg, 'mean': r['mean']}) best_idea_cfg = min(idea_sweep, key=lambda z: z['mean'])['cfg'] # Full 8-seed results at the best idea configuration, with full per-run diagnostics. idea_runs = [run(s, best_idea_cfg['lr'], True) for s in SEEDS] idea_res = {'mean': float(np.mean([r['metric'] for r in idea_runs])), 'std': float(np.std([r['metric'] for r in idea_runs])), 'per_seed': [r['metric'] for r in idea_runs], 'n': 8, 'cfg': best_idea_cfg, 'runs': idea_runs, 'sweep': idea_sweep} rep = make_report('tabular', 'mlp_tiny', base, idea_res) # Signature uses only trained-model observed numbers, not the toy identity. obs = [r['observed_log_gap_slope'] for r in idea_runs] pred = [r['predicted_minus_q2'] for r in idea_runs] rel = [abs(a-b)/max(abs(b),1e-12) for a,b in zip(obs,pred) if np.isfinite(a) and np.isfinite(b)] rep['mechanism_signature'] = { 'quantity': 'tail log(training loss gap) slope versus -mean(q^2), measured during trained NN runs', 'observed_mean': float(np.mean(obs)), 'predicted_mean': float(np.mean(pred)), 'relative_error_mean': float(np.mean(rel)) if rel else None, 'tolerance': 0.20, 'confirmed': bool(rel and np.mean(rel) <= 0.20), 'switch_fraction': float(np.mean([r['switched'] for r in idea_runs]))} rep['audit'] = {'baseline_lr_grid': LR_GRID, 'idea_lr_grid': LR_GRID, 'baseline_selected_lr': best_lr, 'idea_selected_lr': best_idea_cfg['lr'], 'epochs': EPOCHS, 'batch': BATCH, 'clip_norm': CLIP, 'domain_rationale': 'tabular is the prescribed structural track for optimizer and learning-rate schedule ideas'} print(json.dumps(rep, indent=2)) if __name__ == '__main__': main()