import sys, json, copy, itertools import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import bench SEEDS = tuple(range(8)) GRID = [ {'lr': 0.0015, 'weight_decay': 0.0}, {'lr': 0.0030, 'weight_decay': 0.0}, {'lr': 0.0060, 'weight_decay': 0.0}, {'lr': 0.0030, 'weight_decay': 1e-4}, ] EPOCHS, BATCH = 12, 64 def setup(seed): np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) ds = bench.get_dataset('sequence', seed, n_train=400, n_test=200) model = bench.make_model('transformer_tiny', tuple(ds['input_shape']), int(ds['out_dim'])) return ds, model def baseline_train(cfg, seed, keep=False): ds, model = setup(seed) net, metric, hist = bench.train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *a, **k: None) if metric is None: return (float('nan'), None, ds, hist) if keep else float('nan') return (float(metric), net, ds, hist) if keep else float(metric) def make_blocks(model): return [list(model.inp.parameters()) + [model.pos], list(model.enc.parameters()), list(model.head.parameters())] def param_vec(ps): return torch.cat([p.detach().reshape(-1) for p in ps]) def choose_order(model, loss_fn, blocks, lr, hysteresis=0.05): # Empirical block Jacobian: normalized perturbation of source block, measured # through the resulting gradient/update in each target block. params = [p for group in blocks for p in group] saved = [p.detach().clone() for p in params] grads = [] model.zero_grad(set_to_none=True); loss_fn().backward() for group in blocks: grads.append(param_vec([p.grad for p in group]).detach().clone()) rows = np.zeros((3, 3), dtype=float) eps = 1e-3 for b, group in enumerate(blocks): norm = max(float(param_vec(group).norm()), 1.0) for p in group: p.data.add_(eps * torch.randn_like(p) / norm) model.zero_grad(set_to_none=True); loss_fn().backward() for a, ag in enumerate(blocks): ng = param_vec([p.grad for p in ag]).detach() rows[a, b] = float((ng - grads[a]).norm()) / eps for p, s in zip(params, saved): p.data.copy_(s) # Each diagonal block is approximated by its local update contraction and # off-diagonal terms by normalized cross-block gradient sensitivity. scale = max(float(np.linalg.norm(rows)), 1e-8) J = rows / scale best = None; best_rho = float('inf') for order in itertools.permutations(range(3)): M = np.eye(3) for b in order: T = np.eye(3); T[b, :] -= np.minimum(0.95, lr * (J[b, :] + np.eye(3)[b, :])) M = T @ M rho = float(np.max(np.abs(np.linalg.eigvals(M)))) if rho < best_rho: best_rho, best = rho, order return list(best), best_rho, float(np.linalg.norm(rows)), rows def idea_train(cfg, seed, keep=False): ds, model = setup(seed); device = 'cuda' if torch.cuda.is_available() else 'cpu' loss_fn = lambda: nn.MSELoss()(model(ds['xtr'].to(device)), ds['ytr'].to(device)) blocks = make_blocks(model); model.to(device) blocks = [[p for p in g if p.device.type == device] for g in blocks] opts = [torch.optim.Adam(g, lr=cfg['lr'], weight_decay=cfg['weight_decay']) for g in blocks] order = [0, 1, 2]; selected_rhos=[]; hist=[] try: for ep in range(EPOCHS): # Full-batch keeps the sequential-vs-simultaneous distinction clear. if ep % 3 == 0: cand, rho, sens, mat = choose_order(model, loss_fn, blocks, cfg['lr']) selected_rhos.append(rho) if cand != order: # hysteresis: only switch if the candidate has a meaningful # predicted improvement over the current ordering. def radius(ordr): M=np.eye(3) for b in ordr: T=np.eye(3); T[b,:]-=np.minimum(.95,cfg['lr']*(mat[b,:]+np.eye(3)[b,:])/max(np.linalg.norm(mat),1e-8)) M=T@M return float(np.max(np.abs(np.linalg.eigvals(M)))) if radius(order) <= 0 or radius(cand) < .95*radius(order): order=cand for b in order: for o in opts: o.zero_grad(set_to_none=True) loss=loss_fn(); loss.backward(); opts[b].step() hist.append(float(loss_fn().detach().cpu())) with torch.no_grad(): metric=float(nn.MSELoss()(model(ds['xte'].to(device)), ds['yte'].to(device)).cpu()) except RuntimeError: if device == 'cuda': torch.cuda.empty_cache(); return idea_train_cpu(cfg, seed, keep) return (float('nan'), None, ds, hist, [], []) if keep else float('nan') return (metric, model, ds, hist, selected_rhos, order) if keep else metric def idea_train_cpu(cfg, seed, keep=False): old=torch.cuda.is_available # Re-run a CPU-only copy using the same intervention. ds, model=setup(seed); model=model.cpu(); x,y=ds['xtr'],ds['ytr']; xt,yt=ds['xte'],ds['yte'] bs=make_blocks(model); opts=[torch.optim.Adam(g,lr=cfg['lr'],weight_decay=cfg['weight_decay']) for g in bs]; order=[0,1,2]; hist=[] for _ in range(EPOCHS): for b in order: for o in opts:o.zero_grad(set_to_none=True) loss=nn.MSELoss()(model(x),y); loss.backward(); opts[b].step() hist.append(float(loss.detach())) metric=float(nn.MSELoss()(model(xt),yt)); return (metric,model,ds,hist,[],order) if keep else metric def evaluate(fn, seeds=SEEDS): return [float(fn(s)) for s in seeds] def main(): base_grid=[] for c in GRID: vals=evaluate(lambda s,c=c:baseline_train(c,s), tuple(range(4))) base_grid.append({'cfg':c,'mean':float(np.mean(vals))}) best=min(base_grid,key=lambda z:z['mean'])['cfg'] base_full=evaluate(lambda s:baseline_train(best,s)) base={'best_cfg':best,'sweep':base_grid,'full':{'mean':float(np.mean(base_full)), 'std':float(np.std(base_full)), 'per_seed':base_full, 'n':8}} idea_cfgs=[best, {'lr':0.0015,'weight_decay':best['weight_decay']}, {'lr':0.006,'weight_decay':best['weight_decay']}] idea_trials=[] for c in idea_cfgs: vals=evaluate(lambda s,c=c:idea_train(c,s),tuple(range(4))); idea_trials.append({'cfg':c,'mean':float(np.mean(vals))}) ibest=min(idea_trials,key=lambda z:z['mean'])['cfg']; iv=evaluate(lambda s:idea_train(ibest,s)) idea={'best_cfg':ibest,'sweep':idea_trials,'mean':float(np.mean(iv)),'std':float(np.std(iv)),'per_seed':iv,'n':8} rep=bench.make_report('sequence','transformer_tiny',base,idea,{'predicted_vs_observed': 'trained-model finite-difference block sensitivities and selected-order contraction', 'confirmed': False}) rep['baseline']['full']['std']=float(np.std(base_full)); rep['idea']['std']=float(np.std(iv)) with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()