Gramian-balanced neural SSM compression / experiment.py
Mechanism confirmed, baseline not beaten
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()