Phase-Delay Spectral Margin for Attractor RNNs / experiment.py

Mechanism confirmed, baseline not beaten

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