import json, math, os, sys 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 = tuple(range(4)) EPOCHS = 18 BATCH = 128 class OrderAdaptiveIntegral(torch.optim.Optimizer): """Global-order nested integrated-gradient optimizer for this benchmark.""" def __init__(self, params, lr, max_order=2, beta=.95, rho=.98, interval=2, patience=2, ramp_steps=8, gains=(.05, .0005), state_clip=10.0, a0=1.0): super().__init__(params, dict(lr=lr)) self.max_order, self.beta, self.rho = max_order, beta, rho self.interval, self.patience, self.ramp_steps = interval, patience, ramp_steps self.gains, self.state_clip, self.a0 = tuple(gains), state_clip, a0 self.order, self.ema, self.bad = 0, None, 0 self.step_count, self.activations = 0, [] for group in self.param_groups: for p in group['params']: self.state[p]['z'] = {} self.state[p]['gamma'] = {} @torch.no_grad() def step(self): ss, count = 0.0, 0 for group in self.param_groups: for p in group['params']: if p.grad is not None: ss += float(p.grad.detach().pow(2).sum()) count += p.numel() gn = math.sqrt(ss) / max(math.sqrt(count), 1.0) self.ema = gn if self.ema is None else self.beta*self.ema + (1-self.beta)*gn self.step_count += 1 if self.step_count % self.interval == 0 and self.order < self.max_order: previous = getattr(self, '_last_decision_ema', None) if previous is not None and self.ema > self.rho * previous: self.bad += 1 else: self.bad = 0 self._last_decision_ema = self.ema if self.bad >= self.patience: self.order += 1 self.bad = 0 self.activations.append({'step': self.step_count, 'order': self.order, 'ema': float(self.ema)}) for group in self.param_groups: for p in group['params']: self.state[p]['z'][self.order] = torch.zeros_like(p) self.state[p]['gamma'][self.order] = 0.0 for group in self.param_groups: lr = group['lr'] for p in group['params']: if p.grad is None: continue st, g = self.state[p], p.grad update = self.a0 * g for r in range(1, self.order + 1): z = st['z'][r] z.add_(g if r == 1 else st['z'][r-1], alpha=lr) zn = float(z.norm()) if np.isfinite(zn) and zn > self.state_clip: z.mul_(self.state_clip / zn) st['gamma'][r] = min(1.0, st['gamma'][r] + 1.0/max(1, self.ramp_steps)) update = update + st['gamma'][r] * self.gains[r-1] * z p.add_(update, alpha=-lr) def train_one(track, model_name, cfg, seed, adaptive): torch.manual_seed(seed) np.random.seed(seed) ds = get_dataset(track, seed, n_train=400, n_test=400) model = make_model(model_name, ds['input_shape'], ds['out_dim']) device = 'cuda' if torch.cuda.is_available() else 'cpu' lossf = nn.MSELoss() try: model = model.to(device) xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device) if adaptive: opt = OrderAdaptiveIntegral(model.parameters(), lr=cfg['lr'], max_order=2, interval=2, patience=2, ramp_steps=8) else: opt = torch.optim.SGD(model.parameters(), lr=cfg['lr'], momentum=cfg['momentum']) hist, grad_hist, order_hist = [], [], [] for _ in range(EPOCHS): model.train(); perm = torch.randperm(len(xtr), device=device); total = 0. for i in range(0, len(xtr), BATCH): idx = perm[i:i+BATCH] loss = lossf(model(xtr[idx]), ytr[idx]) opt.zero_grad(); loss.backward() gn = math.sqrt(sum(float(p.grad.detach().pow(2).sum()) for p in model.parameters() if p.grad is not None)) opt.step(); total += float(loss) * len(idx) grad_hist.append(gn) order_hist.append(getattr(opt, 'order', 0)) hist.append(total / len(xtr)) model.eval() with torch.no_grad(): metric = float(lossf(model(ds['xte'].to(device)), ds['yte'].to(device))) return metric, {'model': model, 'history': hist, 'grad_hist': grad_hist, 'order_hist': order_hist, 'activations': getattr(opt, 'activations', [])} except RuntimeError: if device == 'cuda': torch.cuda.empty_cache() return train_one_cpu(track, model_name, cfg, seed, adaptive) raise def train_one_cpu(track, model_name, cfg, seed, adaptive): torch.manual_seed(seed); np.random.seed(seed) ds = get_dataset(track, seed, n_train=400, n_test=400) model = make_model(model_name, ds['input_shape'], ds['out_dim']) opt = (OrderAdaptiveIntegral(model.parameters(), lr=cfg['lr'], max_order=2, interval=2, patience=2, ramp_steps=8) if adaptive else torch.optim.SGD(model.parameters(), lr=cfg['lr'], momentum=cfg['momentum'])) lossf = nn.MSELoss(); hist=[]; grads=[]; orders=[] for _ in range(EPOCHS): perm=torch.randperm(len(ds['xtr'])); total=0. for i in range(0,len(perm),BATCH): idx=perm[i:i+BATCH]; loss=lossf(model(ds['xtr'][idx]),ds['ytr'][idx]); opt.zero_grad(); loss.backward() grads.append(math.sqrt(sum(float(p.grad.detach().pow(2).sum()) for p in model.parameters() if p.grad is not None))) opt.step(); orders.append(getattr(opt,'order',0)); total += float(loss)*len(idx) hist.append(total/len(perm)) with torch.no_grad(): metric=float(lossf(model(ds['xte']),ds['yte'])) return metric, {'model':model,'history':hist,'grad_hist':grads,'order_hist':orders,'activations':getattr(opt,'activations',[])} def main(): track, model_name = 'tabular', 'mlp_tiny' # Union of baseline and idea step sizes; baseline also sweeps its central momentum knob. grid = [{'lr': lr, 'momentum': m} for lr in (0.003, 0.01, 0.03) for m in (0.0, 0.9)] base = sweep_baseline(lambda c: lambda s: train_one(track, model_name, c, s, False)[0], grid, seeds=SWEEP_SEEDS) best_lr, best_m = base['best_cfg']['lr'], base['best_cfg']['momentum'] idea_grid = [best_lr, 0.003 if best_lr != 0.003 else 0.01, 0.03 if best_lr != 0.03 else 0.01] idea_cfgs = [{'lr': x, 'momentum': best_m} for x in idea_grid] # Evaluate all three candidate idea settings on the same full paired seeds. candidates=[] for c in idea_cfgs: vals=[] for s in SEEDS: vals.append(train_one(track, model_name, c, s, True)[0]) candidates.append({'cfg':c,'result':{'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)}}) best=min(candidates,key=lambda x:x['result']['mean']) sig=[] for s in SEEDS: _, info=train_one(track, model_name, best['cfg'], s, True) gh=np.asarray(info['grad_hist']); oh=np.asarray(info['order_hist']) active=np.where(oh>0)[0] pre=float(np.mean(gh[:max(1, len(gh)//3)])); post=float(np.mean(gh[-max(1, len(gh)//3):])) sig.append({'seed':s,'activated':bool(len(active)),'first_activation':int(active[0]) if len(active) else None,'grad_ratio':post/max(pre,1e-12)}) idea_res=best['result'] baseline=base # Signature prediction: easy tabular training should usually stay order zero; measure trained behavior. frac=float(np.mean([x['activated'] for x in sig])) confirmed=bool(frac < 0.5 and all(np.isfinite(x['grad_ratio']) for x in sig)) report=make_report(track, model_name, baseline, idea_res, {'mechanism_signature':{ 'prediction':'on easy well-conditioned tabular training adaptive order remains mostly p=0 while residual contracts', 'observed_activation_fraction':frac,'per_seed':sig,'confirmed':confirmed}, 'protocol_note':'baseline sweep uses all lr values tried by idea and sweeps momentum; idea uses three lr settings.'}) report['idea_candidates']=candidates with open('bench_report.json','w') as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__ == '__main__': main()