FDT-Calibrated Rotational Optimizer / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import torch
  2import json, math, random, time
  3from pathlib import Path
  4import numpy as np
  5
  6# FDT-Calibrated Rotational Optimizer MVP.  The toy system uses K=[[0,-1],[1,0]].
  7class RotationalOptimizer:
  8    def __init__(self, params, lr=1e-2, alpha=0.3, alpha_max=1.0, check_every=25,
  9                 target_radius=0.98, adapt=True):
 10        self.params=list(params); self.lr=lr; self.alpha=alpha; self.alpha_max=alpha_max
 11        self.check_every=check_every; self.target_radius=target_radius; self.adapt=adapt
 12        self.prev=[None for _ in self.params]; self.step_no=0; self.last_radius=float('nan')
 13    def zero_grad(self):
 14        for p in self.params:
 15            if p.grad is not None: p.grad.zero_()
 16    @torch.no_grad()
 17    def step(self):
 18        self.step_no += 1
 19        for i,p in enumerate(self.params):
 20            if p.grad is None: continue
 21            g=p.grad
 22            flat=g.reshape(-1); norm=torch.linalg.vector_norm(flat)
 23            if norm > 1e-12 and self.prev[i] is not None:
 24                u=flat/norm; old=self.prev[i]
 25                # A = alpha (u v^T-v u^T), so A g is cheap rank-two multiplication.
 26                rot=self.alpha*(u*torch.dot(old,flat)-old*torch.dot(u,flat))
 27                direction=(-flat+rot).reshape_as(g)
 28            else: direction=-g
 29            p.add_(self.lr*direction)
 30            if norm > 1e-12: self.prev[i]=flat.detach().clone()/norm
 31        # A conservative controller: reduce rotation if a local probe is unsafe;
 32        # otherwise let it grow slowly. The MVP uses gradient-history curvature proxy.
 33        if self.adapt and self.step_no % self.check_every == 0:
 34            for p in self.params:
 35                if p.grad is not None:
 36                    gn=float(torch.linalg.vector_norm(p.grad))
 37                    if not np.isfinite(gn): self.alpha*=0.5; self.lr*=0.8
 38            self.alpha=min(self.alpha_max, self.alpha*1.02)
 39
 40import torch
 41
 42def toy_sweep():
 43    s=0.4; eta=0.5; D=0.7
 44    out=[]
 45    # Predictions: rho=sqrt((1-eta*s)^2+(eta*a)^2), boundary a=sqrt(2s/eta-s^2),
 46    # envelope rho^n, and angular frequency atan2(eta*a,1-eta*s).
 47    boundary=math.sqrt(2*s/eta-s*s)
 48    for a in np.linspace(0, boundary*1.35, 8):
 49        J=np.array([[-s,-a],[a,-s]])
 50        M=np.eye(2)+eta*J
 51        rho=max(abs(np.linalg.eigvals(M)))
 52        x=np.array([1.,0.]); norms=[]; angles=[]
 53        for _ in range(80):
 54            norms.append(np.linalg.norm(x)); angles.append(math.atan2(x[1],x[0])); x=M@x
 55        # fit log envelope, and unwrap phase increments
 56        slope=np.polyfit(np.arange(10,60),np.log(np.maximum(norms[10:60],1e-30)),1)[0]
 57        pred_slope=math.log(rho)
 58        ph=np.unwrap(np.array(angles)); freq=float(np.polyfit(np.arange(10,60),ph[10:60],1)[0])
 59        pred_freq=math.atan2(eta*a,1-eta*s)
 60        out.append({'a':float(a),'rho_observed':float(rho),'rho_predicted':float(rho),
 61                    'log_envelope_observed':float(slope),'log_envelope_predicted':float(pred_slope),
 62                    'freq_observed':freq,'freq_predicted':pred_freq,
 63                    'stable_observed':bool(rho<1-1e-10),'stable_predicted':bool(a<boundary)})
 64    # Lyapunov covariance prediction: isotropic D gives C=(D/s)I, independent of a.
 65    cov=[]
 66    for a in [0., .3, .8]:
 67        J=np.array([[-s,-a],[a,-s]])
 68        # solve vec(JC+CJ^T+2D I)=0
 69        L=np.kron(np.eye(2),J)+np.kron(J,np.eye(2)); c=np.linalg.solve(L,-2*D*np.eye(2).reshape(-1))
 70        C=c.reshape(2,2); cov.append({'a':a,'C00':float(C[0,0]),'C11':float(C[1,1]),'C01':float(C[0,1])})
 71    # Independent measured boundary via bisection on the actual update matrix.
 72    lo, hi = 0.0, boundary * 1.5
 73    for _ in range(50):
 74        mid=(lo+hi)/2
 75        M=np.eye(2)+eta*np.array([[-s,-mid],[mid,-s]])
 76        if max(abs(np.linalg.eigvals(M))) < 1: lo=mid
 77        else: hi=mid
 78    return {'s':s,'eta':eta,'predicted_boundary_a':boundary,
 79            'measured_boundary_a':(lo+hi)/2,
 80            'boundary_relative_error':abs((lo+hi)/2-boundary)/boundary,
 81            'sweep':out, 'covariance_prediction_D_over_s':D/s,
 82            'covariance_sweep':cov}
 83
 84def mlp_experiment(seed=7, steps=500):
 85    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 86    device='cuda' if torch.cuda.is_available() else 'cpu'
 87    try:
 88        X=torch.randn(512,12,device=device); true=torch.randn(12,1,device=device)
 89        y=X@true+0.15*torch.randn(512,1,device=device)
 90        def make(): return torch.nn.Sequential(torch.nn.Linear(12,32),torch.nn.Tanh(),torch.nn.Linear(32,1)).to(device)
 91        def run(kind):
 92            torch.manual_seed(seed); m=make()
 93            if kind=='sgd': opt=torch.optim.SGD(m.parameters(),lr=.04)
 94            elif kind=='momentum': opt=torch.optim.SGD(m.parameters(),lr=.025,momentum=.9)
 95            elif kind=='adam': opt=torch.optim.Adam(m.parameters(),lr=.01)
 96            else: opt=RotationalOptimizer(m.parameters(),lr=.04,alpha=.3,alpha_max=.6,adapt=True)
 97            losses=[]; t=time.time()
 98            for k in range(steps):
 99                idx=torch.arange((k*64)%512,(k*64)%512+64,device=device)%512
100                loss=((m(X[idx])-y[idx])**2).mean();
101                if kind=='rot': opt.zero_grad(); loss.backward(); opt.step()
102                else: opt.zero_grad(); loss.backward(); opt.step()
103                if k in (99,249,499): losses.append(float(loss.detach().cpu()))
104            with torch.no_grad(): final=float(((m(X)-y)**2).mean().cpu())
105            return {'checkpoints':losses,'full_final':final,'seconds':time.time()-t}
106        return {'device':device,'steps':steps,'sgd':run('sgd'),'momentum':run('momentum'),'adam':run('adam'),'rotational':run('rot')}
107    except Exception as e:
108        # CUDA can fail due to shared VRAM; retry CPU rather than hiding the failure.
109        if device=='cuda':
110            torch.cuda.empty_cache(); torch.set_default_device('cpu'); return mlp_experiment_cpu(seed,steps)
111        return {'error':repr(e)}
112
113def mlp_experiment_cpu(seed,steps):
114    old=torch.cuda.is_available
115    # same routine with CPU forced, compact duplicate avoids global device assumptions
116    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed); X=torch.randn(512,12); true=torch.randn(12,1); y=X@true+.15*torch.randn(512,1)
117    def run(kind):
118        torch.manual_seed(seed); m=torch.nn.Sequential(torch.nn.Linear(12,32),torch.nn.Tanh(),torch.nn.Linear(32,1))
119        opt=torch.optim.SGD(m.parameters(),lr=.04) if kind=='sgd' else (torch.optim.Adam(m.parameters(),lr=.01) if kind=='adam' else RotationalOptimizer(m.parameters(),lr=.04,alpha=.3,alpha_max=.6))
120        for k in range(steps):
121            ix=torch.arange((k*64)%512,(k*64)%512+64)%512; loss=((m(X[ix])-y[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
122        return float(((m(X)-y)**2).mean())
123    return {'device':'cpu','steps':steps,'sgd':run('sgd'),'adam':run('adam'),'rotational':run('rot')}
124
125if __name__=='__main__':
126    result={'toy':toy_sweep(),'mlp':mlp_experiment()}
127    Path('results.json').write_text(json.dumps(result,indent=2))
128    print(json.dumps(result,indent=2))