1import sys, json, copy, itertools
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6import bench
  7
  8SEEDS = tuple(range(8))
  9GRID = [
 10    {'lr': 0.0015, 'weight_decay': 0.0},
 11    {'lr': 0.0030, 'weight_decay': 0.0},
 12    {'lr': 0.0060, 'weight_decay': 0.0},
 13    {'lr': 0.0030, 'weight_decay': 1e-4},
 14]
 15EPOCHS, BATCH = 12, 64
 16
 17
 18def setup(seed):
 19    np.random.seed(seed); torch.manual_seed(seed)
 20    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 21    ds = bench.get_dataset('sequence', seed, n_train=400, n_test=200)
 22    model = bench.make_model('transformer_tiny', tuple(ds['input_shape']), int(ds['out_dim']))
 23    return ds, model
 24
 25
 26def baseline_train(cfg, seed, keep=False):
 27    ds, model = setup(seed)
 28    net, metric, hist = bench.train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'],
 29        batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *a, **k: None)
 30    if metric is None: return (float('nan'), None, ds, hist) if keep else float('nan')
 31    return (float(metric), net, ds, hist) if keep else float(metric)
 32
 33
 34def make_blocks(model):
 35    return [list(model.inp.parameters()) + [model.pos], list(model.enc.parameters()), list(model.head.parameters())]
 36
 37
 38def param_vec(ps):
 39    return torch.cat([p.detach().reshape(-1) for p in ps])
 40
 41
 42def choose_order(model, loss_fn, blocks, lr, hysteresis=0.05):
 43    # Empirical block Jacobian: normalized perturbation of source block, measured
 44    # through the resulting gradient/update in each target block.
 45    params = [p for group in blocks for p in group]
 46    saved = [p.detach().clone() for p in params]
 47    grads = []
 48    model.zero_grad(set_to_none=True); loss_fn().backward()
 49    for group in blocks: grads.append(param_vec([p.grad for p in group]).detach().clone())
 50    rows = np.zeros((3, 3), dtype=float)
 51    eps = 1e-3
 52    for b, group in enumerate(blocks):
 53        norm = max(float(param_vec(group).norm()), 1.0)
 54        for p in group: p.data.add_(eps * torch.randn_like(p) / norm)
 55        model.zero_grad(set_to_none=True); loss_fn().backward()
 56        for a, ag in enumerate(blocks):
 57            ng = param_vec([p.grad for p in ag]).detach()
 58            rows[a, b] = float((ng - grads[a]).norm()) / eps
 59        for p, s in zip(params, saved): p.data.copy_(s)
 60    # Each diagonal block is approximated by its local update contraction and
 61    # off-diagonal terms by normalized cross-block gradient sensitivity.
 62    scale = max(float(np.linalg.norm(rows)), 1e-8)
 63    J = rows / scale
 64    best = None; best_rho = float('inf')
 65    for order in itertools.permutations(range(3)):
 66        M = np.eye(3)
 67        for b in order:
 68            T = np.eye(3); T[b, :] -= np.minimum(0.95, lr * (J[b, :] + np.eye(3)[b, :]))
 69            M = T @ M
 70        rho = float(np.max(np.abs(np.linalg.eigvals(M))))
 71        if rho < best_rho: best_rho, best = rho, order
 72    return list(best), best_rho, float(np.linalg.norm(rows)), rows
 73
 74
 75def idea_train(cfg, seed, keep=False):
 76    ds, model = setup(seed); device = 'cuda' if torch.cuda.is_available() else 'cpu'
 77    loss_fn = lambda: nn.MSELoss()(model(ds['xtr'].to(device)), ds['ytr'].to(device))
 78    blocks = make_blocks(model); model.to(device)
 79    blocks = [[p for p in g if p.device.type == device] for g in blocks]
 80    opts = [torch.optim.Adam(g, lr=cfg['lr'], weight_decay=cfg['weight_decay']) for g in blocks]
 81    order = [0, 1, 2]; selected_rhos=[]; hist=[]
 82    try:
 83        for ep in range(EPOCHS):
 84            # Full-batch keeps the sequential-vs-simultaneous distinction clear.
 85            if ep % 3 == 0:
 86                cand, rho, sens, mat = choose_order(model, loss_fn, blocks, cfg['lr'])
 87                selected_rhos.append(rho)
 88                if cand != order:
 89                    # hysteresis: only switch if the candidate has a meaningful
 90                    # predicted improvement over the current ordering.
 91                    def radius(ordr):
 92                        M=np.eye(3)
 93                        for b in ordr:
 94                            T=np.eye(3); T[b,:]-=np.minimum(.95,cfg['lr']*(mat[b,:]+np.eye(3)[b,:])/max(np.linalg.norm(mat),1e-8))
 95                            M=T@M
 96                        return float(np.max(np.abs(np.linalg.eigvals(M))))
 97                    if radius(order) <= 0 or radius(cand) < .95*radius(order): order=cand
 98            for b in order:
 99                for o in opts: o.zero_grad(set_to_none=True)
100                loss=loss_fn(); loss.backward(); opts[b].step()
101            hist.append(float(loss_fn().detach().cpu()))
102        with torch.no_grad(): metric=float(nn.MSELoss()(model(ds['xte'].to(device)), ds['yte'].to(device)).cpu())
103    except RuntimeError:
104        if device == 'cuda':
105            torch.cuda.empty_cache(); return idea_train_cpu(cfg, seed, keep)
106        return (float('nan'), None, ds, hist, [], []) if keep else float('nan')
107    return (metric, model, ds, hist, selected_rhos, order) if keep else metric
108
109
110def idea_train_cpu(cfg, seed, keep=False):
111    old=torch.cuda.is_available
112    # Re-run a CPU-only copy using the same intervention.
113    ds, model=setup(seed); model=model.cpu(); x,y=ds['xtr'],ds['ytr']; xt,yt=ds['xte'],ds['yte']
114    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=[]
115    for _ in range(EPOCHS):
116        for b in order:
117            for o in opts:o.zero_grad(set_to_none=True)
118            loss=nn.MSELoss()(model(x),y); loss.backward(); opts[b].step()
119        hist.append(float(loss.detach()))
120    metric=float(nn.MSELoss()(model(xt),yt)); return (metric,model,ds,hist,[],order) if keep else metric
121
122
123def evaluate(fn, seeds=SEEDS): return [float(fn(s)) for s in seeds]
124
125def main():
126    base_grid=[]
127    for c in GRID:
128        vals=evaluate(lambda s,c=c:baseline_train(c,s), tuple(range(4)))
129        base_grid.append({'cfg':c,'mean':float(np.mean(vals))})
130    best=min(base_grid,key=lambda z:z['mean'])['cfg']
131    base_full=evaluate(lambda s:baseline_train(best,s))
132    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}}
133    idea_cfgs=[best, {'lr':0.0015,'weight_decay':best['weight_decay']}, {'lr':0.006,'weight_decay':best['weight_decay']}]
134    idea_trials=[]
135    for c in idea_cfgs:
136        vals=evaluate(lambda s,c=c:idea_train(c,s),tuple(range(4))); idea_trials.append({'cfg':c,'mean':float(np.mean(vals))})
137    ibest=min(idea_trials,key=lambda z:z['mean'])['cfg']; iv=evaluate(lambda s:idea_train(ibest,s))
138    idea={'best_cfg':ibest,'sweep':idea_trials,'mean':float(np.mean(iv)),'std':float(np.std(iv)),'per_seed':iv,'n':8}
139    rep=bench.make_report('sequence','transformer_tiny',base,idea,{'predicted_vs_observed': 'trained-model finite-difference block sensitivities and selected-order contraction', 'confirmed': False})
140    rep['baseline']['full']['std']=float(np.std(base_full)); rep['idea']['std']=float(np.std(iv))
141    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
142    print(json.dumps(rep,indent=2))
143
144if __name__=='__main__': main()