1import json, math
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1180
  6rng = np.random.default_rng(SEED)
  7
  8
  9def torus_dist(q, x):
 10    d = np.abs(x - q)
 11    d = np.minimum(d, 1.0 - d)
 12    return np.sqrt((d * d).sum(axis=-1))
 13
 14
 15def levina_bickel(dist, m=20, eps=1e-12):
 16    # dist excludes the query itself and is sorted along the last axis.
 17    rr = dist[..., :m]
 18    r = rr[..., -1:]
 19    logs = np.log((r + eps) / (rr[..., :-1] + eps))
 20    k = 1.0 / (logs.mean(axis=-1) + eps)
 21    return np.clip(k, 1.0, dist.shape[-1])
 22
 23
 24def kernel_weights(d, r, khat):
 25    # The r^-khat factor is retained explicitly, though it cancels in rows.
 26    g = np.exp(-np.square(d / (r + 1e-12))) / np.power(r + 1e-12, khat)
 27    return g / (g.sum(axis=-1, keepdims=True) + 1e-15)
 28
 29
 30def math_verification():
 31    out = {}
 32    # Prediction 1: on a locally uniform k-dimensional cloud, kernel mass
 33    # scales as r^k, so log mass slope is k; weighted radius scales as r.
 34    rows = []
 35    for k in range(1, 6):
 36        n, nq = 1400, 40
 37        x = rng.random((n, k))
 38        qi = rng.choice(n, nq, replace=False)
 39        q = x[qi]
 40        ds = np.stack([np.sort(torus_dist(v, x)[np.arange(n) != i]) for v, i in zip(q, qi)])
 41        # Use a fixed radius range to measure continuum scaling; omit very tiny radii.
 42        radii = np.geomspace(0.055, 0.20, 7)
 43        masses, wr = [], []
 44        for rad in radii:
 45            w = np.exp(-np.square(ds / rad)).mean(axis=1)
 46            a = np.exp(-np.square(ds / rad))
 47            a /= a.sum(axis=1, keepdims=True)
 48            masses.append(w.mean())
 49            wr.append(np.sqrt((a * ds * ds).sum(axis=1)).mean())
 50        mass_slope = np.polyfit(np.log(radii), np.log(masses), 1)[0]
 51        radius_slope = np.polyfit(np.log(radii), np.log(wr), 1)[0]
 52        rows.append({"dimension": k, "predicted_mass_slope": k,
 53                     "observed_mass_slope": float(mass_slope),
 54                     "predicted_radius_slope": 1.0,
 55                     "observed_radius_slope": float(radius_slope)})
 56    out["continuum_scaling"] = rows
 57
 58    # Prediction 2: multiplying the kernel by r^-p leaves raw interaction
 59    # mass scale-invariant exactly when p equals intrinsic dimension k.
 60    rows = []
 61    for k in [1, 2, 3, 4]:
 62        n, nq = 1800, 50
 63        x = rng.random((n, k)); qi = rng.choice(n, nq, replace=False)
 64        ds = np.stack([np.sort(torus_dist(x[i], x)[np.arange(n) != i]) for i in qi])
 65        radii = np.geomspace(0.06, 0.16, 7)
 66        for p in [0.0, float(k), float(k + 1)]:
 67            masses = [np.exp(-np.square(ds / r)).sum(axis=1).mean() / r**p for r in radii]
 68            slope = np.polyfit(np.log(radii), np.log(masses), 1)[0]
 69            rows.append({"dimension": k, "exponent_p": p,
 70                         "predicted_slope": k-p, "observed_slope": float(slope)})
 71    out["dimension_scaling"] = rows
 72
 73    # Prediction 3: bandwidth multiplier gamma gives effective support ~gamma^k
 74    # and weighted distance ~gamma (away from finite-domain saturation).
 75    rows = []
 76    for k in [1, 2, 3, 4]:
 77        n, nq, m = 1800, 45, 20
 78        x = rng.random((n, k)); qi = rng.choice(n, nq, replace=False)
 79        ds = np.stack([np.sort(torus_dist(x[i], x)[np.arange(n) != i]) for i in qi])
 80        r0 = ds[:, m-1]
 81        gammas = np.geomspace(0.65, 2.0, 7)
 82        eff, radius = [], []
 83        for gamma in gammas:
 84            w = kernel_weights(ds, gamma * r0[:, None], np.full((nq, 1), float(k)))
 85            eff.append((1.0 / np.square(w).sum(axis=1)).mean())
 86            radius.append(np.sqrt((w * ds * ds).sum(axis=1)).mean())
 87        slope_eff = np.polyfit(np.log(gammas), np.log(eff), 1)[0]
 88        slope_rad = np.polyfit(np.log(gammas), np.log(radius), 1)[0]
 89        rows.append({"dimension": k, "predicted_effective_support_slope": k,
 90                     "observed_effective_support_slope": float(slope_eff),
 91                     "predicted_radius_slope": 1.0,
 92                     "observed_radius_slope": float(slope_rad)})
 93    out["bandwidth_scaling"] = rows
 94
 95    # Prediction 4: the Levina--Bickel estimator should recover local
 96    # intrinsic dimension from neighbor-distance log ratios.
 97    rows = []
 98    for k in [1, 2, 3, 4]:
 99        n, nq = 1600, 45
100        x = rng.random((n, k))
101        qi = rng.choice(n, nq, replace=False)
102        ds = np.stack([np.sort(torus_dist(x[i], x)[np.arange(n) != i]) for i in qi])
103        for m in [10, 20, 32]:
104            estimates = levina_bickel(ds, m=m)
105            rows.append({"dimension": k, "m": m,
106                         "predicted_mean_dimension": float(k),
107                         "observed_mean_dimension": float(estimates.mean()),
108                         "observed_std": float(estimates.std())})
109    out["intrinsic_dimension_estimation"] = rows
110    return out
111
112
113def softmax(a):
114    a = a - a.max(axis=-1, keepdims=True)
115    e = np.exp(a)
116    return e / e.sum(axis=-1, keepdims=True)
117
118
119def attention_comparison():
120    # Smooth regression on a 2-D manifold with noisy representations.
121    n, d = 700, 2
122    xy = rng.random((n, d))
123    target = np.sin(2*np.pi*xy[:, 0]) + 0.7*np.cos(2*np.pi*xy[:, 1])
124    z = xy + 0.22*rng.normal(size=(n, d))
125    values = target + 0.45*rng.normal(size=n)
126    # Hold out query points; keys/values are the remaining observations.
127    qn = 120; qxy, qtarget = xy[:qn], target[:qn]
128    keys, kvals = z[qn:], values[qn:]
129    zd = np.stack([torus_dist(q, keys) for q in z[:qn]])
130    order = np.argsort(zd, axis=1); sd = np.take_along_axis(zd, order, axis=1)
131    m = 20; r = sd[:, m-1]
132    ak = kernel_weights(zd, r[:, None], np.full((qn,1), 2.0))
133    pred_kernel = ak @ kvals
134    # Uniform aggregation is the locality-free attention baseline.
135    pred_uniform = np.full(qn, kvals.mean())
136    # Standard dot-product attention with fixed random projections.
137    Wq = rng.normal(size=(d, d)); Wk = rng.normal(size=(d, d))
138    logits = (z[:qn] @ Wq) @ (keys @ Wk).T / math.sqrt(d)
139    ad = softmax(logits); pred_dot = ad @ kvals
140    def mse(p): return float(np.mean((p-qtarget)**2))
141    def entropy(a):
142        return float(np.mean(-(a*np.log(a+1e-15)).sum(axis=1)))
143    return {"kernel_mse": mse(pred_kernel), "dot_product_mse": mse(pred_dot),
144            "uniform_mse": mse(pred_uniform), "kernel_entropy": entropy(ak),
145            "dot_product_entropy": entropy(ad), "kernel_mean_effective_neighbors": float(np.mean(1/(ak*ak).sum(1)))}
146
147
148def main():
149    result = {"seed": SEED, "math": math_verification(), "mini_experiment": attention_comparison()}
150    Path("results.json").write_text(json.dumps(result, indent=2))
151    print(json.dumps(result, indent=2))
152
153if __name__ == "__main__":
154    main()