Robust Oscillatory RNN via Cyclic Loop-Gain Certification / cyclic_rnn_experiment.py
Mechanism confirmed, baseline not beaten
1import json, math, random
2from pathlib import Path
3import numpy as np
4from scipy.optimize import differential_evolution, minimize_scalar
5
6SEED = 7
7np.random.seed(SEED); random.seed(SEED)
8
9
10def cycle_matrix(weights):
11 n = len(weights); A = np.zeros((n,n), dtype=complex)
12 # y_i receives from predecessor i-1; product of the directed cycle is alpha
13 for i, w in enumerate(weights): A[i, (i-1) % n] = w
14 return A
15
16
17def determinant_identity_check(n=5):
18 rng = np.random.default_rng(SEED)
19 weights = rng.uniform(.3, 1.4, n)
20 alpha = np.prod(weights)
21 errs=[]
22 for _ in range(30):
23 z = rng.normal()+1j*rng.normal()
24 h = rng.normal()+1j*rng.normal()
25 hs = h*(1 + .3*(rng.normal(size=n)+1j*rng.normal(size=n)))
26 D=np.diag(hs); A=cycle_matrix(weights)
27 lhs=np.linalg.det(np.eye(n)-D@A)
28 rhs=1-alpha*np.prod(hs)
29 errs.append(abs(lhs-rhs))
30 return float(max(errs)), float(np.mean(errs))
31
32# Stable unit h(z)=b/(z-a), with positive feedback and cycle alpha.
33def q_abs(omega, phases, n, alpha=.5, a=.8, b=1., r=0.):
34 z=np.exp(1j*omega); h=b/(z-a)
35 return abs(1-alpha*(h**n)*np.prod(1+r*np.exp(1j*np.asarray(phases))))
36
37def critical_radius(n, alpha=.5, a=.8, b=1.):
38 # Minimize the scalar characteristic residual over unit-circle frequency and
39 # independent uncertainty phases. This is a finite-grid approximation to rc.
40 def obj(x): return q_abs(x[0], x[1:], n, alpha, a, b, x[-1] if False else 0.)
41 # Instead search r jointly and find where min residual crosses 0; phases have
42 # an exact useful reduction: product factors are optimized by DE at each r.
43 def minres(r):
44 def f(x): return q_abs(x[0], x[1:], n, alpha, a, b, r)
45 bounds=[(0,2*np.pi)]+[(0,2*np.pi)]*n
46 out=differential_evolution(f,bounds,seed=SEED,popsize=8,maxiter=100,tol=1e-8,polish=True)
47 return out.fun
48 # Direct zeros are numerically ill-conditioned, so locate the first r for
49 # which a boundary root exists using scalar minimization of residual.
50 rs=np.linspace(0,.99,100)
51 vals=np.array([minres(float(r)) for r in rs])
52 # report the first local near-zero crossing; residual has a numerical floor
53 k=int(np.argmin(vals))
54 # A robust scalar prediction for the destabilizing/stabilizing real direction.
55 # Root z=a+(alpha*b^n)^(1/n)*(1-r), reaching z=1.
56 predicted=1-(1-a)/(alpha**(1/n)*b)
57 return predicted, float(rs[k]), float(vals[k]), [(float(rs[i]),float(vals[i])) for i in range(0,len(rs),10)]
58
59
60def root_radius_for_real_gain(n, r, alpha=.5, a=.8, b=1.):
61 # Exact roots under all factors (1-r): (z-a)^n=alpha*b^n*(1-r)^n.
62 c=(alpha**(1/n))*b*(1-r)
63 roots=np.array([a+c*np.exp(2j*np.pi*k/n) for k in range(n)])
64 return float(np.max(np.abs(roots))), roots
65
66
67def mechanism_sweeps():
68 detmax, detmean=determinant_identity_check()
69 alpha=.5; a=.8; b=1.
70 ns=[2,3,4,6,8]
71 scaling=[]
72 for n in ns:
73 pred=1-(1-a)/(alpha**(1/n)*b)
74 # binary search observed real-gain boundary max pole radius <= 1
75 lo,hi=0.,.999
76 for _ in range(60):
77 mid=(lo+hi)/2
78 if root_radius_for_real_gain(n,mid,alpha,a,b)[0] > 1: lo=mid
79 else: hi=mid
80 scaling.append({'n':n,'predicted_rc':pred,'observed_rc':hi,'abs_error':abs(pred-hi)})
81 radii=[]
82 n=4; pred=scaling[2]['predicted_rc']
83 for r in [0., .4, .7, pred-.02, pred+.02, .9]:
84 rad,_=root_radius_for_real_gain(n,r,alpha,a,b)
85 radii.append({'r':r,'max_pole_radius':rad,'unstable':rad>1+1e-10})
86 return {'determinant_max_abs_error':detmax,'determinant_mean_abs_error':detmean,
87 'critical_radius_scaling':scaling,'radius_transition':radii}
88
89
90def tiny_rnn():
91 # Small deterministic comparison; cyclic B is the proposed topology, dense B baseline.
92 try:
93 import torch
94 torch.manual_seed(SEED); np.random.seed(SEED)
95 device='cuda' if torch.cuda.is_available() else 'cpu'
96 class RNN(torch.nn.Module):
97 def __init__(self,n,cyclic):
98 super().__init__(); self.n=n; self.cyclic=cyclic
99 self.B=torch.nn.Parameter(torch.zeros(n,n)); self.C=torch.nn.Parameter(torch.randn(n,1)*.15); self.O=torch.nn.Parameter(torch.randn(1,n)*.15)
100 mask=torch.zeros(n,n)
101 for i in range(n): mask[i,(i-1)%n]=1
102 self.register_buffer('mask',mask)
103 torch.nn.init.orthogonal_(self.B)
104 def forward(self,u,gain=None):
105 h=torch.zeros(u.shape[0],self.n,device=u.device); ys=[]
106 B=self.B*self.mask if self.cyclic else self.B
107 if gain is not None: B=B*gain
108 for t in range(u.shape[1]):
109 h=torch.tanh(h@B.T+u[:,t,:]@self.C.T); ys.append(h@self.O.T)
110 return torch.stack(ys,1)
111 n=8; T=48; batch=64
112 t=torch.arange(T+1,device=device).float()[None,:]
113 u=torch.sin(.23*t).repeat(batch,1).unsqueeze(-1); target=u[:,1:,:]; inp=u[:,:-1,:]
114 out={}
115 for cyclic in [False,True]:
116 model=RNN(n,cyclic).to(device); opt=torch.optim.Adam(model.parameters(),lr=.02)
117 for _ in range(250):
118 opt.zero_grad(); loss=((model(inp)-target)**2).mean(); loss.backward(); opt.step()
119 with torch.no_grad():
120 clean=float(((model(inp)-target)**2).mean())
121 vals=[]
122 for r in [.1,.3,.5]:
123 g=torch.exp(torch.randn(n,device=device)*r-r*r/2).view(1,n)
124 vals.append(float(((model(inp,gain=g)-target)**2).mean()))
125 out['cyclic' if cyclic else 'dense']={'clean_mse':clean,'gain_noise_mse':vals}
126 return out
127 except Exception as e:
128 return {'error':repr(e)}
129
130if __name__=='__main__':
131 result={'seed':SEED,'math_and_mechanism':mechanism_sweeps(),'tiny_rnn':tiny_rnn()}
132 Path('results.json').write_text(json.dumps(result,indent=2))
133 print(json.dumps(result,indent=2))