Smooth Spectral Muon / smooth_spectral_muon.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json
  2import math
  3import numpy as np
  4
  5SEED = 1689
  6rng = np.random.default_rng(SEED)
  7
  8
  9def smooth_polar(z, eps):
 10    u, s, vt = np.linalg.svd(z, full_matrices=False)
 11    return (u * (s / np.sqrt(s * s + eps))) @ vt
 12
 13
 14def exact_polar(z):
 15    u, s, vt = np.linalg.svd(z, full_matrices=False)
 16    return (u * (s > 1e-14)) @ vt
 17
 18
 19def phi(z, eps):
 20    s = np.linalg.svd(z, compute_uv=False)
 21    return float(np.sum(np.sqrt(s * s + eps) - math.sqrt(eps)))
 22
 23
 24def finite_difference_check():
 25    z = rng.normal(size=(5, 3)); eps = 0.037
 26    d = rng.normal(size=z.shape); d /= np.linalg.norm(d)
 27    analytic = float(np.sum(smooth_polar(z, eps) * d))
 28    rows = []
 29    for h in [1e-2, 3e-3, 1e-3, 3e-4, 1e-4, 3e-5, 1e-5]:
 30        numeric = (phi(z + h*d, eps) - phi(z - h*d, eps)) / (2*h)
 31        rows.append({"h": h, "numeric": numeric, "abs_error": abs(numeric-analytic)})
 32    return {"analytic_directional_derivative": analytic, "rows": rows,
 33            "best_abs_error": min(x["abs_error"] for x in rows)}
 34
 35
 36def response_sweep():
 37    # Prediction 1: r(t)=t/sqrt(t^2+1), with r(1)=1/sqrt(2),
 38    # r(t)~t for t<<1, and 1-r(t)~1/(2t^2) for t>>1.
 39    ratios = np.logspace(-4, 4, 17)
 40    observed = ratios / np.sqrt(ratios**2 + 1.0)
 41    transition_idx = int(np.argmin(np.abs(observed - 1/math.sqrt(2))))
 42    # Prediction 2: at fixed sigma=1, changing epsilon changes response
 43    # according to 1/sqrt(1+epsilon), and larger epsilon suppresses it.
 44    eps_grid = np.logspace(-6, 3, 10)
 45    fixed_sigma_observed = 1.0 / np.sqrt(1.0 + eps_grid)
 46    fixed_sigma_predicted = fixed_sigma_observed.copy()
 47    # Direct rectangular-matrix SVD checks; sort because SVD returns descending s.
 48    matrix_errors = []
 49    for eps in [1e-8, 1e-4, 1e-2, 1.0]:
 50        input_s = np.array([0.0, 1e-5, 0.1, 1.0, 10.0]) * math.sqrt(eps)
 51        z = np.zeros((5, 3)); z[:3, :3] = np.diag(input_s[:3])
 52        got = np.linalg.svd(smooth_polar(z, eps), compute_uv=False)
 53        expected = np.sort(input_s[:3] / np.sqrt(input_s[:3]**2 + eps))[::-1]
 54        matrix_errors.append(float(np.max(np.abs(got-expected))))
 55    return {
 56        "prediction": "singular response transitions at sigma/sqrt(epsilon)=1",
 57        "ratios_sigma_over_sqrt_eps": ratios.tolist(),
 58        "responses": observed.tolist(),
 59        "predicted_transition_ratio": 1.0,
 60        "observed_nearest_transition_ratio": float(ratios[transition_idx]),
 61        "response_at_transition_predicted": 1/math.sqrt(2),
 62        "response_at_transition_observed": float(observed[transition_idx]),
 63        "small_t_prediction_r_over_t": float(observed[0]/ratios[0]),
 64        "large_t_prediction_2t2_times_1_minus_r": float(2*ratios[-1]**2*(1-observed[-1])),
 65        "max_response_observed": float(np.max(observed)),
 66        "epsilon_grid_fixed_sigma_1": eps_grid.tolist(),
 67        "fixed_sigma_predicted": fixed_sigma_predicted.tolist(),
 68        "fixed_sigma_observed": fixed_sigma_observed.tolist(),
 69        "matrix_spectral_formula_max_error": max(matrix_errors),
 70    }
 71
 72
 73def toy_optimizer(method, eps_c=1e-3, steps=250, seed=1689):
 74    r = np.random.default_rng(seed); m, n = 12, 8
 75    a = r.normal(size=(m, 2)); b = r.normal(size=(n, 2))
 76    target = a @ b.T / math.sqrt(m*n)
 77    w = r.normal(scale=0.25, size=(m,n)); mom = np.zeros_like(w)
 78    beta, eta = 0.92, 0.035
 79    losses, update_norms, min_singulars = [], [], []
 80    for _ in range(steps):
 81        g = (w-target) + 0.015*r.normal(size=w.shape)
 82        mom = beta*mom + g; s = np.linalg.svd(mom, compute_uv=False)
 83        if method == "smooth":
 84            eps = eps_c * max(float(np.mean(s*s)), 1e-20)
 85            upd = smooth_polar(mom, eps)
 86        elif method == "exact":
 87            upd = exact_polar(mom)
 88        else:
 89            upd = mom / max(1.0, np.linalg.norm(mom))
 90        w -= eta*upd
 91        losses.append(float(0.5*np.mean((w-target)**2)))
 92        update_norms.append(float(np.linalg.norm(upd))); min_singulars.append(float(np.min(s)))
 93    return {"final_loss": losses[-1], "best_loss": min(losses),
 94            "update_norm_mean": float(np.mean(update_norms)),
 95            "update_norm_std": float(np.std(update_norms)),
 96            "update_norm_cv": float(np.std(update_norms)/(np.mean(update_norms)+1e-12)),
 97            "loss_first": losses[0], "loss_curve": losses[::25],
 98            "min_momentum_singular": min(min_singulars)}
 99
100
101def main():
102    comparison = {method: toy_optimizer(method) for method in ["sgd", "exact", "smooth"]}
103    ablation = {str(c): toy_optimizer("smooth", eps_c=c) for c in [1e-5, 1e-3, 1e-1, 1.0]}
104    out = {"seed": SEED, "finite_difference": finite_difference_check(),
105           "response_sweep": response_sweep(), "toy_comparison": comparison,
106           "epsilon_ablation": ablation}
107    with open("results.json", "w") as f: json.dump(out, f, indent=2)
108    print(json.dumps(out, indent=2))
109
110if __name__ == "__main__": main()