Robust Oscillatory RNN via Cyclic Loop-Gain Certification / cyclic_rnn_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  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))