Gramian-balanced neural SSM compression / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json
  2import numpy as np
  3from scipy.linalg import solve_discrete_lyapunov, svd
  4
  5SEED = 7
  6
  7
  8def compress(Y, tau=1e-12):
  9    if Y.shape[1] == 0:
 10        return np.zeros((Y.shape[0], 0))
 11    U, s, _ = svd(Y, full_matrices=False)
 12    keep = s > (tau * s[0] if s[0] else 0.0)
 13    return U[:, keep] * s[keep]
 14
 15
 16def gramian_factors(F, B, C, K=100, tau=1e-12):
 17    Zp = np.zeros((F.shape[0], 0))
 18    Zq = np.zeros_like(Zp)
 19    for _ in range(K):
 20        Zp = compress(np.concatenate((F @ Zp, B), axis=1), tau)
 21        Zq = compress(np.concatenate((F.T @ Zq, C.T), axis=1), tau)
 22    return Zp, Zq
 23
 24
 25def balance(F, B, C, r, K=100):
 26    Zp, Zq = gramian_factors(F, B, C, K)
 27    U, s_all, Vt = svd(Zq.T @ Zp, full_matrices=False)
 28    rr = min(r, len(s_all))
 29    keep = s_all[:rr] > max(s_all[0] * 1e-13, 1e-15)
 30    rr = int(np.sum(keep))
 31    U, V, s = U[:, :rr], Vt.T[:, :rr], s_all[:rr]
 32    invsqrt = np.diag(1.0 / np.sqrt(s))
 33    Vr = Zp @ V @ invsqrt
 34    Wr = Zq @ U @ invsqrt
 35    return Wr.T @ F @ Vr, Wr.T @ B, C @ Vr, s_all, Vr, Wr
 36
 37
 38def response(F, B, C, ws):
 39    I = np.eye(F.shape[0])
 40    vals = []
 41    for w in ws:
 42        vals.append(C @ np.linalg.solve(np.exp(1j*w) * I - F, B))
 43    return np.asarray(vals)
 44
 45
 46def relerr(A, B):
 47    return float(np.linalg.norm(A - B) / max(np.linalg.norm(A), 1e-15))
 48
 49
 50def main():
 51    rng = np.random.default_rng(SEED)
 52    results = {}
 53
 54    # Diagonal SSM: P_ii=1/(1-a_i^2), and finite-horizon error is a_i^(2K).
 55    a = np.array([.20, .55, .80, .90, .96, .985, .995])
 56    F = np.diag(a)
 57    B = np.ones((len(a), 1))
 58    C = np.ones((1, len(a)))
 59    P = solve_discrete_lyapunov(F, B @ B.T)
 60    Q = solve_discrete_lyapunov(F.T, C.T @ C)
 61    stein_p = relerr(P, F @ P @ F.T + B @ B.T)
 62    stein_q = relerr(Q, F.T @ Q @ F + C.T @ C)
 63    results['stein_relative_residuals'] = {'P': stein_p, 'Q': stein_q}
 64
 65    # Prediction 1: finite-horizon Gramian tail is approximately rho^(2K).
 66    finite = np.zeros_like(P)
 67    rows = []
 68    for K in [1, 2, 4, 8, 16, 32, 64, 128]:
 69        finite = sum(np.linalg.matrix_power(F, i) @ B @ B.T @ np.linalg.matrix_power(F.T, i) for i in range(K))
 70        tail = relerr(P, finite)
 71        # Exact relative spectral-norm tail for this diagonal system.
 72        tail_exact = np.outer(a, a)**K / (1 - np.outer(a, a))
 73        exact_tail = np.linalg.norm(tail_exact) / np.linalg.norm(P)
 74        rows.append({'K': K, 'observed_tail': tail, 'predicted_exact_tail': float(exact_tail),
 75                     'asymptotic_slowest_mode': float(a[-1]**(2*K))})
 76    results['prediction_1_finite_horizon_geometric_tail'] = rows
 77
 78    # Prediction 2: controllability/observability and Hankel singular values diverge as 1/(1-rho^2).
 79    scaling = []
 80    for rho in [.50, .70, .80, .90, .95, .98, .99, .995]:
 81        f = np.array([[rho]])
 82        p = solve_discrete_lyapunov(f, np.ones((1, 1)))[0, 0]
 83        q = solve_discrete_lyapunov(f, np.ones((1, 1)))[0, 0]
 84        scaling.append({'rho': rho, 'P': float(p), 'Q': float(q), 'predicted_1_over_1_minus_rho2': 1.0/(1-rho*rho)})
 85    results['prediction_2_near_instability_scaling'] = scaling
 86
 87    # Prediction 3: balanced truncation error is bounded by 2*sum(discarded HSVs).
 88    n = 12
 89    eigs = np.array([.15, .25, .35, .45, .55, .62, .68, .74, .80, .86, .91, .96])
 90    R = rng.normal(size=(n, n)); U0, _, _ = np.linalg.svd(R)
 91    f = U0 @ np.diag(eigs) @ U0.T
 92    b = rng.normal(size=(n, 2)); c = rng.normal(size=(2, n))
 93    full = response(f, b, c, np.linspace(0, np.pi, 301))
 94    trunc_rows = []
 95    for r in [2, 4, 6, 8, 10]:
 96        fr, br, cr, hsv, Vr, Wr = balance(f, b, c, r, K=160)
 97        red = response(fr, br, cr, np.linspace(0, np.pi, 301))
 98        err = float(np.max([np.linalg.norm(x, 2) for x in (full-red)]))
 99        bound = float(2*np.sum(hsv[r:])) if r < len(hsv) else 0.0
100        # HSV list from the balanced realization is sorted; compare the actual retained rank.
101        trunc_rows.append({'r': r, 'max_frequency_error': err, 'bound_2_sum_discarded_HSV': bound,
102                           'biorthogonality_error': relerr(Wr.T @ Vr, np.eye(Vr.shape[1]))})
103    results['prediction_3_balanced_truncation_bound'] = trunc_rows
104
105    # Compression mini-experiment: Gramian balancing versus magnitude state pruning.
106    r = 4
107    fr, br, cr, hsv, Vr, Wr = balance(f, b, c, r, K=160)
108    gram_err = float(np.max([np.linalg.norm(x, 2) for x in (full-response(fr, br, cr, np.linspace(0, np.pi, 301)))]))
109    # Prune states with smallest row/column energy, then preserve the same coordinates.
110    score = np.sum(b*b, axis=1) + np.sum(c*c, axis=0)
111    keep = np.argsort(score)[-r:]
112    fp = f[np.ix_(keep, keep)]; bp = b[keep]; cp = c[:, keep]
113    prune_err = float(np.max([np.linalg.norm(x, 2) for x in (full-response(fp, bp, cp, np.linspace(0, np.pi, 301)))]))
114    results['compression_comparison'] = {'rank': r, 'gramian_max_frequency_error': gram_err,
115                                         'magnitude_pruning_max_frequency_error': prune_err,
116                                         'full_state_update_cost_units': n*n + n*2 + 2*n,
117                                         'reduced_projection_state_dimension': r,
118                                         'ideal_state_state_transition_cost_ratio': (r*r)/(n*n)}
119
120    # Verify the factor iteration itself against exact Gramians.
121    zp, zq = gramian_factors(f, b, c, K=160)
122    results['factor_iteration_relative_errors'] = {'P': relerr(zp@zp.T, solve_discrete_lyapunov(f,b@b.T)),
123                                                    'Q': relerr(zq@zq.T, solve_discrete_lyapunov(f.T,c.T@c))}
124    with open('results.json', 'w') as out:
125        json.dump(results, out, indent=2)
126    print(json.dumps(results, indent=2))
127
128
129if __name__ == '__main__':
130    main()