ProCreations/repro-formal-problem-solving
0
1#!/usr/bin/env python32"""Clean-room exact audit of the three WIRE automatic claims."""3 4from __future__ import annotations5 6import hashlib7import json8import platform9import time10from pathlib import Path11 12import matplotlib13 14matplotlib.use("Agg")15import matplotlib.pyplot as plt16import numpy as np17import scipy18from numpy.polynomial.hermite import hermgauss19 20np.seterr(all="ignore")21 22 23HERE = Path(__file__).resolve().parent24OUTPUT_DIR = HERE / "outputs" / "wire"25RESULTS = OUTPUT_DIR / "wire_results.json"26FIGURE = OUTPUT_DIR / "wire_audit.png"27PAPER_SHA = "3eb6899ac0da995483dfba1eafe1ed625a1673d638d498882bcf741709f3415f"28OFFICIAL_COMMIT = "4ac067eb38272543b0cdd7591d630399ff37bce4"29 30 31def laplacian_from_edges(n, edges):32 a = np.zeros((n, n), dtype=float)33 for i, j in edges:34 a[i, j] = a[j, i] = 1.035 return np.diag(a.sum(axis=1)) - a36 37 38def path_laplacian(n):39 return laplacian_from_edges(n, [(i, i + 1) for i in range(n - 1)])40 41 42def cycle_laplacian(n):43 return laplacian_from_edges(n, [(i, (i + 1) % n) for i in range(n)])44 45 46def star_laplacian(n):47 return laplacian_from_edges(n, [(0, i) for i in range(1, n)])48 49 50def grid_laplacian(nx, ny):51 edges = []52 for y in range(ny):53 for x in range(nx):54 i = y * nx + x55 if x + 1 < nx:56 edges.append((i, i + 1))57 if y + 1 < ny:58 edges.append((i, i + nx))59 return laplacian_from_edges(nx * ny, edges)60 61 62def random_connected_laplacian(n, rng, p=0.25):63 edges = {(i, i + 1) for i in range(n - 1)}64 for i in range(n):65 for j in range(i + 2, n):66 if rng.random() < p:67 edges.add((i, j))68 return laplacian_from_edges(n, sorted(edges))69 70 71def graph_library():72 rng = np.random.default_rng(2606)73 out = []74 for n in [6, 9, 13]:75 out.extend([(f"path-{n}", path_laplacian(n)), (f"cycle-{n}", cycle_laplacian(n)), (f"star-{n}", star_laplacian(n))])76 for nx, ny in [(3, 3), (3, 5), (4, 6)]:77 out.append((f"grid-{nx}x{ny}", grid_laplacian(nx, ny)))78 for n in [8, 12, 18]:79 out.append((f"random-{n}", random_connected_laplacian(n, rng)))80 return out81 82 83def spectral_coordinates(L, resistance_scaled=False):84 vals, vecs = np.linalg.eigh(L)85 keep = vals > 1e-1086 vals, vecs = vals[keep], vecs[:, keep]87 if resistance_scaled:88 vecs = vecs / np.sqrt(vals)[None, :]89 return vals, vecs90 91 92def rotate_rows(x, theta):93 """Apply independent 2D rotations to consecutive feature pairs."""94 pairs = x.reshape(x.shape[0], -1, 2)95 c, s = np.cos(theta), np.sin(theta)96 out = np.empty_like(pairs)97 out[..., 0] = pairs[..., 0] * c - pairs[..., 1] * s98 out[..., 1] = pairs[..., 0] * s + pairs[..., 1] * c99 return out.reshape(x.shape)100 101 102def claim1_spectral_structure():103 rows = []104 rng = np.random.default_rng(1807)105 for name, L in graph_library():106 vals, vecs = spectral_coordinates(L)107 m = min(5, vecs.shape[1])108 coords = vecs[:, :m]109 omega = np.linspace(0.7, 1.3, m)110 theta = coords @ omega111 score = np.cos(theta[:, None] - theta[None, :])112 113 perm = rng.permutation(len(L))114 P = np.eye(len(L))[perm]115 Lp = P @ L @ P.T116 vp = P @ vecs[:, :m]117 residual = np.linalg.norm(Lp @ vp - vp * vals[:m][None, :], ord="fro")118 score_p = np.cos((P @ theta)[:, None] - (P @ theta)[None, :])119 equivariance_error = np.max(np.abs(score_p - P @ score @ P.T))120 121 q = np.tile([1.0, 0.0], (len(L), 1))122 q_rot = rotate_rows(q, theta[:, None])123 direct_score = q_rot @ q_rot.T124 relative_rotation_error = np.max(np.abs(direct_score - score))125 offdiag = score[~np.eye(len(L), dtype=bool)]126 rows.append(127 {128 "graph": name,129 "nodes": len(L),130 "spectral_dimensions": m,131 "permuted_eigen_residual_fro": float(residual),132 "permutation_equivariance_error": float(equivariance_error),133 "absolute_vs_relative_rotation_error": float(relative_rotation_error),134 "offdiagonal_score_standard_deviation": float(np.std(offdiag)),135 "distinct_scores_rounded_1e-10": int(len(np.unique(np.round(offdiag, 10)))),136 }137 )138 return {139 "protocol": "15 connected graphs; Laplacian eigenfeatures drive pairwise rotary logits; fixed node permutations are checked algebraically",140 "rows": rows,141 "graphs": len(rows),142 "max_permuted_eigen_residual_fro": max(r["permuted_eigen_residual_fro"] for r in rows),143 "max_permutation_equivariance_error": max(r["permutation_equivariance_error"] for r in rows),144 "max_relative_rotation_error": max(r["absolute_vs_relative_rotation_error"] for r in rows),145 "min_offdiagonal_score_standard_deviation": min(r["offdiagonal_score_standard_deviation"] for r in rows),146 "min_distinct_scores": min(r["distinct_scores_rounded_1e-10"] for r in rows),147 }148 149 150def claim2_grids_and_resistance():151 path_rows = []152 for n in [4, 5, 8, 13, 21, 34, 55, 64]:153 vals, vecs = np.linalg.eigh(path_laplacian(n))154 analytic = -np.cos(np.pi * (np.arange(n) + 0.5) / n)155 analytic /= np.linalg.norm(analytic)156 corr = abs(float(analytic @ vecs[:, 1]))157 eigenvalue_error = abs(vals[1] - (2.0 - 2.0 * np.cos(np.pi / n)))158 path_rows.append({"n": n, "absolute_mode_correlation": corr, "eigenvalue_error": float(eigenvalue_error)})159 160 grid_rows = []161 for nx, ny in [(3, 4), (3, 7), (4, 4), (4, 9), (5, 6), (6, 8), (7, 9), (8, 10)]:162 L = grid_laplacian(nx, ny)163 xx = np.tile(np.arange(nx), ny)164 yy = np.repeat(np.arange(ny), nx)165 ux = np.cos(np.pi * (xx + 0.5) / nx)166 uy = np.cos(np.pi * (yy + 0.5) / ny)167 lamx = 2.0 - 2.0 * np.cos(np.pi / nx)168 lamy = 2.0 - 2.0 * np.cos(np.pi / ny)169 rx = np.linalg.norm(L @ ux - lamx * ux) / np.linalg.norm(ux)170 ry = np.linalg.norm(L @ uy - lamy * uy) / np.linalg.norm(uy)171 grid_rows.append({"grid": f"{nx}x{ny}", "x_mode_residual": float(rx), "y_mode_residual": float(ry)})172 173 hx, hw = hermgauss(80)174 resistance_rows, slopes = [], []175 omega_grid = np.asarray([0.01, 0.02, 0.04, 0.08, 0.16])176 for name, L in graph_library():177 vals, coords = spectral_coordinates(L, resistance_scaled=True)178 pinv = np.linalg.pinv(L, hermitian=True)179 diag = np.diag(pinv)180 R_pinv = diag[:, None] + diag[None, :] - 2.0 * pinv181 R_coords = np.sum((coords[:, None, :] - coords[None, :, :]) ** 2, axis=2)182 resistance_error = np.max(np.abs(R_pinv - R_coords))183 pairs = [(i, j) for i in range(len(L)) for j in range(i + 1, len(L))]184 # Cover low, median, and high resistance pairs deterministically.185 pairs.sort(key=lambda ij: R_pinv[ij])186 selected = [pairs[0], pairs[len(pairs) // 2], pairs[-1]]187 for i, j in selected:188 R = float(R_pinv[i, j])189 exact = np.exp(-0.5 * omega_grid**2 * R)190 leading = 1.0 - 0.5 * omega_grid**2 * R191 residuals = np.abs(exact - leading)192 slope = float(np.polyfit(np.log(omega_grid), np.log(residuals), 1)[0])193 slopes.append(slope)194 # Independent Gauss-Hermite expectation E cos(omega*sqrt(R)*Z).195 gh = np.array([np.sum(hw * np.cos(w * np.sqrt(R) * np.sqrt(2) * hx)) / np.sqrt(np.pi) for w in omega_grid])196 resistance_rows.append(197 {198 "graph": name,199 "pair": [i, j],200 "effective_resistance": R,201 "spectral_vs_pseudoinverse_max_error_graph": float(resistance_error),202 "small_omega_remainder_loglog_slope": slope,203 "max_gauss_hermite_vs_exact_expectation_error": float(np.max(np.abs(gh - exact))),204 "remainder_over_omega4_at_smallest_omega": float(residuals[0] / omega_grid[0] ** 4),205 }206 )207 return {208 "path_grid_protocol": "closed-form first path/grid Laplacian modes, up to sign and normalization",209 "path_rows": path_rows,210 "grid_rows": grid_rows,211 "min_path_mode_correlation": min(r["absolute_mode_correlation"] for r in path_rows),212 "max_path_eigenvalue_error": max(r["eigenvalue_error"] for r in path_rows),213 "max_grid_mode_residual": max(max(r["x_mode_residual"], r["y_mode_residual"]) for r in grid_rows),214 "resistance_protocol": "Laplacian-pseudoinverse resistance vs squared resistance-scaled spectral distance; exact Gaussian WIRE expectation and independent 80-node Gauss-Hermite oracle",215 "resistance_rows": resistance_rows,216 "pairs": len(resistance_rows),217 "max_spectral_vs_pseudoinverse_resistance_error": max(r["spectral_vs_pseudoinverse_max_error_graph"] for r in resistance_rows),218 "remainder_slope_range": [min(slopes), max(slopes)],219 "max_gauss_hermite_expectation_error": max(r["max_gauss_hermite_vs_exact_expectation_error"] for r in resistance_rows),220 }221 222 223def claim3_linear_attention():224 rng = np.random.default_rng(314159)225 rows = []226 for n in [8, 32, 128, 512, 1024, 2048]:227 d, dv, m = 32, 7, 6228 q, k, v = rng.normal(size=(n, d)), rng.normal(size=(n, d)), rng.normal(size=(n, dv))229 coords, omega = rng.normal(size=(n, m)), rng.normal(size=(d // 2, m))230 theta = coords @ omega.T231 qr, kr = rotate_rows(q, theta), rotate_rows(k, theta)232 associative = qr @ (kr.T @ v)233 pairwise = (qr @ kr.T) @ v234 relerr = np.linalg.norm(associative - pairwise) / np.linalg.norm(pairwise)235 dense_elements = n * n236 associative_elements = d * dv237 rows.append(238 {239 "tokens": n,240 "feature_dim": d,241 "value_dim": dv,242 "relative_associativity_error": float(relerr),243 "dense_attention_elements": dense_elements,244 "associative_intermediate_elements": associative_elements,245 "intermediate_memory_ratio": dense_elements / associative_elements,246 "dense_leading_multiply_adds": int(2 * n * n * d),247 "associative_leading_multiply_adds": int(2 * n * d * dv),248 }249 )250 return {251 "identity": "(Q_rot K_rot^T)V = Q_rot(K_rot^T V); WIRE acts on Q and K before the associative linear-attention contraction",252 "rows": rows,253 "max_relative_associativity_error": max(r["relative_associativity_error"] for r in rows),254 "largest_n_memory_ratio": rows[-1]["intermediate_memory_ratio"],255 "largest_n_operation_ratio": rows[-1]["dense_leading_multiply_adds"] / rows[-1]["associative_leading_multiply_adds"],256 }257 258 259def make_figure(c2, c3):260 fig, axes = plt.subplots(1, 3, figsize=(15, 4.4))261 ns = [r["n"] for r in c2["path_rows"]]262 corr_err = [1 - r["absolute_mode_correlation"] for r in c2["path_rows"]]263 axes[0].semilogy(ns, np.maximum(corr_err, 1e-17), "o-")264 axes[0].set(xlabel="path nodes", ylabel="1 − |mode correlation|", title="Grid modes recover RoPE coordinates")265 axes[0].grid(alpha=0.25)266 267 slopes = [r["small_omega_remainder_loglog_slope"] for r in c2["resistance_rows"]]268 axes[1].hist(slopes, bins=10, color="#6d5dfc", alpha=0.85)269 axes[1].axvline(4, color="black", linestyle="--", label="O(ω⁴)")270 axes[1].set(xlabel="fitted remainder exponent", ylabel="graph-pair count", title="Effective-resistance expansion")271 axes[1].legend()272 273 nr = [r["tokens"] for r in c3["rows"]]274 mr = [r["intermediate_memory_ratio"] for r in c3["rows"]]275 axes[2].loglog(nr, mr, "o-", color="#e45756")276 axes[2].set(xlabel="tokens N", ylabel="dense / associative intermediate", title="WIRE keeps linear attention associative")277 axes[2].grid(alpha=0.25, which="both")278 fig.tight_layout()279 fig.savefig(FIGURE, dpi=180)280 281 282def main():283 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)284 t0 = time.time()285 c1 = claim1_spectral_structure()286 c2 = claim2_grids_and_resistance()287 c3 = claim3_linear_attention()288 289 assertions = {290 "claim1_all_graphs_structurally_nonconstant": c1["min_offdiagonal_score_standard_deviation"] > 1e-5 and c1["min_distinct_scores"] >= 3,291 "claim1_equivariance_and_rotation_identities": c1["max_permutation_equivariance_error"] < 1e-12 and c1["max_relative_rotation_error"] < 1e-12,292 "claim2_path_and_grid_modes": c2["min_path_mode_correlation"] > 1 - 1e-12 and c2["max_grid_mode_residual"] < 1e-12,293 "claim2_resistance_identity": c2["max_spectral_vs_pseudoinverse_resistance_error"] < 1e-11,294 "claim2_omega4_remainder": c2["remainder_slope_range"][0] > 3.94 and c2["remainder_slope_range"][1] < 4.01,295 "claim2_expectation_oracle": c2["max_gauss_hermite_expectation_error"] < 1e-12,296 "claim3_linear_attention_associativity": c3["max_relative_associativity_error"] < 1e-11,297 "claim3_subquadratic_intermediate": c3["largest_n_memory_ratio"] > 10000,298 }299 assert all(assertions.values()), assertions300 make_figure(c2, c3)301 302 payload = {303 "paper": {304 "title": "Rotary Position Encodings for Graphs",305 "openreview": "trn64znfNx",306 "arxiv": "2509.22259",307 "paper_pdf_sha256": PAPER_SHA,308 "official_repository": "https://github.com/cederikhoefs/Graph-RoPE",309 "official_commit": OFFICIAL_COMMIT,310 "relationship": "clean-room NumPy/SciPy implementation of paper equations; author code not executed",311 },312 "automatic_claims": {313 "claim_1_spectral_structural_rotation": c1,314 "claim_2_grid_rope_and_effective_resistance": c2,315 "claim_3_linear_attention_compatibility": c3,316 },317 "assertions": assertions,318 "limitations": [319 "Claims are audited at the mathematical operator level; this does not retrain the paper's more than 200 graph models.",320 "The effective-resistance statement is the paper's random-frequency small-omega expectation, not exact invariance for one learned frequency draw.",321 "The linear-attention check targets the bilinear numerator factorization; normalization can be computed separately as in standard linear attention.",322 ],323 "environment": {"python": platform.python_version(), "numpy": np.__version__, "scipy": scipy.__version__},324 "runtime_seconds": time.time() - t0,325 }326 RESULTS.write_text(json.dumps(payload, indent=2) + "\n")327 print(json.dumps({"assertions": assertions, "runtime_seconds": payload["runtime_seconds"]}, indent=2))328 329 330if __name__ == "__main__":331 main()332 