Universal Clock Regularization for Recurrent Dynamics / clock_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
 1import math, random, json
 2from pathlib import Path
 3import numpy as np
 4
 5SEED=7
 6np.random.seed(SEED); random.seed(SEED)
 7
 8def metrics(phi,dt):
 9    w=np.diff(np.unwrap(phi))/dt
10    return float(w.mean()),float(w.var()),float(np.sqrt(w.var())/(abs(w.mean())+1e-12))
11
12def math_sweeps():
13    omega,dt=1.7,.002
14    t=np.arange(0,30*2*np.pi/omega,dt); th=omega*t
15    scaling=[]
16    for m in [1,2,3,5]:
17      for a in [0,.01,.03,.06]:
18        _,v,e=metrics(th+a*np.sin(m*th),dt)
19        scaling.append({'m':m,'a':a,'observed_var':v,'predicted_var':(omega*a*m)**2/2,'observed_E':e,'predicted_E':abs(a*m)/math.sqrt(2)})
20    noise=[]
21    for sig in [0,.001,.003,.01,.03]:
22      vals=[]
23      for r in range(30):
24        rng=np.random.default_rng(SEED+r)
25        vals.append(metrics(th+rng.normal(0,sig,len(th)),dt)[1])
26      pred=2*sig**2/dt**2
27      noise.append({'sigma':sig,'observed_var':float(np.mean(vals)),'predicted_var':pred})
28    boundary=[]
29    for inc in [0.5,1,2,2.8,3.0,3.13,3.15,3.3,4.0]:
30      ph=np.arange(1000)*inc
31      mw,v,e=metrics(ph,1.0)
32      boundary.append({'increment':inc,'estimated_mean':mw,'error':abs(mw-inc),'E':e})
33    return scaling,noise,boundary
34
35def gru_experiment():
36  try:
37    import torch
38    import torch.nn as nn
39    torch.manual_seed(SEED); device='cuda' if torch.cuda.is_available() else 'cpu'
40    dt=.1; omega=1.0; T=45; ntrain=96; ntest=24
41    def make(n,seed):
42      rng=np.random.default_rng(seed); x=[]
43      for i in range(n):
44        phase=rng.uniform(0,2*np.pi); amp=1+rng.normal(0,.04)
45        tt=np.arange(T+1)*dt
46        z=np.stack([amp*np.sin(omega*tt+phase),amp*np.cos(omega*tt+phase)],1)
47        z+=rng.normal(0,.025,z.shape); x.append(z)
48      return torch.tensor(np.array(x),dtype=torch.float32,device=device)
49    train,test=make(ntrain,10),make(ntest,20)
50    class Model(nn.Module):
51      def __init__(self,clock):
52        super().__init__(); self.clock=clock; self.gru=nn.GRU(2,16,batch_first=True); self.out=nn.Linear(16,2)
53        self.phase=nn.Linear(16,2) if clock else None
54      def forward(self,x):
55        h,_=self.gru(x); y=self.out(h); return y,h
56    def run(clock, dev):
57      torch.manual_seed(SEED); m=Model(clock).to(dev); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
58      tr=train.to(dev); te=test.to(dev)
59      # teacher-forced one-step prediction; phase head is intentionally auxiliary.
60      for ep in range(220):
61        inp,tar=tr[:,:-1],tr[:,1:]; y,h=m(inp); loss=((y-tar)**2).mean()
62        if clock:
63          q=m.phase(h); ph=torch.atan2(q[...,0],q[...,1]); d=torch.atan2(torch.sin(torch.diff(ph,dim=1)),torch.cos(torch.diff(ph,dim=1)))/dt
64          # target natural rate is known in this synthetic benchmark; cycle term prevents zero-clock collapse.
65          loss=loss+.08*d.var(dim=1).mean()+.08*(d.mean(dim=1)-omega).pow(2).mean()
66        opt.zero_grad(); loss.backward(); opt.step()
67      with torch.no_grad():
68        y,h=m(te[:,:-1]); one=float(((y-te[:,1:])**2).mean().cpu())
69        # assess phase consistency of learned phase head; baseline uses geometric atan2.
70        if clock: ph=torch.atan2(m.phase(h)[...,0],m.phase(h)[...,1])
71        else: ph=torch.atan2(h[...,0],h[...,1])
72        ds=torch.atan2(torch.sin(torch.diff(ph,dim=1)),torch.cos(torch.diff(ph,dim=1)))/dt
73        return {'one_step_mse':one,'clock_E':float((ds.std(1)/(ds.mean(1).abs()+1e-6)).mean().cpu()),'device':str(dev)}
74    try:
75      return {'baseline':run(False,device),'clock':run(True,device)}
76    except Exception as cuda_error:
77      # Required robust fallback after any CUDA/cuDNN allocation or execution error.
78      if device != 'cpu':
79        return {'baseline':run(False,'cpu'),'clock':run(True,'cpu'),'cuda_error':str(cuda_error)}
80      raise
81  except Exception as e:
82    return {'error':str(e),'fallback':'math-only'}
83
84def main():
85  s,n,b=math_sweeps(); out={'scaling':s,'noise':n,'unwrap_boundary':b,'gru':gru_experiment()}
86  Path('results.json').write_text(json.dumps(out,indent=2))
87  print(json.dumps(out,indent=2))
88if __name__=='__main__': main()