Positive-real rational resolvent mixer / resolvent_mixer_experiment.py
Failed on benchmark
1import json
2import numpy as np
3
4
5def make_certified(r=8, d=6, seed=0):
6 rng = np.random.default_rng(seed)
7 # C=B^T makes the KYP off-diagonal blocks vanish. Positive A,D
8 # then give a strict certificate while retaining a nontrivial rational term.
9 A = 0.35 * np.eye(r)
10 B = rng.normal(size=(r, d)) / np.sqrt(d)
11 C = B.T
12 D = 0.12 * np.eye(d)
13 P1 = np.diag([1.0] * (r // 2) + [0.0] * (r - r // 2))
14 P2 = np.eye(r) - P1
15 return A, B, C, D, P1, P2
16
17
18def kyp_matrix(A, B, C, D):
19 return np.block([[A + A.T, C.T - B], [C - B.T, D + D.T]])
20
21
22def H_of_z(z1, z2, A, B, C, D, P1, P2):
23 M = A + z1 * P1 + z2 * P2
24 return D + C @ np.linalg.solve(M, B)
25
26
27def resolvent(x, eta, H):
28 return np.linalg.solve(np.eye(H.shape[0]) + eta * H, x)
29
30
31def verify(seed=0):
32 A, B, C, D, P1, P2 = make_certified(seed=seed)
33 eig_kyp = np.linalg.eigvalsh(kyp_matrix(A, B, C, D))
34 rng = np.random.default_rng(seed + 11)
35 min_herm = np.inf
36 max_sigma = 0.0
37 max_energy_ratio = 0.0
38 samples = []
39 for z1, z2 in np.exp(rng.uniform(np.log(1e-3), np.log(20.0), size=(500, 2))):
40 H = H_of_z(z1, z2, A, B, C, D, P1, P2)
41 herm_eig = np.linalg.eigvalsh((H + H.T) / 2).min()
42 R = np.linalg.inv(np.eye(H.shape[0]) + 0.8 * H)
43 min_herm = min(min_herm, herm_eig)
44 max_sigma = max(max_sigma, np.linalg.svd(R, compute_uv=False)[0])
45 x = rng.normal(size=H.shape[0])
46 max_energy_ratio = max(max_energy_ratio, np.linalg.norm(R @ x) / np.linalg.norm(x))
47 # The control uses the same positive H but an explicit additive residual.
48 H0 = H_of_z(0.5, 0.5, A, B, C, D, P1, P2)
49 x = rng.normal(size=H0.shape[0])
50 implicit_norms, additive_norms = [np.linalg.norm(x)], [np.linalg.norm(x)]
51 for _ in range(24):
52 x = resolvent(x, 0.8, H0)
53 implicit_norms.append(np.linalg.norm(x))
54 x = rng.normal(size=H0.shape[0])
55 for _ in range(24):
56 x = x + 0.8 * (H0 @ x)
57 additive_norms.append(np.linalg.norm(x))
58 return {
59 "kyp_min_eigenvalue": float(eig_kyp.min()),
60 "min_hermitian_part_eigenvalue": float(min_herm),
61 "max_resolvent_singular_value": float(max_sigma),
62 "max_sampled_energy_ratio": float(max_energy_ratio),
63 "implicit_norm_start_end": [float(implicit_norms[0]), float(implicit_norms[-1])],
64 "additive_norm_start_end": [float(additive_norms[0]), float(additive_norms[-1])],
65 "implicit_norm_max": float(max(implicit_norms)),
66 "additive_norm_max": float(max(additive_norms)),
67 }
68
69
70def toy_training(seed=123, n=128, steps=40, depth=8, eta=0.8):
71 # Same fixed operator and same initial batch. Target is the stable zero state;
72 # this isolates whether the update itself suppresses activation growth.
73 rng = np.random.default_rng(seed)
74 A, B, C, D, P1, P2 = make_certified(seed=seed)
75 H = H_of_z(0.5, 0.5, A, B, C, D, P1, P2)
76 x0 = rng.normal(size=(n, H.shape[0]))
77 baseline = x0.copy()
78 idea = x0.copy()
79 rows = []
80 for step in range(steps + 1):
81 rows.append({
82 "step": step,
83 "baseline_loss": float(np.mean(baseline ** 2)),
84 "idea_loss": float(np.mean(idea ** 2)),
85 "baseline_activation_rms": float(np.sqrt(np.mean(baseline ** 2))),
86 "idea_activation_rms": float(np.sqrt(np.mean(idea ** 2))),
87 "baseline_max_norm": float(np.max(np.linalg.norm(baseline, axis=1))),
88 "idea_max_norm": float(np.max(np.linalg.norm(idea, axis=1))),
89 })
90 if step == steps:
91 break
92 # Standard explicit residual versus the proposed implicit resolvent.
93 baseline = baseline + eta * (baseline @ H.T)
94 idea = np.stack([resolvent(v, eta, H) for v in idea])
95 # Explicit residual controls at smaller learning/update coefficients.
96 sweep = {}
97 for explicit_eta in (0.05, 0.1, 0.2, 0.4, 0.8):
98 z = x0.copy()
99 for _ in range(steps):
100 z = z + explicit_eta * (z @ H.T)
101 sweep[str(explicit_eta)] = {
102 "final_rms": float(np.sqrt(np.mean(z ** 2))),
103 "max_rms_during_run": float(np.sqrt(np.mean(x0 ** 2))) if explicit_eta == 0 else float(np.sqrt(np.mean(z ** 2))),
104 }
105 return {"operator_eigenvalues": np.linalg.eigvalsh(H).tolist(), "trace": rows, "explicit_step_sweep_final_rms": sweep}
106
107
108def main():
109 out = {"verification": verify(), "toy_training": toy_training()}
110 with open("results.json", "w") as f:
111 json.dump(out, f, indent=2)
112 print(json.dumps(out, indent=2))
113
114
115if __name__ == "__main__":
116 main()