1"""Gauge-Patched Local Experts: toy math checks and a tiny optimization study.
  2
  3Run: python3 gauge_patched_experts.py
  4Outputs: results.json and a concise terminal report.
  5"""
  6import json, math
  7from pathlib import Path
  8import numpy as np
  9import torch
 10
 11SEED = 1270
 12np.random.seed(SEED)
 13torch.manual_seed(SEED)
 14torch.set_num_threads(4)
 15
 16
 17def rot(theta):
 18    c, s = np.cos(theta), np.sin(theta)
 19    return np.array([[c, -s], [s, c]], dtype=float)
 20
 21
 22def holder(vals, xs, beta):
 23    # Dense pairwise Holder seminorm, adequate for this 1D sanity check.
 24    dx = np.abs(xs[:, None] - xs[None, :])
 25    dv = np.abs(vals[:, None] - vals[None, :])
 26    mask = dx > 1e-12
 27    return float(np.max(dv[mask] / (dx[mask] ** beta)))
 28
 29
 30def math_checks():
 31    beta = 0.75
 32    # Prediction 1: exp(skew) is orthogonal, hence preserves every feature norm.
 33    norm_errors = []
 34    for theta in np.linspace(-3.0, 3.0, 31):
 35        h = np.array([0.7, -1.3])
 36        g = rot(theta)  # exact SO(2) exponential
 37        norm_errors.append(abs(np.linalg.norm(g @ h) - np.linalg.norm(h)))
 38    p1 = {"prediction": "orthogonal transition preserves feature norm", "max_abs_norm_error": max(norm_errors)}
 39
 40    # Prediction 2: a transition extended from identity across a gap d has
 41    # Holder size proportional to d^{-beta}. Use a smooth cubic bump.
 42    theta = 0.4
 43    ds = np.array([0.04, 0.06, 0.09, 0.13, 0.19, 0.28])
 44    hs = []
 45    xs = np.linspace(-1.0, 1.0, 1601)
 46    for d in ds:
 47        # transition changes only across [0,d], identity before and rot(theta) after
 48        t = np.clip(xs / d, 0, 1)
 49        bump = t*t*(3 - 2*t)
 50        # angle field; Holder norm of matrix entries is equivalent up to constants
 51        vals = np.array([rot(theta*a)[0, 0] for a in bump])
 52        hs.append(holder(vals, xs, beta))
 53    slope = float(np.polyfit(np.log(ds), np.log(hs), 1)[0])
 54    p2 = {"prediction": "extension Holder norm scales as gap^{-beta}", "beta": beta,
 55          "observed_loglog_slope": slope, "predicted_slope": -beta,
 56          "gaps": ds.tolist(), "holder_norms": [float(x) for x in hs]}
 57
 58    # Prediction 3: for fixed h, ||h - R(eps)h||^2 is quadratic in eps.
 59    h = np.array([1.2, -0.8])
 60    eps = np.array([0.005, 0.01, 0.02, 0.04, 0.08, 0.16])
 61    residuals = np.array([np.linalg.norm(h - rot(e) @ h)**2 for e in eps])
 62    slope3 = float(np.polyfit(np.log(eps), np.log(residuals), 1)[0])
 63    ratios = residuals / (eps**2 * np.dot(h, h))
 64    p3 = {"prediction": "small-angle patch residual is quadratic", "observed_loglog_slope": slope3,
 65          "predicted_slope": 2.0, "residual_over_eps2_norm2": [float(x) for x in ratios]}
 66    return [p1, p2, p3]
 67
 68
 69def tiny_experiment():
 70    """Two local experts see differently rotated coordinates of the same target.
 71    Both variants have the same trainable expert offsets and task loss. The idea
 72    adds the overlap alignment penalty using the known orthogonal transition.
 73    """
 74    torch.manual_seed(SEED)
 75    n = 512
 76    z = torch.randn(n, 2)
 77    true_w = torch.tensor([[1.0, -0.4], [0.3, 0.8]])
 78    y = z @ true_w.T
 79    angle = 0.9
 80    G = torch.tensor(rot(angle), dtype=torch.float32)
 81    # Expert 0 receives z, expert 1 receives G z; target is shared latent output.
 82    x0, x1 = z, z @ G.T
 83    torch.manual_seed(SEED + 77)
 84    init0 = torch.randn(2,2)*0.5
 85    init1 = torch.randn(2,2)*0.5
 86    def run(lam, steps=500):
 87        # Matched initialization isolates the effect of patch regularization.
 88        W0 = torch.nn.Parameter(init0.clone())
 89        W1 = torch.nn.Parameter(init1.clone())
 90        opt = torch.optim.Adam([W0,W1], lr=0.04)
 91        for _ in range(steps):
 92            h0, h1 = x0 @ W0.T, x1 @ W1.T
 93            task = ((h0-y)**2).mean() + ((h1-y)**2).mean()
 94            # h0 ~= G^{-1} h1, equivalently h0 ~= G.T h1.
 95            patch = ((h0 - h1 @ G)**2).mean()
 96            loss = task + lam * patch
 97            opt.zero_grad(); loss.backward(); opt.step()
 98        with torch.no_grad():
 99            h0, h1 = x0 @ W0.T, x1 @ W1.T
100            task = float((((h0-y)**2).mean()+((h1-y)**2).mean()).item()/2)
101            disagreement = float(((h0-h1@G)**2).mean().item())
102        return task, disagreement
103    lambdas = [0.0, 0.01, 0.1, 1.0, 10.0]
104    sweep = []
105    for lam in lambdas:
106        task, disagreement = run(lam)
107        sweep.append({"lambda_patch": lam, "task_mse": task, "overlap_mse": disagreement})
108    return {"metric": "mean task MSE; aligned overlap MSE", "matched_initialization": True,
109            "baseline": sweep[0], "gauge_patched_lambda_1": sweep[3], "lambda_sweep": sweep}
110
111
112def main():
113    checks = math_checks()
114    exp = tiny_experiment()
115    out = {"seed": SEED, "math_checks": checks, "tiny_experiment": exp}
116    Path("results.json").write_text(json.dumps(out, indent=2))
117    print(json.dumps(out, indent=2))
118
119if __name__ == "__main__":
120    main()