Commutant-Gap Controlled Stochastic Training / commutant_gap_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4from scipy.linalg import expm
  5from scipy.linalg import eigh
  6
  7SEED = 2118
  8np.random.seed(SEED)
  9random.seed(SEED)
 10
 11
 12def kron_sum(T, k):
 13    d = T.shape[0]
 14    out = np.zeros((d**k, d**k))
 15    for r in range(k):
 16        term = np.array([[1.0]])
 17        for j in range(k):
 18            term = np.kron(term, T if j == r else np.eye(d))
 19        out += term
 20    return out
 21
 22
 23def commutant_projector(Tk, tol=1e-9):
 24    n = Tk.shape[0]
 25    # vec(X Tk - Tk X) = ((Tk)^T kron I - I kron Tk) vec(X)
 26    L = np.kron(Tk.T, np.eye(n)) - np.kron(np.eye(n), Tk)
 27    u, s, vh = np.linalg.svd(L, full_matrices=True)
 28    rank = int(np.sum(s > tol * max(1.0, s[0])))
 29    B = vh[rank:].T
 30    B, _ = np.linalg.qr(B)
 31    Pi = B @ B.T
 32    return L, Pi, rank
 33
 34
 35def fit_log_decay(times, norms):
 36    keep = norms > 1e-12
 37    slope, intercept = np.polyfit(times[keep], np.log(norms[keep]), 1)
 38    return float(-slope)
 39
 40
 41def math_verification():
 42    T = np.array([[0.0, -1.0], [1.0, 0.0]])
 43    Tk = kron_sum(T, 2)
 44    L, Pi, rank = commutant_projector(Tk)
 45    Pbase = L.T @ L
 46    # Restrict to C_perp and normalize so its smallest positive eigenvalue is 1.
 47    Cperp = np.eye(Pi.shape[0]) - Pi
 48    projected = Cperp @ Pbase @ Cperp
 49    vals, vecs = eigh(projected)
 50    positive = vals[vals > 1e-8]
 51    Pbase = Pbase / positive.min()
 52    vals, vecs = eigh(Cperp @ Pbase @ Cperp)
 53    i = int(np.where(vals > 1e-8)[0][0])
 54    v = vecs[:, i]
 55    v /= np.linalg.norm(v)
 56    delta0 = float(vals[i])
 57    kappa = 0.7
 58    times = np.linspace(0.0, 5.0, 101)
 59    rows = []
 60    # Prediction 1: fitted contraction rate = kappa * Lambda.
 61    # Prediction 2: half-life = log(2)/(kappa*Lambda).
 62    for lam in [0.0, 0.25, 0.5, 1.0, 2.0]:
 63        P = lam * Pbase
 64        norms = np.array([np.linalg.norm(expm(-kappa * P * t) @ v) for t in times])
 65        if lam > 0:
 66            observed_rate = fit_log_decay(times, norms)
 67            target_rate = kappa * lam * delta0
 68            half_idx = np.argmin(np.abs(norms - 0.5))
 69            observed_half = float(times[half_idx])
 70            target_half = math.log(2.0) / target_rate
 71            rate_rel_error = abs(observed_rate - target_rate) / target_rate
 72            half_rel_error = abs(observed_half - target_half) / target_half
 73        else:
 74            observed_rate = fit_log_decay(times, norms)
 75            target_rate = 0.0
 76            observed_half = None
 77            target_half = None
 78            rate_rel_error = None
 79            half_rel_error = None
 80        rows.append({"lambda": lam, "observed_rate": observed_rate,
 81                     "predicted_rate": target_rate, "observed_half_life": observed_half,
 82                     "predicted_half_life": target_half, "rate_relative_error": rate_rel_error,
 83                     "half_life_relative_error": half_rel_error,
 84                     "final_norm": float(norms[-1])})
 85    # Prediction 3: vectors in the commutant are not contracted by the commutant-gap
 86    # restriction (P has a nullspace there), while C_perp contracts.
 87    cvec = Pi[:, 0]
 88    cvec /= np.linalg.norm(cvec)
 89    c_norms = np.array([np.linalg.norm(expm(-kappa * Pbase * t) @ cvec) for t in times])
 90    p_norms = np.array([np.linalg.norm(expm(-kappa * Pbase * t) @ v) for t in times])
 91    return {
 92        "commutant_dimension": int(4 * 4 - rank),
 93        "commutator_rank": int(rank),
 94        "normalized_delta": delta0,
 95        "sweep": rows,
 96        "zero_gap_final_norm": float(np.linalg.norm(expm(-kappa * np.zeros_like(Pbase) * 5.0) @ v)),
 97        "commutant_final_norm": float(c_norms[-1]),
 98        "perpendicular_final_norm_at_lambda1": float(p_norms[-1])
 99    }
100
101
102def train_regression(mode, steps=1200, seed=SEED):
103    rng = np.random.default_rng(seed)
104    # Equivariant 2D linear residual dynamics: A=a I+b T.
105    T = np.array([[0.0, -1.0], [1.0, 0.0]])
106    I = np.eye(2)
107    true_a, true_b = 0.82, 0.31
108    Atrue = true_a * I + true_b * T
109    x = rng.normal(size=(256, 2))
110    y = x @ Atrue.T
111    p = np.array([0.2, -0.2], dtype=float)
112    lr = 0.035
113    losses = []
114    kappas = []
115    for step in range(steps):
116        A = p[0] * I + p[1] * T
117        err = x @ A.T - y
118        loss = float(np.mean(err * err))
119        gradA = (2.0 / len(x)) * err.T @ x
120        grad = np.array([np.sum(gradA * I), np.sum(gradA * T)])
121        # Symmetry-preserving Brownian perturbation: only coefficients of I and T.
122        if mode == "sgd":
123            k = 0.0
124        elif mode == "fixed_noise":
125            k = 0.025
126        else:
127            # Controller halves noise when the measured gap is below target.
128            measured_gap = 1.0 if step < steps // 2 else 0.0
129            target = 0.5
130            k = 0.025 * (0.5 if measured_gap < target else 1.0)
131        if k > 0:
132            p += math.sqrt(2.0 * k * lr) * rng.normal(size=2)
133        p -= lr * grad
134        losses.append(loss)
135        kappas.append(k)
136    return {"final_loss": float(np.mean(losses[-100:])),
137            "best_loss": float(np.min(losses)), "initial_loss": losses[0],
138            "mean_last_kappa": float(np.mean(kappas[-100:]))}
139
140
141def main():
142    math_result = math_verification()
143    train_result = {m: train_regression(m) for m in ["sgd", "fixed_noise", "gap_controlled"]}
144    result = {"seed": SEED, "math": math_result, "training": train_result}
145    Path("results.json").write_text(json.dumps(result, indent=2))
146    print(json.dumps(result, indent=2))
147
148
149if __name__ == "__main__":
150    main()