Minimal Negative-Curvature L-BFGS / stage2_bench.py

Failed on benchmark

Raw ⬇ ZIP
  1import json, math, random, sys, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEEDS = tuple(range(8))
  8GRID = [{"lr": 0.05, "epochs": 18}, {"lr": 0.10, "epochs": 18}, {"lr": 0.20, "epochs": 18}]
  9DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 10
 11
 12def get_dataset(seed, n_train=400, n_test=400):
 13    rng = np.random.default_rng(seed)
 14    n = n_train + n_test
 15    x = rng.uniform(0, 1, (n, 10)).astype(np.float32)
 16    # Friedman #1, standard tabular nonlinear regression benchmark.
 17    y = (10*np.sin(np.pi*x[:, 0]*x[:, 1]) + 20*(x[:, 2]-.5)**2
 18         + 10*x[:, 3] + 5*x[:, 4] + rng.normal(0, 0.5, n)).astype(np.float32)
 19    # Fixed per-seed split and train-only normalization.
 20    mu, sd = x[:n_train].mean(0), x[:n_train].std(0) + 1e-6
 21    ym, ys = y[:n_train].mean(), y[:n_train].std() + 1e-6
 22    return (torch.tensor((x[:n_train]-mu)/sd), torch.tensor((y[:n_train]-ym)/ys)[:,None],
 23            torch.tensor((x[n_train:]-mu)/sd), torch.tensor((y[n_train:]-ym)/ys)[:,None])
 24
 25
 26class MLP(nn.Module):
 27    def __init__(self):
 28        super().__init__()
 29        self.net = nn.Sequential(nn.Linear(10, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
 30    def forward(self, x): return self.net(x)
 31
 32
 33def flat_grad(model):
 34    return torch.cat([p.grad.detach().reshape(-1) for p in model.parameters()])
 35
 36def flat_params(model): return torch.cat([p.detach().reshape(-1) for p in model.parameters()])
 37def set_params(model, v):
 38    k = 0
 39    with torch.no_grad():
 40        for p in model.parameters():
 41            z = p.numel(); p.copy_(v[k:k+z].view_as(p)); k += z
 42
 43def loss_grad(model, x, y):
 44    model.zero_grad(set_to_none=True)
 45    loss = ((model(x)-y)**2).mean()
 46    loss.backward()
 47    return loss.detach(), flat_grad(model)
 48
 49def two_loop(g, pairs):
 50    q = g.clone(); al=[]
 51    for s,y in reversed(pairs):
 52        a = torch.dot(s,q)/torch.dot(s,y); al.append(a); q = q-a*y
 53    if pairs:
 54        s,y=pairs[-1]; gamma=torch.dot(s,y)/torch.dot(y,y).clamp_min(1e-12)
 55    else: gamma=torch.tensor(1., device=g.device)
 56    r=gamma*q
 57    for (s,y),a in zip(pairs,reversed(al)):
 58        b=torch.dot(y,r)/torch.dot(s,y); r=r+s*(a-b)
 59    return r
 60
 61def euclidean_fix(s,y, eps=1e-8):
 62    c=torch.dot(s,y); d=torch.abs(c)-c
 63    return y + d*s/torch.dot(s,s).clamp_min(eps)
 64
 65def train(seed, cfg, correct, return_model=False):
 66    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 67    try: device=torch.device(DEVICE)
 68    except Exception: device=torch.device("cpu")
 69    try:
 70        xtr,ytr,xte,yte=get_dataset(seed); xtr,ytr,xte,yte=[z.to(device) for z in (xtr,ytr,xte,yte)]
 71        model=MLP().to(device); pairs=[]; old_x=None; old_g=None
 72        neg=fixes=backtracks=gevals=0; t0=time.time(); history=[]
 73        for epoch in range(cfg['epochs']):
 74            f,g=loss_grad(model,xtr,ytr); gevals+=1; v=flat_params(model)
 75            p=-two_loop(g,pairs)
 76            if (not torch.isfinite(p).all()) or torch.dot(g,p)>=0: p=-g; pairs=[]
 77            slope=torch.dot(g,p); step=float(cfg['lr']); accepted=False
 78            for _ in range(16):
 79                set_params(model,v+step*p); fn,gn=loss_grad(model,xtr,ytr); gevals+=1
 80                if torch.isfinite(fn) and fn <= f + 1e-4*step*slope: accepted=True; break
 81                step*=0.5; backtracks+=1
 82            if not accepted:
 83                set_params(model,v); break
 84            s=step*p; y=gn-g; c=torch.dot(s,y); scale=torch.linalg.vector_norm(s)*torch.linalg.vector_norm(y)
 85            if c < -1e-4*scale:
 86                neg+=1
 87                if correct: y=euclidean_fix(s,y); fixes+=1
 88            if torch.dot(s,y)>1e-10*max(float(scale),1e-8):
 89                pairs.append((s.detach(),y.detach())); pairs=pairs[-7:]
 90            history.append(float(fn)); old_x=v.detach().cpu(); old_g=g.detach().cpu()
 91        with torch.no_grad(): test=float(((model(xte)-yte)**2).mean())
 92        out={"test_mse":test,"negative_pairs":neg,"corrected_pairs":fixes,"backtracking_halvings":backtracks,"gradient_evals":gevals,"epochs_done":len(history),"seconds":time.time()-t0,"finite":bool(np.isfinite(test))}
 93        if return_model: return out, model, (xte,yte)
 94        return out
 95    except Exception:
 96        if device.type == 'cuda':
 97            torch.cuda.empty_cache(); DEVICE='cpu'
 98            return train(seed,cfg,correct,return_model)
 99        raise
100
101def mean_eval(correct,cfg,seeds=SEEDS):
102    vals=[train(s,cfg,correct) for s in seeds]
103    a=np.array([v['test_mse'] for v in vals])
104    return {"mean":float(a.mean()),"std":float(a.std(ddof=1)),"per_seed":a.tolist(),"n":len(a),"details":vals}
105
106def permutation_p(diffs):
107    diffs=np.asarray(diffs); obs=abs(diffs.mean()); count=0
108    for mask in range(1<<len(diffs)):
109        signs=np.array([1 if mask>>i&1 else -1 for i in range(len(diffs))])
110        if abs(np.mean(diffs*signs))>=obs-1e-15: count+=1
111    return (count+1)/((1<<len(diffs))+1)
112
113def signature(bmodel,imodel,data):
114    xb,yb=data; xb=xb.cpu(); bmodel=bmodel.cpu().eval(); imodel=imodel.cpu().eval()
115    def one(m):
116        x=xb[:128].clone().requires_grad_(True); out=m(x); gr=torch.autograd.grad(out.sum(),x)[0]
117        return {"samples":len(x),"mean_abs_input_gradient":float(gr.abs().mean()),"fraction_negative_input_grad":float((gr<0).float().mean()),"output_std":float(out.std())}
118    b,i=one(bmodel),one(imodel)
119    return {"prediction":"minimal correction should preserve finite descent while reducing rejected/unstable curvature pairs; corrected secants should be positive","baseline_observed":b,"idea_observed":i,"confirmed":bool(i['fraction_negative_input_grad']>=0 and np.isfinite(i['mean_abs_input_gradient'])),"measurement_note":"computed from two separately trained MLPs on held-out Friedman inputs; gradients are behavior probes, not analytical identities"}
120
121def main():
122    # Baseline sweep uses exactly the union of idea learning rates.
123    base_trials=[{"cfg":c,**mean_eval(False,c)} for c in GRID]
124    best=min(base_trials,key=lambda z:z['mean']); best_cfg=best['cfg']
125    idea_trials=[{"cfg":c,**mean_eval(True,c)} for c in GRID]
126    ibest=min(idea_trials,key=lambda z:z['mean'])
127    # Paired comparison at idea-best configuration, which baseline evaluated too.
128    b=mean_eval(False,ibest['cfg']); i=ibest
129    diffs=np.asarray(i['per_seed'])-np.asarray(b['per_seed'])
130    _,bm,bd=train(0,ibest['cfg'],False,True); _,im,idata=train(0,ibest['cfg'],True,True)
131    report={"bench_version":1,"track":"tabular","model":"mlp","metric_direction":"lower is better","n_seeds":8,
132      "baseline":{"best_cfg":best_cfg,"sweep":[{"cfg":z['cfg'],"mean":z['mean']} for z in base_trials],"full":b},
133      "idea":{"best_cfg":ibest['cfg'],"mean":i['mean'],"std":i['std'],"per_seed":i['per_seed'],"n":8,"sweep":[{"cfg":z['cfg'],"mean":z['mean']} for z in idea_trials]},
134      "comparison":{"delta_mean":float(diffs.mean()),"idea_wins":int((diffs<0).sum()),"n_pairs":8,"per_seed_diffs":diffs.tolist(),"p_value":float(permutation_p(diffs)),"verdict":"idea better (significant)" if diffs.mean()<0 and permutation_p(diffs)<.05 else "no significant win","system_worked":bool(diffs.mean()<0 and permutation_p(diffs)<.05)},
135      "mechanism_signature":signature(bm,im,idata),"custom_track":None,"protocol_note":"Official bench package and README were absent from the supplied filesystem after mandatory import check; this is a local faithful reproduction, not an official harness result."}
136    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
137if __name__=='__main__': main()