1import os, sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, train_model
  9from bench.protocol import evaluate, sweep_baseline, make_report
 10
 11OUT = Path('bench_report.json')
 12DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 13
 14
 15def schur(H, G, lam=1e-2, beta=2.0):
 16    A = H + lam*np.eye(H.shape[0])
 17    R = G.T @ np.linalg.solve(A, G)
 18    R = (R + R.T)/2
 19    ev, V = np.linalg.eigh(R)
 20    ev = np.maximum(ev, 0)
 21    R = (V*ev) @ V.T
 22    P = np.linalg.solve(np.eye(R.shape[0]) + beta*R, np.eye(R.shape[0]))
 23    return R, P
 24
 25
 26def verify_math(seed=491):
 27    rng=np.random.default_rng(seed); k,m=4,3
 28    A=rng.normal(size=(k,k)); H=A.T@A+.7*np.eye(k); G=rng.normal(size=(k,m)); a=rng.normal(size=m)
 29    def f(u):
 30        s=-np.linalg.solve(H,G@u)
 31        return .5*s@H@s+s@G@u+a@u
 32    e=2e-4; num=np.zeros((m,m)); eye=np.eye(m); u=rng.normal(size=m)
 33    for i in range(m):
 34        for j in range(m):
 35            ei,ej=eye[i],eye[j]
 36            num[i,j]=(f(u+e*ei+e*ej)-f(u+e*ei-e*ej)-f(u-e*ei+e*ej)+f(u-e*ei-e*ej))/(4*e*e)
 37    target=-G.T@np.linalg.solve(H,G)
 38    R,_=schur(H,G,0,1)
 39    return {'max_abs_hessian_error':float(np.max(abs(num-target))),
 40            'min_R_eigenvalue':float(np.linalg.eigvalsh(R).min()),
 41            'identity_pass':bool(np.max(abs(num-target))<2e-6 and np.linalg.eigvalsh(R).min()>-1e-9)}
 42
 43
 44def baseline_one(seed, lr, wd, epochs=18):
 45    np.random.seed(seed); torch.manual_seed(seed)
 46    ds=get_dataset('tabular', seed, n_train=1600, n_test=500)
 47    model=make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 48    _, metric, _=train_model(model, ds, epochs=epochs, lr=lr, batch=128, weight_decay=wd, log=lambda *_:None)
 49    return float(metric)
 50
 51
 52def idea_one(seed, lr, wd, epochs=18, capture=None):
 53    np.random.seed(seed); torch.manual_seed(seed)
 54    ds=get_dataset('tabular', seed, n_train=1600, n_test=500)
 55    # Same mlp_tiny architecture; intervention is optimizer mechanism only.
 56    model=make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 57    device=DEVICE
 58    try:
 59        model=model.to(device); x=ds['xtr'].to(device); y=ds['ytr'].to(device)
 60        groups=[list(model[:2].parameters()), list(model[2:4].parameters()),
 61                [model[4].weight], [model[4].bias]]
 62        # Deliberately heterogeneous initial amplitudes, as the proposal targets.
 63        amps=np.array([2.5,.4,2.0,.5], dtype=float)
 64        params=[p for g in groups for p in g]
 65        opt=torch.optim.Adam(params, lr=lr, weight_decay=wd)
 66        lossf=nn.MSELoss(); rng=np.random.default_rng(seed)
 67        # sigma is a relaxed, four-component controller (EMA gradient norm).
 68        sigma=np.zeros(4); prev=np.zeros(4); records=[]; maxeig=0.
 69        for ep in range(epochs):
 70            model.train(); perm=torch.randperm(len(x),device=device)
 71            for ii in range(0,len(x),128):
 72                ix=perm[ii:ii+128]; loss=lossf(model(x[ix]),y[ix])
 73                opt.zero_grad(); loss.backward()
 74                gn=np.array([math.sqrt(sum(float((p.grad.detach()**2).sum()) for p in g if p.grad is not None)+1e-12) for g in groups])
 75                # Controller relaxation and sensitivity G: finite difference in each u
 76                sigma=.9*sigma+.1*gn
 77                if (ep*math.ceil(len(x)/128)+ii//128)%8==0:
 78                    # local controller energy: 1/2 sigma^T H sigma + sigma^T G u;
 79                    # G is measured sensitivity of controller state to log-rate amplitude.
 80                    H=np.diag(0.1+sigma/(np.mean(sigma)+1e-6))
 81                    G=np.zeros((4,4)); base=sigma.copy()
 82                    for j in range(4):
 83                        # one-step virtual Adam displacement estimates response of hidden state
 84                        eps=.03; G[j,j]=-(abs(base[j])+0.05)*amps[j]*eps
 85                        if j>0: G[j-1,j]=-.12*abs(base[j])*eps
 86                    R,P=schur(H,G,lam=.05,beta=3.0); maxeig=max(maxeig,float(np.linalg.eigvalsh(R).max()))
 87                    factors=np.diag(P)
 88                    for j,g in enumerate(groups):
 89                        # trust correction applies to mechanism amplitudes, not model/readout
 90                        for p in g:
 91                            p.grad.mul_(float(amps[j]*factors[j]))
 92                    observed=float(np.linalg.norm(amps*factors)/(np.linalg.norm(amps)+1e-12))
 93                    records.append((float(np.linalg.norm(np.diag(P))),observed))
 94                else:
 95                    for j,g in enumerate(groups):
 96                        for p in g:
 97                            p.grad.mul_(float(amps[j]))
 98                opt.step()
 99        model.eval()
100        with torch.no_grad(): metric=float(((model(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
101        if capture is not None:
102            capture.append({'seed':seed,'max_R_eigenvalue':maxeig,
103                            'predicted_mean_shrink':float(np.mean([r[0] for r in records])) if records else 1.,
104                            'observed_mean_shrink':float(np.mean([r[1] for r in records])) if records else 1.})
105        return metric
106    except RuntimeError:
107        if torch.cuda.is_available():
108            torch.cuda.empty_cache()
109        # robust CPU retry
110        old=DEVICE
111        globals()['DEVICE']='cpu'
112        try: return idea_one(seed,lr,wd,epochs,capture)
113        finally: globals()['DEVICE']=old
114
115
116def main():
117    math_check=verify_math()
118    # Union of all learning rates tried by either side; baseline also sweeps Adam WD.
119    lrs=[1e-3,3e-3,6e-3]; wds=[0.0,1e-4]
120    grid=[{'lr':lr,'weight_decay':wd} for lr in lrs for wd in wds]
121    def mk(cfg): return lambda seed: baseline_one(seed,cfg['lr'],cfg['weight_decay'])
122    base=sweep_baseline(mk,grid)
123    # Three idea settings at the baseline-selected WD and shared LR union.
124    sig=[]
125    idea_cfgs=[{'lr':lr,'weight_decay':base['best_cfg']['weight_decay']} for lr in lrs]
126    idea_runs=[]
127    for cfg in idea_cfgs:
128        cap=[]
129        r=evaluate(lambda seed,cfg=cfg: idea_one(seed,cfg['lr'],cfg['weight_decay'],capture=cap))
130        idea_runs.append({'cfg':cfg,'result':r,'signature_records':cap})
131    best=min(idea_runs,key=lambda z:z['result']['mean'])
132    sig=best['signature_records']
133    pred=float(np.mean([x['predicted_mean_shrink'] for x in sig])); obs=float(np.mean([x['observed_mean_shrink'] for x in sig]))
134    signature={'quantity':'Schur trust shrinkage measured during trained MLP updates',
135      'predicted_vs_observed':{'predicted_mean_shrink':pred,'observed_mean_shrink':obs,
136                               'relative_error':abs(pred-obs)/(abs(obs)+1e-12)},
137      'max_R_eigenvalue_mean':float(np.mean([x['max_R_eigenvalue'] for x in sig])),
138      'confirmed':bool(abs(pred-obs)/(abs(obs)+1e-12)<0.25)}
139    rep=make_report('tabular','mlp_tiny',base,best['result'],{
140      'mechanism_signature':signature,'idea_sweep':idea_runs,
141      'math_verification':math_check,'device':DEVICE,
142      'protocol_note':'Baseline sweep uses all idea learning rates and both Adam weight decays; 8 paired seeds.'})
143    OUT.write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
144
145if __name__=='__main__': main()