import json, math, random, time from pathlib import Path import numpy as np import torch from sklearn.datasets import make_moons SEED=133 random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED) torch.set_num_threads(4) device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') # ---------- exact/iterated finite-dimensional spectral sanity check ---------- def spectral_check(): lam=np.array([1., 2., 4., 7.]) rho=3.0 p=(lam-rho)**2 u0=np.array([.3,.4,.5,.7]); u0=u0/np.linalg.norm(u0) # exact normalized semigroup, avoiding underflow by subtracting minimum exponent times=np.array([0., .5, 1., 2., 4., 8.]) traj=[] for t in times: z=u0*np.exp(-t*(p-p.min())); z=z/np.linalg.norm(z); traj.append(z) traj=np.array(traj) selected=np.where(p==p.min())[0] # Discrete displayed update with a stable small eta u=u0.copy(); eta=.04 iter_errors=[] for _ in range(300): q=(lam-rho)**2*u u=u-eta*q; u=u/np.linalg.norm(u) iter_errors.append(np.linalg.norm(u[selected[0]+1:]) if False else np.linalg.norm(u[[0,3]])) # Selected eigenspaces are indices 1,2 (lambda 2 and 4), with equal minimum. unwanted=np.array([0,3]); selected=np.array([1,2]) final_unwanted=np.linalg.norm(u[unwanted]) ratio=.0 # exact ratio of unwanted/selected should decay at gap 1 ratios=[] for z in traj: ratios.append(np.linalg.norm(z[unwanted])/np.linalg.norm(z[selected])) return {'eigenvalues':lam.tolist(),'rho':rho,'polynomial_values':p.tolist(), 'selected_indices':selected.tolist(),'exact_ratio_t0':ratios[0], 'exact_ratio_t8':ratios[-1],'discrete_final_unwanted':float(final_unwanted), 'gap':float(np.min(p[unwanted])-p.min()), 'claim_observed':bool(ratios[-1] < ratios[0]*math.exp(-7.5))} # ---------- tiny neural experiment ---------- def flat(xs): return torch.cat([x.reshape(-1) for x in xs]) def set_flat(model, vec): pos=0 with torch.no_grad(): for p in model.parameters(): n=p.numel(); p.copy_(vec[pos:pos+n].view_as(p)); pos+=n def make_model(): return torch.nn.Sequential(torch.nn.Linear(2,16),torch.nn.Tanh(),torch.nn.Linear(16,2)).to(device) def loss_at(model,x,y): return torch.nn.functional.cross_entropy(model(x),y) def grad_vec(model, loss, create_graph=False): gs=torch.autograd.grad(loss, tuple(model.parameters()), create_graph=create_graph, retain_graph=True, allow_unused=False) return flat(gs) def hvp(model, loss, u, create_graph=False): gs=torch.autograd.grad(loss, tuple(model.parameters()), create_graph=True, retain_graph=True) dot=(flat(gs)*u).sum() hs=torch.autograd.grad(dot, tuple(model.parameters()), create_graph=create_graph, retain_graph=True) return flat(hs) def perturb(model, d): old=flat([p.detach() for p in model.parameters()]).clone() set_flat(model, old+d) return old def train(method, X, Y, steps=150, gamma=.025, lr=.03): model=make_model(); opt=torch.optim.SGD(model.parameters(),lr=lr) n=len(X); gen=torch.Generator(device=device); gen.manual_seed(SEED+7) losses=[]; dirs=[]; rayleigh=[] for step in range(steps): idx=torch.randint(n,(64,),generator=gen,device=device); x=X[idx]; y=Y[idx] opt.zero_grad(set_to_none=True) base=loss_at(model,x,y) g=grad_vec(model,base,create_graph=(method=='cb-sam')) gn=g.detach().norm().clamp_min(1e-12) if method=='sgd': base.backward(); opt.step(); losses.append(float(base.detach().cpu())); continue if method=='sam': u=(g.detach()/gn) else: # Use a robust, reproducible target near the current gradient curvature. # A short random-probe median is the stated rho-estimation option. probes=[] for j in range(2): z=torch.randn(g.numel(),device=device); z=z/z.norm() hz=hvp(model,base,z,create_graph=False).detach() probes.append(float((z*hz).sum().cpu())) rho=float(np.median(probes)) u=g.detach()/gn for _ in range(3): v=hvp(model,base,u,create_graph=True) w=hvp(model,base,v,create_graph=False) q=w-2*rho*v.detach()+(rho*rho)*u u=(u-.03*q.detach()); u=u/u.norm().clamp_min(1e-12) hu=hvp(model,base,u,create_graph=False).detach() rayleigh.append(float((u*hu).sum().cpu())) old=perturb(model,gamma*u) pert_loss=loss_at(model,x,y) opt.zero_grad(set_to_none=True); pert_loss.backward(); set_flat(model,old); opt.step() # Restore original parameters before applying the perturbed-loss gradient. losses.append(float(base.detach().cpu())); dirs.append(float(u.detach().dot(g.detach()).cpu())) with torch.no_grad(): test_loss=loss_at(model,X,Y).item(); pred=model(X).argmax(1); acc=(pred==Y).float().mean().item() # random perturbation loss increase, a small directional-sharpness proxy vals=[] w=flat([p.detach() for p in model.parameters()]); for j in range(8): z=torch.randn_like(w); z=z/z.norm(); old=w.clone(); set_flat(model,w+gamma*z) vals.append(loss_at(model,X,Y).item()-test_loss); set_flat(model,w) sharp=float(np.mean(vals)) return {'train_loss_last':float(losses[-1]) if losses else None,'test_loss':test_loss,'accuracy':acc, 'random_loss_increase':sharp,'mean_u_dot_g':float(np.mean(dirs)) if dirs else None, 'mean_u_rayleigh':float(np.mean(rayleigh)) if rayleigh else None} def main(): check=spectral_check() # fixed full dataset, identical minibatch sequence across methods xa,ya=make_moons(n_samples=256,noise=.18,random_state=SEED) X=torch.tensor(xa,dtype=torch.float32,device=device); Y=torch.tensor(ya,dtype=torch.long,device=device) results={} t=time.time() for m in ('sgd','sam','cb-sam'): # reset seed before each model for a matched initialization torch.manual_seed(SEED); results[m]=train(m,X,Y) out={'device':str(device),'spectral_check':check,'results':results,'seconds':time.time()-t} Path('results.json').write_text(json.dumps(out,indent=2)) print(json.dumps(out,indent=2)) if __name__=='__main__': try: main() except Exception as e: if device.type=='cuda': print('CUDA failed, rerun CPU:',repr(e)); device=torch.device('cpu'); main() else: raise