1import json, math, time
  2import numpy as np
  3import torch
  4
  5SEED = 17
  6np.random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_default_dtype(torch.float64)
  8DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
  9try:
 10    torch.tensor([1.0], device=DEVICE).sum().item()
 11except Exception:
 12    DEVICE = "cpu"
 13
 14
 15def core_sanity():
 16    # A residual with one nearly unidentifiable direction. Full inversion amplifies
 17    # its tiny residual, while the rank-safe inverse discards that direction.
 18    s = torch.tensor([1.0, 1e-9])
 19    J = torch.diag(s)
 20    r = torch.tensor([0.1, 0.1])
 21    full = -torch.linalg.pinv(J) @ r
 22    keep = s >= 1e-6 * s[0]
 23    safe = -(torch.diag(torch.where(keep, 1/s, torch.zeros_like(s))) @ r)
 24    # A first-order linear residual after each proposed step.
 25    full_res = r + J @ full
 26    safe_res = r + J @ safe
 27    return {"singular_values": s.tolist(), "full_step_norm": float(torch.linalg.norm(full)),
 28            "truncated_step_norm": float(torch.linalg.norm(safe)),
 29            "full_linearized_residual": float(torch.linalg.norm(full_res)),
 30            "truncated_linearized_residual": float(torch.linalg.norm(safe_res)),
 31            "kept_rank": int(keep.sum())}
 32
 33
 34def make_data(ntrain=96, ntest=256):
 35    xtr = torch.linspace(-1, 1, ntrain, device=DEVICE).reshape(-1,1)
 36    xte = torch.linspace(-1, 1, ntest, device=DEVICE).reshape(-1,1)
 37    def target(x):
 38        return torch.where(x < 0, torch.sin(4*x), 0.55*torch.sin(7*x)+0.25*x)
 39    ytr = target(xtr) + 0.025*torch.randn_like(xtr)
 40    return xtr, ytr, xte, target(xte)
 41
 42
 43def unpack(z, width):
 44    # tanh hidden layer, feature vector includes a constant output bias
 45    a = z[:width].reshape(width,1)
 46    b = z[width:2*width]
 47    return a, b
 48
 49def features(z, x, width):
 50    a,b = unpack(z,width)
 51    h = torch.tanh(x @ a.T + b)
 52    return torch.cat([h, torch.ones((x.shape[0],1),device=x.device)], dim=1)
 53
 54def projected(z, x, y, width, tau=1e-6):
 55    Phi = features(z,x,width)
 56    U,S,Vh = torch.linalg.svd(Phi, full_matrices=False)
 57    cutoff = tau*S[0] if S.numel() else 0.
 58    inv = torch.where(S >= cutoff, 1/S, torch.zeros_like(S))
 59    w = Vh.T @ (inv[:,None] * (U.T @ y))
 60    r = Phi @ w-y
 61    return w, r, S
 62
 63def vp_gn(x,y,xt,yt,width=16,steps=35,rho=1e-4,tau=1e-6):
 64    z = (0.4*torch.randn(2*width,device=DEVICE)).requires_grad_()
 65    hist=[]; accepted=[]; conds=[]
 66    t0=time.perf_counter()
 67    for it in range(steps):
 68        w,r,S = projected(z,x,y,width,tau)
 69        loss=0.5*torch.sum(r*r)
 70        def residual(zz): return projected(zz,x,y,width,tau)[0] # not used; avoid differentiating solve
 71        # Jacobian of residual with output w held fixed (Gauss-Newton VP/envelope step)
 72        ww=w.detach()
 73        def fixed_res(zz): return (features(zz,x,width) @ ww-y).reshape(-1)
 74        J=torch.autograd.functional.jacobian(fixed_res,z,create_graph=False).reshape(x.shape[0], -1)
 75        uj,sj,vhj=torch.linalg.svd(J,full_matrices=False)
 76        conds.append(float(sj[0]/max(float(sj[-1]),1e-30)))
 77        keep=sj >= rho*sj[0]
 78        coeff = torch.where(keep, 1.0/sj, torch.zeros_like(sj))
 79        rhs = torch.mv(uj.T, r.reshape(-1))
 80        delta = -torch.mv(vhj.T, coeff * rhs)
 81        base=float(loss)
 82        alpha=1.0; ok=False
 83        for _ in range(9):
 84            with torch.no_grad():
 85                ztry=z+alpha*delta
 86                wtry,rtry,_=projected(ztry,x,y,width,tau)
 87                if float(0.5*torch.sum(rtry*rtry)) < base:
 88                    z=ztry.detach().requires_grad_(); ok=True; break
 89            alpha*=0.5
 90        if not ok: alpha=0.0
 91        accepted.append(alpha)
 92        with torch.no_grad():
 93            wp,rp,_=projected(z,x,y,width,tau)
 94            pred=features(z,xt,width)@wp
 95            hist.append((float(0.5*torch.sum(rp*rp)/len(y)),float(torch.mean((pred-yt)**2)),int(keep.sum())))
 96    return {"train_loss":hist[-1][0],"test_mse":hist[-1][1],"initial_train_loss":hist[0][0],
 97            "accepted_steps":sum(a>0 for a in accepted),"median_alpha":float(np.median(accepted)),
 98            "max_jacobian_condition":max(conds),"history":hist,"seconds":time.perf_counter()-t0}
 99
100def adam_run(x,y,xt,yt,width=16,steps=350):
101    z=(0.4*torch.randn(2*width,device=DEVICE)).requires_grad_()
102    w=(0.1*torch.randn(width+1,device=DEVICE)).requires_grad_()
103    opt=torch.optim.Adam([z,w],lr=0.025)
104    hist=[]; t0=time.perf_counter()
105    for _ in range(steps):
106        opt.zero_grad()
107        pred=features(z,x,width)@w
108        loss=torch.mean((pred-y)**2)/2
109        loss.backward(); opt.step()
110        with torch.no_grad(): hist.append((float(loss),float(torch.mean((features(z,xt,width)@w-yt)**2))))
111    return {"train_loss":hist[-1][0],"test_mse":hist[-1][1],"initial_train_loss":hist[0][0],
112            "history":hist,"seconds":time.perf_counter()-t0}
113
114def main():
115    sanity=core_sanity()
116    x,y,xt,yt=make_data()
117    # identical initial seed streams are controlled separately, but methods have
118    # different parameterizations; report repeated fixed-seed single runs.
119    torch.manual_seed(SEED+1); adam=adam_run(x,y,xt,yt)
120    torch.manual_seed(SEED+2); vp=vp_gn(x,y,xt,yt)
121    out={"device":DEVICE,"seed":SEED,"core_sanity":sanity,"adam":adam,"rank_safe_vp_gn":vp}
122    with open("results.json","w") as f: json.dump(out,f,indent=2)
123    print(json.dumps({k:v for k,v in out.items() if k not in ('adam','rank_safe_vp_gn')},indent=2))
124    print(json.dumps({"adam":{k:adam[k] for k in adam if k!='history'},"rank_safe_vp_gn":{k:vp[k] for k in vp if k!='history'}},indent=2))
125if __name__=='__main__': main()