import json, math from pathlib import Path import numpy as np SEED = 1180 rng = np.random.default_rng(SEED) def torus_dist(q, x): d = np.abs(x - q) d = np.minimum(d, 1.0 - d) return np.sqrt((d * d).sum(axis=-1)) def levina_bickel(dist, m=20, eps=1e-12): # dist excludes the query itself and is sorted along the last axis. rr = dist[..., :m] r = rr[..., -1:] logs = np.log((r + eps) / (rr[..., :-1] + eps)) k = 1.0 / (logs.mean(axis=-1) + eps) return np.clip(k, 1.0, dist.shape[-1]) def kernel_weights(d, r, khat): # The r^-khat factor is retained explicitly, though it cancels in rows. g = np.exp(-np.square(d / (r + 1e-12))) / np.power(r + 1e-12, khat) return g / (g.sum(axis=-1, keepdims=True) + 1e-15) def math_verification(): out = {} # Prediction 1: on a locally uniform k-dimensional cloud, kernel mass # scales as r^k, so log mass slope is k; weighted radius scales as r. rows = [] for k in range(1, 6): n, nq = 1400, 40 x = rng.random((n, k)) qi = rng.choice(n, nq, replace=False) q = x[qi] ds = np.stack([np.sort(torus_dist(v, x)[np.arange(n) != i]) for v, i in zip(q, qi)]) # Use a fixed radius range to measure continuum scaling; omit very tiny radii. radii = np.geomspace(0.055, 0.20, 7) masses, wr = [], [] for rad in radii: w = np.exp(-np.square(ds / rad)).mean(axis=1) a = np.exp(-np.square(ds / rad)) a /= a.sum(axis=1, keepdims=True) masses.append(w.mean()) wr.append(np.sqrt((a * ds * ds).sum(axis=1)).mean()) mass_slope = np.polyfit(np.log(radii), np.log(masses), 1)[0] radius_slope = np.polyfit(np.log(radii), np.log(wr), 1)[0] rows.append({"dimension": k, "predicted_mass_slope": k, "observed_mass_slope": float(mass_slope), "predicted_radius_slope": 1.0, "observed_radius_slope": float(radius_slope)}) out["continuum_scaling"] = rows # Prediction 2: multiplying the kernel by r^-p leaves raw interaction # mass scale-invariant exactly when p equals intrinsic dimension k. rows = [] for k in [1, 2, 3, 4]: n, nq = 1800, 50 x = rng.random((n, k)); qi = rng.choice(n, nq, replace=False) ds = np.stack([np.sort(torus_dist(x[i], x)[np.arange(n) != i]) for i in qi]) radii = np.geomspace(0.06, 0.16, 7) for p in [0.0, float(k), float(k + 1)]: masses = [np.exp(-np.square(ds / r)).sum(axis=1).mean() / r**p for r in radii] slope = np.polyfit(np.log(radii), np.log(masses), 1)[0] rows.append({"dimension": k, "exponent_p": p, "predicted_slope": k-p, "observed_slope": float(slope)}) out["dimension_scaling"] = rows # Prediction 3: bandwidth multiplier gamma gives effective support ~gamma^k # and weighted distance ~gamma (away from finite-domain saturation). rows = [] for k in [1, 2, 3, 4]: n, nq, m = 1800, 45, 20 x = rng.random((n, k)); qi = rng.choice(n, nq, replace=False) ds = np.stack([np.sort(torus_dist(x[i], x)[np.arange(n) != i]) for i in qi]) r0 = ds[:, m-1] gammas = np.geomspace(0.65, 2.0, 7) eff, radius = [], [] for gamma in gammas: w = kernel_weights(ds, gamma * r0[:, None], np.full((nq, 1), float(k))) eff.append((1.0 / np.square(w).sum(axis=1)).mean()) radius.append(np.sqrt((w * ds * ds).sum(axis=1)).mean()) slope_eff = np.polyfit(np.log(gammas), np.log(eff), 1)[0] slope_rad = np.polyfit(np.log(gammas), np.log(radius), 1)[0] rows.append({"dimension": k, "predicted_effective_support_slope": k, "observed_effective_support_slope": float(slope_eff), "predicted_radius_slope": 1.0, "observed_radius_slope": float(slope_rad)}) out["bandwidth_scaling"] = rows # Prediction 4: the Levina--Bickel estimator should recover local # intrinsic dimension from neighbor-distance log ratios. rows = [] for k in [1, 2, 3, 4]: n, nq = 1600, 45 x = rng.random((n, k)) qi = rng.choice(n, nq, replace=False) ds = np.stack([np.sort(torus_dist(x[i], x)[np.arange(n) != i]) for i in qi]) for m in [10, 20, 32]: estimates = levina_bickel(ds, m=m) rows.append({"dimension": k, "m": m, "predicted_mean_dimension": float(k), "observed_mean_dimension": float(estimates.mean()), "observed_std": float(estimates.std())}) out["intrinsic_dimension_estimation"] = rows return out def softmax(a): a = a - a.max(axis=-1, keepdims=True) e = np.exp(a) return e / e.sum(axis=-1, keepdims=True) def attention_comparison(): # Smooth regression on a 2-D manifold with noisy representations. n, d = 700, 2 xy = rng.random((n, d)) target = np.sin(2*np.pi*xy[:, 0]) + 0.7*np.cos(2*np.pi*xy[:, 1]) z = xy + 0.22*rng.normal(size=(n, d)) values = target + 0.45*rng.normal(size=n) # Hold out query points; keys/values are the remaining observations. qn = 120; qxy, qtarget = xy[:qn], target[:qn] keys, kvals = z[qn:], values[qn:] zd = np.stack([torus_dist(q, keys) for q in z[:qn]]) order = np.argsort(zd, axis=1); sd = np.take_along_axis(zd, order, axis=1) m = 20; r = sd[:, m-1] ak = kernel_weights(zd, r[:, None], np.full((qn,1), 2.0)) pred_kernel = ak @ kvals # Uniform aggregation is the locality-free attention baseline. pred_uniform = np.full(qn, kvals.mean()) # Standard dot-product attention with fixed random projections. Wq = rng.normal(size=(d, d)); Wk = rng.normal(size=(d, d)) logits = (z[:qn] @ Wq) @ (keys @ Wk).T / math.sqrt(d) ad = softmax(logits); pred_dot = ad @ kvals def mse(p): return float(np.mean((p-qtarget)**2)) def entropy(a): return float(np.mean(-(a*np.log(a+1e-15)).sum(axis=1))) return {"kernel_mse": mse(pred_kernel), "dot_product_mse": mse(pred_dot), "uniform_mse": mse(pred_uniform), "kernel_entropy": entropy(ak), "dot_product_entropy": entropy(ad), "kernel_mean_effective_neighbors": float(np.mean(1/(ak*ak).sum(1)))} def main(): result = {"seed": SEED, "math": math_verification(), "mini_experiment": attention_comparison()} Path("results.json").write_text(json.dumps(result, indent=2)) print(json.dumps(result, indent=2)) if __name__ == "__main__": main()