Team Ai
Apppublic

ProCreations/repro-formal-problem-solving

sourceHugging Faceupdated 2mo agoView on Hugging Face
0likes
audit_wire.py332 linesDownload Raw Back to root
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