Phase-Delay Spectral Margin for Attractor RNNs / experiment.py
Mechanism confirmed, baseline not beaten
1import json, math, random
2from pathlib import Path
3import numpy as np
4import torch
5
6SEED = 3035
7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
8try:
9 device = 'cuda' if torch.cuda.is_available() else 'cpu'
10except Exception:
11 device = 'cpu'
12
13
14def laplacian(A, psi, alpha):
15 C = A * np.cos(psi[None, :] - psi[:, None] - alpha)
16 return np.diag(C.sum(1)) - C
17
18
19def core_verification():
20 n = 6; K = 1.0
21 # Symmetric positive graph: exact real spectrum and one gauge eigenvalue.
22 A = np.full((n, n), 0.18); np.fill_diagonal(A, 0.0)
23 psi = np.zeros(n); alpha = np.zeros((n, n))
24 L = laplacian(A, psi, alpha)
25 ev = np.linalg.eigvals(L)
26 nonzero = np.sort(np.real(ev))[1:]
27 lam_max = float(np.max(np.real(ev)))
28 eta_c = 2.0 / (K * lam_max)
29 # Test actual discrete perturbation amplification around the predicted boundary.
30 q = np.random.default_rng(SEED).normal(size=n); q -= q.mean(); q /= np.linalg.norm(q)
31 observed = []
32 for ratio in [0.8, 1.0, 1.2]:
33 eta = ratio * eta_c
34 M = np.eye(n) - eta*K*L
35 # maximum non-gauge one-step amplification is the relevant observed quantity
36 vals = np.linalg.eigvals(M)
37 amp = float(np.max(np.abs(vals[np.abs(vals-1)>1e-7])))
38 observed.append({'ratio_to_eta_c': ratio, 'max_mode_amplification': amp})
39 return {'eigenvalues': ev.tolist(), 'min_nonzero_real_eigenvalue': float(nonzero[0]),
40 'max_eigenvalue': lam_max, 'predicted_eta_c': eta_c, 'euler_tests': observed,
41 'boundary_crossed': observed[0]['max_mode_amplification'] < 1 and observed[2]['max_mode_amplification'] > 1}
42
43
44class PhaseRNN(torch.nn.Module):
45 def __init__(self, n=8, inp=3, dt=0.12):
46 super().__init__(); self.n=n; self.dt=dt
47 self.a = torch.nn.Parameter(torch.full((n,n), -1.2) + .08*torch.randn(n,n))
48 self.raw_alpha = torch.nn.Parameter(.5*torch.randn(n,n))
49 self.W = torch.nn.Parameter(.15*torch.randn(n, inp))
50 self.b = torch.nn.Parameter(torch.zeros(n))
51 def matrices(self):
52 A = torch.nn.functional.softplus(self.a); A = A * (1-torch.eye(self.n, device=A.device))
53 alpha = math.pi * torch.tanh(self.raw_alpha)
54 return A, alpha
55 def step(self, theta, x):
56 A, alpha = self.matrices()
57 # theta: batch,n; x: batch,input
58 diff = theta[:,None,:] - theta[:,:,None] - alpha[None,:,:]
59 coupling = (A[None,:,:] * torch.sin(diff)).sum(2)
60 omega = x @ self.W.T + self.b
61 return theta + self.dt*(omega + coupling)
62 def rollout(self, x):
63 th = torch.zeros(x.shape[0], self.n, device=x.device); out=[]
64 for t in range(x.shape[1]):
65 th = self.step(th, x[:,t]); out.append(th)
66 return torch.stack(out, 1)
67 def spectral_terms(self, psi, K=1.0, eta=None, gamma=.03):
68 A, alpha = self.matrices()
69 C = A * torch.cos(psi[None,:]-psi[:,None]-alpha)
70 L = torch.diag(C.sum(1)) - C
71 ev = torch.linalg.eigvals(L)
72 # Remove gauge mode by excluding eigenvalue closest to zero.
73 idx = torch.argmin(torch.abs(ev)); keep = torch.ones(self.n, dtype=torch.bool, device=ev.device); keep[idx]=False
74 re = ev.real[keep]
75 margin = torch.nn.functional.softplus(torch.as_tensor(gamma, device=psi.device)-re.min())
76 if eta is None: eta=self.dt
77 amp = torch.abs(1-eta*K*ev)
78 # Gauge is exactly one; exclude it from stability penalty.
79 excess = torch.relu(amp[keep]-1).pow(2).mean()
80 return margin + excess, float(re.min().detach().cpu()), float(amp[keep].max().detach().cpu())
81
82
83def train_case(use_spec, x, y, init_state, steps=350):
84 model=PhaseRNN().to(device); model.load_state_dict(init_state)
85 opt=torch.optim.Adam(model.parameters(), lr=.025)
86 history=[]
87 for it in range(steps):
88 pred=model.rollout(x)
89 task=torch.mean((torch.sin(pred)-y)**2)
90 psi=pred.detach().mean((0,1)) if not use_spec else pred.mean((0,1))
91 spec, margin, amp=model.spectral_terms(psi)
92 loss=task + (.08*spec if use_spec else 0.0)
93 opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
94 if it in (0, steps-1): history.append((float(task.detach().cpu()), margin, amp))
95 with torch.no_grad():
96 pred=model.rollout(x); task=float(torch.mean((torch.sin(pred)-y)**2).cpu())
97 psi=pred.mean((0,1)); spec, margin, amp=model.spectral_terms(psi)
98 # finite perturbation test, quotienting out common phase by zero-mean perturbation
99 th=pred[-1,-1].clone(); d=torch.randn_like(th); d-=d.mean(); d=d/d.norm()*1e-3
100 th2=th+d
101 for _ in range(12): th= model.step(th[None,:], torch.zeros(1,x.shape[2],device=device))[0]; th2=model.step(th2[None,:], torch.zeros(1,x.shape[2],device=device))[0]
102 decay=float((th2-th).norm().cpu()/1e-3)
103 return {'task_mse':task, 'spectral_loss':float(spec.cpu()), 'min_real_eigenvalue':float(margin), 'max_euler_amplification':float(amp), '12_step_perturbation_gain':decay, 'history':history}
104
105
106def main():
107 ver=core_verification()
108 torch.manual_seed(SEED)
109 T=24; B=64; inp=3; n=8
110 x=torch.randn(B,T,inp,device=device)
111 # Stable teacher target: smooth bounded phase observations.
112 teacher=torch.zeros(B,n,device=device); ys=[]
113 Wt=torch.randn(n,inp,device=device)*.18
114 for t in range(T):
115 teacher=teacher + .12*(x[:,t]@Wt.T - .35*torch.sin(teacher))
116 ys.append(torch.sin(teacher))
117 y=torch.stack(ys,1)
118 torch.manual_seed(SEED+1); init=PhaseRNN(n,inp).state_dict()
119 # Recreate same parameter initialization for both cases.
120 baseline=train_case(False,x,y,init)
121 idea=train_case(True,x,y,init)
122 result={'device':device,'verification':ver,'baseline':baseline,'idea':idea,
123 'claim_signal': idea['12_step_perturbation_gain'] < baseline['12_step_perturbation_gain'] and idea['max_euler_amplification'] <= baseline['max_euler_amplification']}
124 Path('results.json').write_text(json.dumps(result, indent=2))
125 print(json.dumps(result, indent=2))
126
127if __name__=='__main__': main()