1import json
  2import math
  3import random
  4from pathlib import Path
  5import numpy as np
  6
  7SEED = 2027
  8rng = np.random.default_rng(SEED)
  9EPS = 1e-8
 10
 11
 12def throttle(residual, q, kappa, eps=EPS):
 13    delta = np.linalg.norm(residual) / (math.sqrt(max(q, 0.0)) + eps)
 14    a = min(1.0, kappa / (delta + eps))
 15    return float(a), float(delta)
 16
 17
 18def math_checks():
 19    # Exact q for J=X. Prediction 1: a transitions at ||r|| = kappa sqrt(q).
 20    X = np.array([[1.0, 0.0], [0.0, 2.0], [1.0, -1.0]])
 21    q = float(np.sum(X * X))
 22    kappa = 0.35
 23    threshold = kappa * math.sqrt(q)
 24    residual_norms = np.array([0.25 * threshold, 0.99 * threshold, 1.01 * threshold, 2.0 * threshold, 5.0 * threshold])
 25    gains = np.array([throttle(np.array([v, 0.0, 0.0]), q, kappa)[0] for v in residual_norms])
 26    predicted = np.minimum(1.0, threshold / (residual_norms + EPS))
 27    transition_error = float(np.max(np.abs(gains - predicted)))
 28    transition_observed = float(residual_norms[np.argmax(gains < 0.999999)])
 29
 30    # Prediction 2: in the throttled regime, a(c r) * c is constant (inverse gain scaling).
 31    base = np.array([3.0, -2.0, 1.0])
 32    scales = np.array([1.0, 2.0, 4.0, 8.0])
 33    gains_scale = np.array([throttle(c * base, q, kappa)[0] for c in scales])
 34    inverse_products = scales * gains_scale
 35    expected_products = np.full_like(inverse_products, kappa * math.sqrt(q) / np.linalg.norm(base))
 36    # only assess the scales predicted to be throttled
 37    scale_error = float(np.max(np.abs(inverse_products - expected_products)))
 38
 39    # Prediction 3: scalar theta'= (1-eta*a(theta))*theta. Plain GD diverges for eta*lambda>2;
 40    # throttle makes the large-state update approximately eta*kappa and remains bounded.
 41    eta = 3.0
 42    kappa1 = 0.4
 43    n = 80
 44    theta_plain, theta_throttle = 1.0, 1.0
 45    plain_path, throttled_path, scalar_gains = [], [], []
 46    for _ in range(n):
 47        plain_path.append(abs(theta_plain))
 48        throttled_path.append(abs(theta_throttle))
 49        a, _ = throttle(np.array([theta_throttle]), 1.0, kappa1)
 50        scalar_gains.append(a)
 51        theta_plain = theta_plain - eta * theta_plain
 52        theta_throttle = theta_throttle - eta * a * theta_throttle
 53    plain_growth = abs(theta_plain)
 54    throttled_final = abs(theta_throttle)
 55    # In this setup the predicted GD multiplier is |1-eta|=2 and throttle should not grow exponentially.
 56    expected_plain_growth = 2.0 ** n
 57    return {
 58        "q": q,
 59        "kappa": kappa,
 60        "transition_predicted_residual_norm": threshold,
 61        "transition_observed_first_throttled_norm": transition_observed,
 62        "transition_max_formula_error": transition_error,
 63        "residual_scales": scales.tolist(),
 64        "gains_at_scales": gains_scale.tolist(),
 65        "scale_times_gain": inverse_products.tolist(),
 66        "scale_invariance_max_error": scale_error,
 67        "scalar_predicted_plain_multiplier": 2.0,
 68        "scalar_plain_final_abs": plain_growth,
 69        "scalar_plain_expected_abs": expected_plain_growth,
 70        "scalar_throttled_final_abs": throttled_final,
 71        "scalar_throttled_max_abs": float(max(throttled_path)),
 72        "scalar_min_gain": float(min(scalar_gains)),
 73    }
 74
 75
 76def make_data(n=512, d=8):
 77    rg = np.random.default_rng(SEED + 1)
 78    X = rg.normal(size=(n, d))
 79    true = rg.normal(size=d)
 80    y = X @ true + 0.03 * rg.normal(size=n)
 81    return X, y, true
 82
 83
 84def run_method(method, X, y, true, lr, steps=240, batch=64, kappa=0.5):
 85    rg = np.random.default_rng(SEED + 10)
 86    theta = np.zeros(X.shape[1])
 87    m = np.zeros_like(theta); v = np.zeros_like(theta)
 88    errors, losses, gains, deltas = [], [], [], []
 89    for t in range(steps):
 90        # Alternating masks hide one of two feature groups for 20 steps.
 91        mask = np.ones(X.shape[1])
 92        group = (t // 20) % 2
 93        mask[group::2] = 0.0
 94        idx = rg.choice(len(X), size=batch, replace=False)
 95        Xm = X[idx] * mask
 96        residual = Xm @ theta - y[idx]
 97        g = Xm.T @ residual / batch
 98        q = float(np.sum(Xm * Xm) / batch)
 99        if method == "throttle":
100            a, delta = throttle(residual, q, kappa)
101            g = a * g
102        else:
103            a, delta = 1.0, np.linalg.norm(residual) / (math.sqrt(q) + EPS)
104        if method == "adam":
105            b1, b2 = 0.9, 0.999
106            m = b1*m + (1-b1)*g; v = b2*v + (1-b2)*g*g
107            step = lr * (m/(1-b1**(t+1))) / (np.sqrt(v/(1-b2**(t+1))) + 1e-8)
108            theta -= step
109        else:
110            theta -= lr * g
111        full_res = X @ theta - y
112        errors.append(float(np.linalg.norm(theta - true)))
113        losses.append(float(np.mean(full_res**2)))
114        gains.append(a); deltas.append(delta)
115    return {"final_error": errors[-1], "min_error": min(errors), "max_error": max(errors),
116            "final_loss": losses[-1], "max_loss": max(losses), "mean_gain": float(np.mean(gains)),
117            "fraction_throttled": float(np.mean(np.array(gains) < 0.999999)),
118            "error_trace": errors, "loss_trace": losses, "gain_trace": gains, "delta_trace": deltas}
119
120
121def mini_experiment():
122    X, y, true = make_data()
123    # This is intentionally near/aggressively beyond the plain linear SGD stability range.
124    results = {}
125    for lr in [0.08, 0.16, 0.24]:
126        results[str(lr)] = {}
127        for method in ["sgd", "adam", "throttle"]:
128            results[str(lr)][method] = run_method(method, X, y, true, lr=lr, kappa=0.5)
129    return results
130
131
132def main():
133    out = {"seed": SEED, "math_checks": math_checks(), "mini_experiment": mini_experiment()}
134    Path("results.json").write_text(json.dumps(out, indent=2))
135    print(json.dumps({"math_checks": out["math_checks"], "summary": {
136        lr: {m: {k: v for k, v in r.items() if k in ["final_error", "max_error", "final_loss", "max_loss", "mean_gain", "fraction_throttled"]}
137             for m, r in d.items()} for lr, d in out["mini_experiment"].items()}}, indent=2))
138
139if __name__ == "__main__":
140    main()