Hill-Floquet Regularization for Periodic RNNs / hill_floquet_experiment.py

Mechanism failed

Raw ⬇ ZIP
  1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5from scipy.linalg import eig, expm
  6from scipy.integrate import solve_ivp
  7
  8# Reproducible toy study of the Fourier Hill construction.
  9np.random.seed(7)
 10T = 2.0 * np.pi
 11omega = 2.0 * np.pi / T
 12
 13
 14def rot(theta):
 15    c, s = np.cos(theta), np.sin(theta)
 16    return np.array([[c, -s], [s, c]])
 17
 18
 19def J_of_t(t, gamma=0.0, amp=1.0, rate=1.0):
 20    # A genuinely noncommuting periodic linearization.  The constant gamma
 21    # shifts every Floquet exponent exactly, which gives a sharp prediction.
 22    R = rot(rate * t)
 23    D = np.diag([gamma + 0.18, gamma - 0.42])
 24    return R @ D @ R.T + amp * np.array([[0.0, np.sin(t)], [0.25*np.cos(t), 0.0]])
 25
 26
 27def monodromy_exponents(gamma=0.0, amp=1.0, rate=1.0, rtol=2e-10, atol=2e-12):
 28    d = 2
 29    def rhs(t, y):
 30        Y = y.reshape(d, d)
 31        return (J_of_t(t, gamma, amp, rate) @ Y).reshape(-1)
 32    sol = solve_ivp(rhs, (0.0, T), np.eye(d).reshape(-1), rtol=rtol, atol=atol,
 33                    method="DOP853")
 34    vals = np.linalg.eigvals(sol.y[:, -1].reshape(d, d))
 35    # Real parts are unambiguous for the dominant multiplier in this toy.
 36    mu = np.log(np.abs(vals)) / T
 37    return np.sort(mu)[::-1], vals
 38
 39
 40def fourier_coeffs(gamma, amp, rate, M):
 41    ts = np.arange(M) * T / M
 42    samples = np.stack([J_of_t(t, gamma, amp, rate) for t in ts], axis=0)
 43    # np.fft convention: samples[m] = sum_k coeff[k] exp(+i 2pi km/M)
 44    coeff = np.fft.fft(samples, axis=0) / M
 45    return coeff
 46
 47
 48def hill_exponents(gamma=0.0, amp=1.0, rate=1.0, N=8, M=None, filtered=False):
 49    if M is None: M = max(4*N + 1, 64)
 50    coeff = fourier_coeffs(gamma, amp, rate, M)
 51    modes = np.arange(-N, N+1)
 52    d = 2
 53    H0 = np.zeros((len(modes)*d, len(modes)*d), dtype=complex)
 54    for ii, n in enumerate(modes):
 55        for jj, ell in enumerate(modes):
 56            k = (n - ell) % M
 57            block = coeff[k].copy()
 58            if n == ell:
 59                block = block - 1j * ell * omega * np.eye(d)
 60            H0[ii*d:(ii+1)*d, jj*d:(jj+1)*d] = block
 61    vals, vecs = np.linalg.eig(H0)
 62    if filtered:
 63        # Finite Hill matrices contain edge-localized spectral pollution.
 64        # Physical Fourier modes have small mass in the outermost shell.
 65        q = d
 66        edge = np.r_[np.arange(q), np.arange(len(modes)*q-q, len(modes)*q)]
 67        mass = np.sum(np.abs(vecs[edge, :])**2, axis=0) / np.sum(np.abs(vecs)**2, axis=0)
 68        keep = mass < 0.20
 69        if not np.any(keep):
 70            keep = np.ones(len(vals), dtype=bool)
 71        vals = vals[keep]
 72    return np.sort(vals.real)[::-1], vals
 73
 74
 75def penalty(gamma, beta=2.0, N=8):
 76    dominant = hill_exponents(gamma=gamma, N=N, filtered=True)[0][0]
 77    # Smooth version of max(0, mu+epsilon)^2; here epsilon=0.
 78    return beta * max(0.0, dominant)**2, dominant
 79
 80
 81def regularized_parameter_fit(beta=2.0, steps=120, lr=0.08, N=8):
 82    # Task loss prefers gamma=+0.20 (unstable), Hill penalty should move it
 83    # to the stability boundary or below it.
 84    gamma = 0.20
 85    for _ in range(steps):
 86        mu = hill_exponents(gamma=gamma, N=N, filtered=True)[0][0]
 87        # Numerically, this construction has d(mu)/d(gamma) approximately 1.
 88        dpen = 2.0 * beta * max(0.0, mu)
 89        grad = 2.0 * (gamma - 0.20) + dpen
 90        gamma -= lr * grad
 91    return gamma, hill_exponents(gamma=gamma, N=N, filtered=True)[0][0]
 92
 93
 94def main():
 95    report = {"seed": 7, "T": T, "omega": omega}
 96
 97    # Prediction 1: adding gamma I shifts every exponent by gamma, slope 1.
 98    gs = np.linspace(-0.45, 0.35, 9)
 99    exact = np.array([monodromy_exponents(gamma=g)[0][0] for g in gs])
100    hill = np.array([hill_exponents(gamma=g, N=8, filtered=True)[0][0] for g in gs])
101    slope_exact, intercept_exact = np.polyfit(gs, exact, 1)
102    slope_hill, intercept_hill = np.polyfit(gs, hill, 1)
103    report["prediction_1_offset_scaling"] = {
104        "predicted": "dominant Floquet exponent shifts as mu(gamma)=mu(0)+gamma; slope=1",
105        "spectrum_filter": "discard eigenvectors with >20% mass in outer Fourier shell",
106        "observed_exact_slope": float(slope_exact),
107        "observed_hill_slope": float(slope_hill),
108        "max_exact_shift_error": float(np.max(np.abs((exact-exact[4])-(gs-gs[4])))),
109        "max_hill_vs_exact": float(np.max(np.abs(hill-exact))),
110    }
111
112    # Prediction 2: stability boundary occurs where dominant exponent crosses 0.
113    # Estimate by linear interpolation of the sweep and compare to exact root.
114    def crossing(x, y):
115        for i in range(len(x)-1):
116            if y[i] * y[i+1] <= 0:
117                return float(x[i] - y[i]*(x[i+1]-x[i])/(y[i+1]-y[i]))
118        return float("nan")
119    report["prediction_2_boundary"] = {
120        "predicted": "rho_F=1 exactly when max Re(mu)=0",
121        "exact_crossing_gamma": crossing(gs, exact),
122        "hill_N8_crossing_gamma": crossing(gs, hill),
123        "exact_rho_at_crossing": float(np.exp(T * 0.0)),
124    }
125
126    # Prediction 3: increasing Fourier truncation should converge to monodromy.
127    exact0 = monodromy_exponents()[0][0]
128    trunc = {}
129    for N in [1, 2, 4, 8, 12]:
130        h = hill_exponents(N=N, filtered=True)[0][0]
131        trunc[str(N)] = {"hill_mu": float(h), "abs_error": float(abs(h-exact0))}
132    report["prediction_3_truncation"] = {
133        "predicted": "smooth periodic coefficients give improving Hill approximation; N=8 close to N=4",
134        "spectrum_filter": "discard edge-localized finite-section roots",
135        "monodromy_mu": float(exact0), "values": trunc,
136        "N4_to_N8_relative_change": float(abs(trunc["8"]["hill_mu"]-trunc["4"]["hill_mu"]) /
137                                            max(abs(trunc["8"]["hill_mu"]), 1e-12)),
138    }
139
140    # Secondary mini experiment: task-only fit remains at the unstable target;
141    # regularized fit should sacrifice task objective to reduce growth.
142    base_gamma = 0.20
143    reg_gamma, reg_mu = regularized_parameter_fit()
144    report["mini_optimization"] = {
145        "task_target_gamma": base_gamma,
146        "baseline_gamma": base_gamma,
147        "baseline_mu": float(hill_exponents(gamma=base_gamma, filtered=True)[0][0]),
148        "regularized_gamma": float(reg_gamma),
149        "regularized_mu": float(reg_mu),
150        "baseline_rho": float(np.exp(T * hill_exponents(gamma=base_gamma, filtered=True)[0][0])),
151        "regularized_rho": float(np.exp(T * reg_mu)),
152    }
153
154    Path("results.json").write_text(json.dumps(report, indent=2))
155    print(json.dumps(report, indent=2))
156
157if __name__ == "__main__":
158    main()