import json, math, random from pathlib import Path import numpy as np import torch import torch.nn as nn SEED = 3155 def seed_all(seed=SEED): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def resolvent_peak_np(W, nfreq=512): n = W.shape[0] grid = np.linspace(0.0, np.pi, nfreq) vals = [] eye = np.eye(n) for th in grid: z = np.exp(1j * th) vals.append(np.linalg.svd(np.linalg.inv(z * eye - W), compute_uv=False)[0]) vals = np.asarray(vals) i = int(np.argmax(vals)) return float(vals[i]), float(grid[i]), vals def toy_verification(): # Both matrices have the same stable eigenvalues, but W_nonnormal has # strong off-diagonal coupling and therefore a much larger resolvent peak. W_normal = np.diag([0.72, 0.55]).astype(float) W_nonnormal = np.array([[0.72, 8.0], [0.0, 0.55]], dtype=float) pn, wn, curve_n = resolvent_peak_np(W_normal) pnn, wnn, curve_nn = resolvent_peak_np(W_nonnormal) # Discrete-time transient response ||W^k||, another direct manifestation # of nonnormal amplification despite spectral radius below one. transient = [] P = np.eye(2) for k in range(31): transient.append(float(np.linalg.svd(P, compute_uv=False)[0])) P = P @ W_nonnormal transient = np.asarray(transient) return { "normal_eigenvalues": np.linalg.eigvals(W_normal).tolist(), "nonnormal_eigenvalues": np.linalg.eigvals(W_nonnormal).tolist(), "normal_rho": float(max(abs(np.linalg.eigvals(W_normal)))), "nonnormal_rho": float(max(abs(np.linalg.eigvals(W_nonnormal)))), "normal_resolvent_peak": pn, "nonnormal_resolvent_peak": pnn, "normal_peak_theta": wn, "nonnormal_peak_theta": wnn, "resolvent_amplification_ratio": pnn / pn, "max_transient_gain": float(transient.max()), "transient_peak_step": int(transient.argmax()), "transient_curve": transient.tolist(), } class TinyRNN(nn.Module): def __init__(self, hidden=8): super().__init__() self.hidden = hidden self.inp = nn.Linear(1, hidden) self.rec = nn.Linear(hidden, hidden, bias=False) self.out = nn.Linear(hidden, 1) nn.init.orthogonal_(self.rec.weight, gain=0.85) def forward(self, x, return_states=False): b, t, _ = x.shape h = torch.zeros(b, self.hidden, device=x.device) states = [] for k in range(t): h = torch.tanh(self.inp(x[:, k]) + self.rec(h)) states.append(h) y = self.out(h) return (y, states) if return_states else y def frequency_penalty(model, theta_grid, tau=0.15): # Exact Jacobian for the linearized recurrent map at h=0 and input u=0. # For tanh, derivative is identity there: J=rec.weight, B=inp.weight[:,0], # C=out.weight, D=0. We use the discrete transfer formula. J = model.rec.weight B = model.inp.weight[:, :1] C = model.out.weight I = torch.eye(J.shape[0], device=J.device, dtype=J.dtype) vals = [] for theta in theta_grid: z = torch.complex(torch.cos(theta), torch.sin(theta)) # Solve (zI-J)v=B in complex arithmetic. M = z * I.to(torch.complex64) - J.to(torch.complex64) v = torch.linalg.solve(M, B.to(torch.complex64)) g = C.to(torch.complex64) @ v vals.append(torch.abs(g).reshape(())) vals = torch.stack(vals) return tau * torch.logsumexp(vals / tau, dim=0), vals.max().detach(), vals def make_batch(batch, length, device): x = torch.randn(batch, length, 1, device=device) # Sequence-to-one delayed sum task; it encourages memory without requiring # a large model or dataset. y = x.sum(dim=1) return x, y def train_one(kind, device, steps=500): seed_all(SEED + (0 if kind == "baseline" else 1)) model = TinyRNN().to(device) opt = torch.optim.Adam(model.parameters(), lr=3e-3) theta = torch.linspace(0, math.pi, 24, device=device) losses, penalties, peaks = [], [], [] for step in range(steps): x, y = make_batch(64, 12, device) pred = model(x) task = ((pred - y) ** 2).mean() if kind == "frequency": reg, peak, _ = frequency_penalty(model, theta) loss = task + 0.002 * reg penalties.append(float(reg.detach().cpu())) else: peak = torch.tensor(float("nan"), device=device) loss = task opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) opt.step() losses.append(float(task.detach().cpu())); peaks.append(float(peak.detach().cpu())) with torch.no_grad(): _, _, gridvals = frequency_penalty(model, theta) final_peak = float(gridvals.max().cpu()) final_task = float(np.mean(losses[-50:])) W = model.rec.weight.detach().cpu().numpy() rho = float(max(abs(np.linalg.eigvals(W)))) return {"final_task_loss": final_task, "final_resolvent_peak": final_peak, "spectral_radius": rho, "loss_start": float(np.mean(losses[:20])), "loss_curve": losses, "penalty_curve": penalties} def main(): seed_all() # CUDA is allowed but failure must gracefully fall back to CPU. device = "cuda" if torch.cuda.is_available() else "cpu" try: if device == "cuda": torch.zeros(1, device=device) except Exception: device = "cpu" result = {"seed": SEED, "device": device, "toy_verification": toy_verification()} result["baseline"] = train_one("baseline", device) result["frequency"] = train_one("frequency", device) result["observed_peak_reduction"] = (result["baseline"]["final_resolvent_peak"] - result["frequency"]["final_resolvent_peak"]) / result["baseline"]["final_resolvent_peak"] result["observed_task_change"] = result["frequency"]["final_task_loss"] - result["baseline"]["final_task_loss"] Path("results.json").write_text(json.dumps(result, indent=2)) print(json.dumps({k: v for k, v in result.items() if k not in ("baseline", "frequency")}, indent=2)) print(json.dumps({"baseline": {k:v for k,v in result["baseline"].items() if not k.endswith("curve")}, "frequency": {k:v for k,v in result["frequency"].items() if not k.endswith("curve")}}, indent=2)) if __name__ == "__main__": main()