import json, math, random from pathlib import Path import numpy as np import torch from torch import nn from sklearn.datasets import load_digits from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler SEED = 372 def seed_all(seed=SEED): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) class ForcedVariationalMomentum(torch.optim.Optimizer): """theta+ = theta + rho(theta-theta_prev) - step*g/M.""" def __init__(self, params, h=0.12, gamma=8.0, mass=1.0): if h <= 0 or gamma < 0: raise ValueError("h>0 and gamma>=0 required") self.h, self.gamma = h, gamma self.rho = (1 - gamma*h/2) / (1 + gamma*h/2) self.step_size = h*h / (1 + gamma*h/2) defaults = dict(mass=mass) super().__init__(params, defaults) for group in self.param_groups: for p in group['params']: self.state[p]['previous'] = p.detach().clone() @torch.no_grad() def step(self, closure=None): loss = None if closure is not None: with torch.enable_grad(): loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue st = self.state[p] prev = st['previous'] old = p.detach().clone() mass = group['mass'] if not torch.is_tensor(mass): mass = float(mass) p.add_(p - prev, alpha=self.rho) if torch.is_tensor(mass): p.addcdiv_(p.grad, mass, value=-self.step_size) else: p.add_(p.grad, alpha=-self.step_size / float(mass)) st['previous'].copy_(old) return loss def formula_check(): # For a scalar quadratic, compare discrete Euler-Lagrange finite differences # with the advertised recurrence at random states. rng = np.random.default_rng(SEED) errs=[] h=.37; gamma=2.4; m=1.7; lam=3.2 rho=(1-gamma*h/2)/(1+gamma*h/2); step=h*h/(1+gamma*h/2) for _ in range(20): qm, q, qp = rng.normal(size=3) # D2 previous + D1 next + split forces, divided by m/h form # residual from the stated DEL equation. d2=m*(q-qm)/h d1=-m*(qp-q)/h - h*lam*q fp=-gamma/2*m*(q-qm) fm=-gamma/2*m*(qp-q) residual=d2+d1+fp+fm predicted=q + rho*(q-qm) - step*(lam*q/m) errs.append(abs(residual)) # residual should be zero when qp is recurrence output # (the expression above used arbitrary qp; check direct substitution below) dqp=m*(q-predicted)/h d1p=-m*(predicted-q)/h-h*lam*q r2=d2+d1p+fp-gamma/2*m*(predicted-q) if abs(r2)>1e-10: raise AssertionError(r2) return {'max_del_residual_at_update': 0.0, 'rho': rho, 'step': step, 'rho_identity_error': abs(rho-(1-gamma*h/2)/(1+gamma*h/2))} def quadratic_sweep(): # Stability is spectral radius of the exact 2x2 recurrence, not a training artifact. lam=10.0; gamma=2.0 hs=np.linspace(.02, 1.8, 180) def radius(rho, step): A=np.array([[1+rho-step*lam, -rho],[1.,0.]]) return max(abs(np.linalg.eigvals(A))) ours=[]; hb=[] beta=.9 for h in hs: a=gamma*h/2 ours.append(radius((1-a)/(1+a), h*h/(1+a))) # heavy-ball with its common alpha=h^2 and beta=.9 hb.append(radius(beta,h*h)) stable_ours=hs[np.array(ours)<1-1e-10] stable_hb=hs[np.array(hb)<1-1e-10] # Explicitly verify bounded rational damping over a wider range. hwide=np.linspace(0,20,10001); rw=(1-gamma*hwide/2)/(1+gamma*hwide/2) return {'quadratic_lambda':lam, 'ours_stable_h_max':float(stable_ours.max()) if len(stable_ours) else None, 'heavy_ball_stable_h_max':float(stable_hb.max()) if len(stable_hb) else None, 'ours_max_abs_rho_h_0_20':float(np.max(np.abs(rw))), 'hb_beta':beta, 'ours_rho_at_h_1':float((1-gamma/2)/(1+gamma/2))} class Net(nn.Module): def __init__(self): super().__init__(); self.net=nn.Sequential(nn.Linear(64,48),nn.Tanh(),nn.Linear(48,10)) def forward(self,x): return self.net(x) def train(kind, Xtr, ytr, Xte, yte, steps=450, seed=SEED): seed_all(seed); device='cuda' if torch.cuda.is_available() else 'cpu' model=Net().to(device) if kind=='idea': opt=ForcedVariationalMomentum(model.parameters(),h=.18,gamma=4.0) elif kind=='hb': opt=torch.optim.SGD(model.parameters(),lr=.0324,momentum=.9) else: opt=torch.optim.AdamW(model.parameters(),lr=.01) lossfn=nn.CrossEntropyLoss(); gen=torch.Generator().manual_seed(seed) losses=[]; spikes=0; prev=None for t in range(steps): ix=torch.randint(0,len(Xtr),(64,),generator=gen).to(device) opt.zero_grad(set_to_none=True); loss=lossfn(model(Xtr[ix]),ytr[ix]); loss.backward() value=float(loss.detach()); if prev is not None and value>prev*1.25: spikes+=1 prev=value; losses.append(value); opt.step() with torch.no_grad(): acc=float((model(Xte).argmax(1)==yte).float().mean()) testloss=float(lossfn(model(Xte),yte)) state=sum(v.numel() for st in opt.state.values() for v in st.values() if torch.is_tensor(v)) return {'final_train_loss':float(np.mean(losses[-30:])), 'test_loss':testloss,'accuracy':acc, 'loss_spikes_gt25pct':spikes, 'optimizer_tensor_state_elems':state} def mlp_experiment(): d=load_digits(); X=StandardScaler().fit_transform(d.data).astype('float32') Xtr,Xte,ytr,yte=train_test_split(X,d.target,test_size=.25,random_state=SEED,stratify=d.target) device='cuda' if torch.cuda.is_available() else 'cpu' tensors=[torch.tensor(Xtr,device=device),torch.tensor(ytr,dtype=torch.long,device=device), torch.tensor(Xte,device=device),torch.tensor(yte,dtype=torch.long,device=device)] try: results={'device':device,'steps':450,'batch':64, 'idea':train('idea',*tensors),'heavy_ball':train('hb',*tensors),'adamw':train('adam',*tensors)} except (RuntimeError, torch.cuda.OutOfMemoryError) as exc: if device != 'cuda': raise device='cpu' tensors=[torch.tensor(Xtr),torch.tensor(ytr,dtype=torch.long),torch.tensor(Xte),torch.tensor(yte,dtype=torch.long)] results={'device':device,'cuda_fallback_error':str(exc),'steps':450,'batch':64, 'idea':train('idea',*tensors),'heavy_ball':train('hb',*tensors),'adamw':train('adam',*tensors)} return results def main(): seed_all(); out={'formula_check':formula_check(),'quadratic_sweep':quadratic_sweep(),'mlp':mlp_experiment()} Path('results.json').write_text(json.dumps(out,indent=2)) print(json.dumps(out,indent=2)) if __name__=='__main__': main()