1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6SEED = 1362
  7rng = np.random.default_rng(SEED)
  8
  9
 10def normalize(x):
 11    return x / np.linalg.norm(x, axis=-1, keepdims=True)
 12
 13
 14def rope(x, positions, omegas):
 15    # x: [n,d], d even; rotations act independently on each pair.
 16    y = x.copy()
 17    for l, w in enumerate(omegas):
 18        a, b = 2*l, 2*l+1
 19        th = w * positions
 20        c, s = np.cos(th), np.sin(th)
 21        y[:, a] = c*x[:, a] - s*x[:, b]
 22        y[:, b] = s*x[:, a] + c*x[:, b]
 23    return y
 24
 25
 26def attention(x, positions, omegas, beta):
 27    q = rope(x, positions, omegas)
 28    scores = q @ q.T
 29    z = np.exp(beta * (scores - scores.max(axis=1, keepdims=True)))
 30    a = z / z.sum(axis=1, keepdims=True)
 31    return a, scores
 32
 33
 34def spherical_step(x, a, eta):
 35    m = a @ x
 36    v = m - (m*x).sum(axis=1, keepdims=True)*x
 37    return normalize(x + eta*v)
 38
 39
 40def angular_diameter(x):
 41    c = np.clip(x @ x.T, -1.0, 1.0)
 42    return float(np.arccos(c).max())
 43
 44
 45def make_consensus(n, d, positions, omegas, beta):
 46    u = normalize(rng.normal(size=(1,d)))[0]
 47    x = np.repeat(u[None,:], n, axis=0)
 48    a, scores = attention(x, positions, omegas, beta)
 49    return x, a, scores, u
 50
 51
 52def stability_trial(a, z, u, eta, steps=8, eps=1e-6):
 53    # At consensus, exp(beta*scores) is symmetric and A is reversible.
 54    # If S=D^(1/2) A D^(-1/2), S w=lambda w, then
 55    # A(D^(-1/2)w)=lambda(D^(-1/2)w). Use the worst transverse mode.
 56    dz = np.sqrt(z)
 57    s = dz[:, None] * a / dz[None, :]
 58    vals, vecs = np.linalg.eigh((s + s.T) / 2)
 59    lam_min = float(vals[0])
 60    mode = vecs[:, 0] / dz
 61    mode /= np.linalg.norm(mode)
 62    tangent = np.zeros_like(u)
 63    tangent[0] = 1.0
 64    if abs(np.dot(tangent, u)) > .9:
 65        tangent[:] = 0; tangent[1] = 1.0
 66    tangent -= np.dot(tangent, u)*u
 67    tangent /= np.linalg.norm(tangent)
 68    x = normalize(u[None, :] + eps*mode[:, None]*tangent[None, :])
 69    initial = np.linalg.norm(x - u[None, :])
 70    for _ in range(steps):
 71        x = spherical_step(x, a, eta)
 72    final = np.linalg.norm(x - u[None, :])
 73    return final / initial, lam_min
 74
 75
 76def main():
 77    n, d = 12, 8
 78    positions = np.arange(n, dtype=float)
 79    omegas = np.array([1.0, 0.37, 0.13, 0.047])
 80    betas = [0.5, 1.0, 2.0, 3.0]
 81    results = {"seed": SEED, "n": n, "d": d, "predictions": {}, "rows": []}
 82
 83    # Prediction 1: scores in [-1,1] imply exact softmax floor.
 84    x_random = normalize(rng.normal(size=(n,d)))
 85    a_rand, scores_rand = attention(x_random, positions, omegas, beta=2.0)
 86    alpha = math.exp(-4.0)/n
 87    floor_observed = float(a_rand.min())
 88    floor_pass = floor_observed + 1e-14 >= alpha
 89    delta = max(0.5*np.abs(a_rand[i]-a_rand[j]).sum() for i in range(n) for j in range(n))
 90    delta_bound = 1-math.exp(-4.0)
 91    results["predictions"]["kernel_floor"] = {
 92        "predicted_alpha": alpha, "observed_min_weight": floor_observed,
 93        "predicted_dobrushin_bound": delta_bound, "observed_dobrushin": float(delta),
 94        "pass": bool(floor_pass and delta <= delta_bound + 1e-12)
 95    }
 96
 97    # Prediction 2: norm-preserving update remains on the sphere.
 98    x0 = normalize(rng.normal(size=(n,d))); aa, _ = attention(x0, positions, omegas, beta=1.5)
 99    x1 = spherical_step(x0, aa, eta=2.0)
100    results["predictions"]["sphere_invariance"] = {
101        "max_norm_error": float(np.max(np.abs(np.linalg.norm(x1,axis=1)-1))),
102        "pass": bool(np.max(np.abs(np.linalg.norm(x1,axis=1)-1)) < 1e-12)
103    }
104
105    # Prediction 3: transverse Euler boundary eta_c=2/(1-lambda_min).
106    # At consensus the score matrix is symmetric, so the similarity transform is A itself.
107    for beta in betas:
108        xcons, a, scores, u = make_consensus(n,d,positions,omegas,beta)
109        z = np.exp(beta * scores).sum(axis=1)
110        dz = np.sqrt(z)
111        sim = dz[:, None] * a / dz[None, :]
112        ev = np.linalg.eigvalsh((sim + sim.T) / 2)
113        lam_min = float(ev[0])
114        eta_c = 2.0/(1.0-lam_min)
115        # sweep relative to predicted boundary; classify by growth of worst mode
116        ratios = []
117        for rel in [0.70, 0.90, 0.99, 1.01, 1.10, 1.30]:
118            eta = rel*eta_c
119            ratio, lam_mode = stability_trial(a, z, u, eta)
120            ratios.append({"relative_eta":rel, "growth_ratio":float(ratio), "grows":bool(ratio>1.01)})
121        # nearest grid transition (first point with clear growth)
122        growing = [r["relative_eta"] for r in ratios if r["grows"]]
123        observed_rel = min(growing) if growing else None
124        # Direct linear prediction for this eigenmode is |1+eta*(lambda_min-1)|.
125        predicted_at_110 = abs(1+1.10*eta_c*(lam_min-1))**8
126        results["rows"].append({
127            "beta":beta, "lambda_min":lam_min, "predicted_eta_c":eta_c,
128            "observed_transition_relative_grid":observed_rel,
129            "predicted_growth_ratio_at_1.10xc":float(predicted_at_110),
130            "sweep":ratios,
131            "dobrushin_bound":1-math.exp(-2*beta),
132            "dobrushin_observed":float(max(0.5*np.abs(a[i]-a[j]).sum() for i in range(n) for j in range(n)))
133        })
134    # Explicitly check the claimed boundary: all below 0.99 should contract, all 1.10 should grow.
135    below_ok = all(r["sweep"][2]["growth_ratio"] < 1 for r in results["rows"])
136    above_ok = all(r["sweep"][4]["growth_ratio"] > 1 for r in results["rows"])
137    results["predictions"]["euler_stability"] = {
138        "prediction": "eta<2/(1-lambda_min) contracts; eta>boundary grows",
139        "all_0.99x_contract": bool(below_ok), "all_1.10x_grow": bool(above_ok),
140        "pass": bool(below_ok and above_ok)
141    }
142    # A compact baseline comparison: unconstrained residual has norm growth, spherical does not.
143    xb = x0.copy()
144    for _ in range(20): xb = xb + 2.0*(aa@xb)
145    xs = x0.copy()
146    for _ in range(20): xs = spherical_step(xs, aa, 2.0)
147    results["baseline_vs_idea"] = {
148        "unconstrained_max_norm_after_20": float(np.linalg.norm(xb,axis=1).max()),
149        "spherical_max_norm_after_20": float(np.linalg.norm(xs,axis=1).max()),
150        "initial_angular_diameter": angular_diameter(x0),
151        "spherical_angular_diameter_after_20": angular_diameter(xs)
152    }
153    out = Path("results.json")
154    out.write_text(json.dumps(results, indent=2))
155    print(json.dumps(results, indent=2))
156
157if __name__ == "__main__":
158    main()