1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from sklearn.datasets import make_moons
  6
  7SEED=133
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11
 12# ---------- exact/iterated finite-dimensional spectral sanity check ----------
 13def spectral_check():
 14    lam=np.array([1., 2., 4., 7.])
 15    rho=3.0
 16    p=(lam-rho)**2
 17    u0=np.array([.3,.4,.5,.7]); u0=u0/np.linalg.norm(u0)
 18    # exact normalized semigroup, avoiding underflow by subtracting minimum exponent
 19    times=np.array([0., .5, 1., 2., 4., 8.])
 20    traj=[]
 21    for t in times:
 22        z=u0*np.exp(-t*(p-p.min())); z=z/np.linalg.norm(z); traj.append(z)
 23    traj=np.array(traj)
 24    selected=np.where(p==p.min())[0]
 25    # Discrete displayed update with a stable small eta
 26    u=u0.copy(); eta=.04
 27    iter_errors=[]
 28    for _ in range(300):
 29        q=(lam-rho)**2*u
 30        u=u-eta*q; u=u/np.linalg.norm(u)
 31        iter_errors.append(np.linalg.norm(u[selected[0]+1:]) if False else np.linalg.norm(u[[0,3]]))
 32    # Selected eigenspaces are indices 1,2 (lambda 2 and 4), with equal minimum.
 33    unwanted=np.array([0,3]); selected=np.array([1,2])
 34    final_unwanted=np.linalg.norm(u[unwanted])
 35    ratio=.0
 36    # exact ratio of unwanted/selected should decay at gap 1
 37    ratios=[]
 38    for z in traj:
 39        ratios.append(np.linalg.norm(z[unwanted])/np.linalg.norm(z[selected]))
 40    return {'eigenvalues':lam.tolist(),'rho':rho,'polynomial_values':p.tolist(),
 41            'selected_indices':selected.tolist(),'exact_ratio_t0':ratios[0],
 42            'exact_ratio_t8':ratios[-1],'discrete_final_unwanted':float(final_unwanted),
 43            'gap':float(np.min(p[unwanted])-p.min()),
 44            'claim_observed':bool(ratios[-1] < ratios[0]*math.exp(-7.5))}
 45
 46# ---------- tiny neural experiment ----------
 47def flat(xs): return torch.cat([x.reshape(-1) for x in xs])
 48def set_flat(model, vec):
 49    pos=0
 50    with torch.no_grad():
 51        for p in model.parameters():
 52            n=p.numel(); p.copy_(vec[pos:pos+n].view_as(p)); pos+=n
 53
 54def make_model():
 55    return torch.nn.Sequential(torch.nn.Linear(2,16),torch.nn.Tanh(),torch.nn.Linear(16,2)).to(device)
 56
 57def loss_at(model,x,y): return torch.nn.functional.cross_entropy(model(x),y)
 58
 59def grad_vec(model, loss, create_graph=False):
 60    gs=torch.autograd.grad(loss, tuple(model.parameters()), create_graph=create_graph, retain_graph=True, allow_unused=False)
 61    return flat(gs)
 62
 63def hvp(model, loss, u, create_graph=False):
 64    gs=torch.autograd.grad(loss, tuple(model.parameters()), create_graph=True, retain_graph=True)
 65    dot=(flat(gs)*u).sum()
 66    hs=torch.autograd.grad(dot, tuple(model.parameters()), create_graph=create_graph, retain_graph=True)
 67    return flat(hs)
 68
 69def perturb(model, d):
 70    old=flat([p.detach() for p in model.parameters()]).clone()
 71    set_flat(model, old+d)
 72    return old
 73
 74def train(method, X, Y, steps=150, gamma=.025, lr=.03):
 75    model=make_model(); opt=torch.optim.SGD(model.parameters(),lr=lr)
 76    n=len(X); gen=torch.Generator(device=device); gen.manual_seed(SEED+7)
 77    losses=[]; dirs=[]; rayleigh=[]
 78    for step in range(steps):
 79        idx=torch.randint(n,(64,),generator=gen,device=device); x=X[idx]; y=Y[idx]
 80        opt.zero_grad(set_to_none=True)
 81        base=loss_at(model,x,y)
 82        g=grad_vec(model,base,create_graph=(method=='cb-sam'))
 83        gn=g.detach().norm().clamp_min(1e-12)
 84        if method=='sgd':
 85            base.backward(); opt.step(); losses.append(float(base.detach().cpu())); continue
 86        if method=='sam':
 87            u=(g.detach()/gn)
 88        else:
 89            # Use a robust, reproducible target near the current gradient curvature.
 90            # A short random-probe median is the stated rho-estimation option.
 91            probes=[]
 92            for j in range(2):
 93                z=torch.randn(g.numel(),device=device); z=z/z.norm()
 94                hz=hvp(model,base,z,create_graph=False).detach()
 95                probes.append(float((z*hz).sum().cpu()))
 96            rho=float(np.median(probes))
 97            u=g.detach()/gn
 98            for _ in range(3):
 99                v=hvp(model,base,u,create_graph=True)
100                w=hvp(model,base,v,create_graph=False)
101                q=w-2*rho*v.detach()+(rho*rho)*u
102                u=(u-.03*q.detach()); u=u/u.norm().clamp_min(1e-12)
103            hu=hvp(model,base,u,create_graph=False).detach()
104            rayleigh.append(float((u*hu).sum().cpu()))
105        old=perturb(model,gamma*u)
106        pert_loss=loss_at(model,x,y)
107        opt.zero_grad(set_to_none=True); pert_loss.backward(); set_flat(model,old); opt.step()
108        # Restore original parameters before applying the perturbed-loss gradient.
109        losses.append(float(base.detach().cpu())); dirs.append(float(u.detach().dot(g.detach()).cpu()))
110    with torch.no_grad():
111        test_loss=loss_at(model,X,Y).item(); pred=model(X).argmax(1); acc=(pred==Y).float().mean().item()
112        # random perturbation loss increase, a small directional-sharpness proxy
113        vals=[]
114        w=flat([p.detach() for p in model.parameters()]);
115        for j in range(8):
116            z=torch.randn_like(w); z=z/z.norm(); old=w.clone(); set_flat(model,w+gamma*z)
117            vals.append(loss_at(model,X,Y).item()-test_loss); set_flat(model,w)
118        sharp=float(np.mean(vals))
119    return {'train_loss_last':float(losses[-1]) if losses else None,'test_loss':test_loss,'accuracy':acc,
120            'random_loss_increase':sharp,'mean_u_dot_g':float(np.mean(dirs)) if dirs else None,
121            'mean_u_rayleigh':float(np.mean(rayleigh)) if rayleigh else None}
122
123def main():
124    check=spectral_check()
125    # fixed full dataset, identical minibatch sequence across methods
126    xa,ya=make_moons(n_samples=256,noise=.18,random_state=SEED)
127    X=torch.tensor(xa,dtype=torch.float32,device=device); Y=torch.tensor(ya,dtype=torch.long,device=device)
128    results={}
129    t=time.time()
130    for m in ('sgd','sam','cb-sam'):
131        # reset seed before each model for a matched initialization
132        torch.manual_seed(SEED); results[m]=train(m,X,Y)
133    out={'device':str(device),'spectral_check':check,'results':results,'seconds':time.time()-t}
134    Path('results.json').write_text(json.dumps(out,indent=2))
135    print(json.dumps(out,indent=2))
136if __name__=='__main__':
137    try: main()
138    except Exception as e:
139        if device.type=='cuda':
140            print('CUDA failed, rerun CPU:',repr(e)); device=torch.device('cpu'); main()
141        else: raise