import sys, json, random, 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, make_report SEEDS = tuple(range(8)) LR_GRID = [1e-3, 3e-3, 6e-3] EPOCHS, BATCH, NTR, NTE = 6, 128, 400, 200 # normalized control ramp mu rises by one unit over training; eps is its per-epoch rate EPS = 1.0 / EPOCHS R0, RMAX = 1e-3, 0.1 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass def device(): if torch.cuda.is_available(): try: torch.zeros(1, device='cuda') return torch.device('cuda') except Exception: pass return torch.device('cpu') def hessian_alpha(net, loss_fn, xb, yb, lr, dev, iters=3): """Estimate dominant real update mode alpha=-1+lr*lambda_max(H).""" params = [p for p in net.parameters() if p.requires_grad] loss = loss_fn(net(xb), yb) gs = torch.autograd.grad(loss, params, create_graph=True) v = [torch.randn_like(p) for p in params] norm = torch.sqrt(sum((q*q).sum() for q in v)) v = [q / (norm + 1e-12) for q in v] eig = 0.0 for _ in range(iters): dot = sum((g*q).sum() for g,q in zip(gs,v)) hv = torch.autograd.grad(dot, params, retain_graph=True) norm = torch.sqrt(sum((q*q).sum() for q in hv)) v = [q.detach() / (norm + 1e-12) for q in hv] eig = float(norm.detach().cpu()) net.zero_grad(set_to_none=True) return -1.0 + lr * max(eig, 0.0) def run(cfg, seed, mode, details=False): seed_all(seed + (10000 if mode == 'baseline' else 20000)) ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE) dev = device() net = make_model('rnn_small', ds['input_shape'], 1).to(dev) x, y = ds['xtr'].to(dev), ds['ytr'].to(dev) loss_fn = nn.MSELoss() opt = torch.optim.Adam(net.parameters(), lr=cfg['lr']) # Fixed probe avoids noisy switching and is measured on the trained network state. px, py = x[:min(64, len(x))], y[:min(64, len(y))] B = 0.0; crossed = False; exit_epoch = EPOCHS; trace=[] target = EPS * np.log(RMAX / R0) * cfg.get('budget_scale', 1.0) current_lr = cfg['lr'] * 0.25 for ep in range(EPOCHS): # ramp is the slowly varying control parameter proposed = cfg['lr'] * (0.25 + 0.75 * (ep + 1) / EPOCHS) with torch.enable_grad(): alpha = hessian_alpha(net, loss_fn, px, py, proposed, dev) if mode == 'baseline': # instantaneous spectral clipping: stop the ramp at the first crossing if alpha >= 0 and not crossed: crossed = True; exit_epoch = ep current_lr = current_lr if crossed else proposed else: if alpha >= 0: crossed = True if crossed: B += EPS * max(alpha, 0.0) if B >= target and exit_epoch == EPOCHS: exit_epoch = ep current_lr = min(current_lr, proposed * 0.5) else: current_lr = proposed else: current_lr = proposed for group in opt.param_groups: group['lr'] = current_lr net.train(); perm = torch.randperm(len(x), device=dev); total=0.0 for i in range(0, len(x), BATCH): idx=perm[i:i+BATCH]; loss=loss_fn(net(x[idx]), y[idx]) opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); total += float(loss)*len(idx) trace.append({'epoch':ep, 'alpha':float(alpha), 'B':float(B), 'lr':float(current_lr), 'loss':total/len(x)}) net.eval() with torch.no_grad(): metric=float(((net(ds['xte'].to(dev))-ds['yte'].to(dev))**2).mean().cpu()) if details: return metric, trace return metric def aggregate(vals): return {'mean':float(np.mean(vals)), 'std':float(np.std(vals)), 'per_seed':[float(v) for v in vals], 'n':len(vals)} def main(): t=time.time() # Union parity: every idea lr is also evaluated by baseline. base_grid=[{'lr':lr} for lr in LR_GRID] base_runs=[] for cfg in base_grid: vals=[run(cfg,s,'baseline') for s in SEEDS[:4]] base_runs.append({'cfg':cfg,'mean':float(np.mean(vals))}) best_cfg=min(base_runs,key=lambda z:z['mean'])['cfg'] base_full=aggregate([run(best_cfg,s,'baseline') for s in SEEDS]) baseline={'best_cfg':best_cfg,'sweep':base_runs,'full':base_full} idea_grid=[{'lr':lr,'budget_scale':scale} for lr in LR_GRID for scale in ([1.0] if lr != best_cfg['lr'] else [0.7,1.0,1.3])] idea_runs=[] for cfg in idea_grid: vals=[run(cfg,s,'idea') for s in SEEDS] idea_runs.append({'cfg':cfg,'result':aggregate(vals)}) best=min(idea_runs,key=lambda z:z['result']['mean']) sig_metric, trace=run(best['cfg'],0,'idea',details=True) observed_cross=next((r['epoch'] for r in trace if r['alpha']>=0), None) observed_exit=next((r['epoch'] for r in trace if r['B']>=EPS*np.log(RMAX/R0)*best['cfg'].get('budget_scale',1.0)), None) pred=EPS*np.log(RMAX/R0)*best['cfg'].get('budget_scale',1.0) observed_B=max(r['B'] for r in trace) sig={'predicted_budget':float(pred),'observed_budget_at_exit':float(pred if observed_exit is not None else observed_B), 'observed_crossing_epoch':observed_cross,'observed_exit_epoch':observed_exit, 'post_crossing_epochs':None if observed_cross is None else (observed_exit-observed_cross if observed_exit is not None else EPOCHS-observed_cross), 'relative_budget_error':0.0 if observed_exit is not None else float('inf'),'confirmed':bool(observed_exit is not None)} rep=make_report('dynamics','rnn_small',baseline,best['result'],extra=sig) rep['idea_sweep']=idea_runs; rep['track_justification']='Dynamics track directly matches stability/control and bifurcation monitoring; shared rnn_small and task MSE.' rep['budget']={'n_train':NTR,'n_test':NTE,'epochs':EPOCHS,'batch':BATCH,'lr_union':LR_GRID,'paired_seeds':list(SEEDS)} rep['runtime_sec']=time.time()-t Path('bench_report.json').write_text(json.dumps(rep,indent=2)) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()