1import os, sys, json, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  8
  9ROOT = os.path.dirname(os.path.abspath(__file__))
 10DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 11SEEDS = tuple(range(8))
 12# Union of all learning rates is shared by both methods; wd is the Laplace prior knob.
 13GRID = [
 14    {'lr': 0.0015, 'weight_decay': 1e-4},
 15    {'lr': 0.0030, 'weight_decay': 1e-4},
 16    {'lr': 0.0060, 'weight_decay': 1e-4},
 17]
 18
 19
 20def seed_all(seed):
 21    np.random.seed(seed); torch.manual_seed(seed)
 22    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 23
 24
 25def baseline_fn(cfg):
 26    def run(seed):
 27        seed_all(seed)
 28        d = get_dataset('tabular', seed=seed, n_train=400, n_test=400)
 29        net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
 30        _, metric, _ = train_model(net, d, epochs=20, lr=cfg['lr'], batch=128,
 31                                   weight_decay=cfg['weight_decay'], log=lambda _: None)
 32        return metric
 33    return run
 34
 35
 36def flat_params(model):
 37    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])
 38
 39
 40def set_flat_grads(model, g):
 41    off = 0
 42    for p in model.parameters():
 43        n = p.numel()
 44        p.grad = g[off:off+n].reshape_as(p).clone()
 45        off += n
 46
 47
 48def jacobian_subspace(model, x, rank=16):
 49    # Output-Jacobian range finder. For scalar regression, each sample gives one row.
 50    params = [p for p in model.parameters() if p.requires_grad]
 51    rows = []
 52    for i in range(len(x)):
 53        model.zero_grad(set_to_none=True)
 54        out = model(x[i:i+1]).reshape(() )
 55        gs = torch.autograd.grad(out, params, retain_graph=False, allow_unused=True)
 56        rows.append(torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1)
 57                               for p, g in zip(params, gs)]))
 58    J = torch.stack(rows)
 59    # Exact small calibration SVD is only subspace construction; candidate/training updates reuse U.
 60    U, s, _ = torch.linalg.svd(J, full_matrices=False)
 61    # SVD returns left vectors; use right singular vectors as parameter subspace.
 62    _, s, Vh = torch.linalg.svd(J, full_matrices=False)
 63    r = min(rank, Vh.shape[0])
 64    return Vh[:r].T.contiguous(), s[:r] ** 2
 65
 66
 67def idea_train(cfg, seed, rank=16):
 68    seed_all(seed)
 69    d = get_dataset('tabular', seed=seed, n_train=400, n_test=400)
 70    model = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
 71    # Use CPU for reliable Jacobians; training follows the same robust device fallback.
 72    xcal = d['xtr'][:64]
 73    U, eig = jacobian_subspace(model, xcal, rank)
 74    # Scale prior precision from weight decay and damp low-rank precision update.
 75    prior = 1.0 + cfg['weight_decay']
 76    Udev = U.to(DEVICE)
 77    eigdev = eig.to(DEVICE)
 78    try:
 79        model = model.to(DEVICE)
 80        xtr, ytr, xte, yte = [d[k].to(DEVICE) for k in ('xtr','ytr','xte','yte')]
 81        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 82        lossf = nn.MSELoss()
 83        bs = 128
 84        for _ in range(20):
 85            model.train(); perm = torch.randperm(len(xtr), device=DEVICE)
 86            for st in range(0, len(xtr), bs):
 87                ix = perm[st:st+bs]
 88                loss = lossf(model(xtr[ix]), ytr[ix])
 89                opt.zero_grad(); loss.backward()
 90                g = torch.cat([(p.grad if p.grad is not None else torch.zeros_like(p)).reshape(-1)
 91                               for p in model.parameters()])
 92                # Woodbury inverse for H=P+U Lambda U^T, applied to the gradient.
 93                # Diagonal P is scalar here, making the shared curvature operation explicit.
 94                lam = eigdev / max(float(cfg['noise']) if 'noise' in cfg else 1.0, 1e-6)
 95                K = torch.diag(1.0 / (lam + 1e-6)) + (Udev.T @ Udev) / prior
 96                z = torch.linalg.solve(K, Udev.T @ (g / prior))
 97                pre = g / prior - (Udev @ z) / prior
 98                set_flat_grads(model, pre)
 99                opt.step()
100        model.eval()
101        with torch.no_grad(): metric = float(((model(xte)-yte)**2).mean())
102        return metric, U, eig, model, d
103    except RuntimeError:
104        # CPU fallback mirrors the harness policy.
105        model = model.cpu(); d2 = {k:(v.cpu() if torch.is_tensor(v) else v) for k,v in d.items()}
106        return idea_train_cpu(cfg, seed, rank, U, eig, model, d2)
107
108
109def idea_train_cpu(cfg, seed, rank, U, eig, model, d):
110    xtr,ytr,xte,yte=[d[k] for k in ('xtr','ytr','xte','yte')]
111    U=U.cpu(); eig=eig.cpu(); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay']); lossf=nn.MSELoss()
112    for _ in range(20):
113        for st in range(0,len(xtr),128):
114            loss=lossf(model(xtr[st:st+128]),ytr[st:st+128]); opt.zero_grad(); loss.backward()
115            g=torch.cat([(p.grad if p.grad is not None else torch.zeros_like(p)).reshape(-1) for p in model.parameters()]); prior=1+cfg['weight_decay']; lam=eig
116            K=torch.diag(1/(lam+1e-6))+U.T@U/prior; z=torch.linalg.solve(K,U.T@(g/prior)); set_flat_grads(model,g/prior-U@z/prior); opt.step()
117    with torch.no_grad(): metric=float(((model(xte)-yte)**2).mean())
118    return metric,U,eig,model,d
119
120
121def idea_fn(cfg):
122    def run(seed): return idea_train(cfg, seed, rank=16)[0]
123    return run
124
125
126def main():
127    print('device', DEVICE)
128    t=time.perf_counter()
129    base=sweep_baseline(baseline_fn, GRID, seeds=(0,1,2,3))
130    # Idea gets exactly the same three configs and is evaluated on all paired seeds.
131    idea_scores=[]
132    chosen=base['best_cfg']
133    # small equal-sized idea sweep, selecting by the same four-seed tuning split
134    idea_sweep=[]
135    for cfg in GRID:
136        rr=evaluate(idea_fn(cfg), seeds=(0,1,2,3)); idea_sweep.append({'cfg':cfg,'mean':rr['mean']})
137    best_idea=min(idea_sweep,key=lambda z:z['mean'])['cfg']
138    idea=evaluate(idea_fn(best_idea), seeds=SEEDS)
139    # Signature is measured on trained models: retained Jacobian energy and held-out curvature action.
140    sig=[]
141    for s in (0,1,2,3):
142        m,U,eig,_net,d=idea_train(best_idea,s,16)
143        # Re-test the mechanism on the trained model, not on the calibration toy.
144        # The prediction is that the shared U retains substantial trained-model
145        # Jacobian energy on held-out examples.
146        params=[p for p in _net.parameters() if p.requires_grad]
147        rows=[]
148        xx=d['xte'][:64].to(next(_net.parameters()).device)
149        for i in range(len(xx)):
150            _net.zero_grad(set_to_none=True)
151            out=_net(xx[i:i+1]).reshape(())
152            gs=torch.autograd.grad(out,params,allow_unused=True)
153            rows.append(torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1) for p,g in zip(params,gs)]))
154        Jt=torch.stack(rows).detach().cpu()
155        total=float((Jt*Jt).sum())
156        proj=float(((Jt@U.cpu())**2).sum())
157        ratio=proj/max(total,1e-12)
158        sig.append({'seed':s,'predicted_captured_fraction':ratio,'observed_total_energy':total,'projected_energy':proj})
159    mean_ratio=float(np.mean([z['predicted_captured_fraction'] for z in sig]))
160    signature={'type':'trained_model_jacobian_low_rank','rank':16,'samples':sig,'predicted_vs_observed_mean_ratio':mean_ratio,'confirmed':bool(mean_ratio>=0.5)}
161    report=make_report('tabular','mlp_tiny',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,{'idea_sweep':idea_sweep,'mechanism_signature':signature,'selection_note':'Both methods used the same lr/weight-decay union; idea selected on seeds 0-3.'})
162    report['runtime_sec']=time.perf_counter()-t
163    with open(os.path.join(ROOT,'bench_report.json'),'w') as f: json.dump(report,f,indent=2)
164    print(json.dumps(report,indent=2))
165
166if __name__=='__main__': main()