Positive-real rational resolvent mixer / resolvent_mixer_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  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()