scripts/claim4_sigreg_audit.py

"""Claim 4 — SIGReg mechanism audit (CPU part).

  (A) Target correctness: the authors' loss matches the empirical
      characteristic function of 1-d projections against the standard normal
      CF e^{-t^2/2}; for N(0, I_512) embeddings the loss should be near the
      finite-sample floor, and it should rise for (i) collapsed (low-rank),
      (ii) anisotropic, (iii) heavy-tailed, and (iv) shifted embeddings.

  (B) Gradient direction sanity: minimizing the loss by gradient descent on a
      free embedding matrix (no network) drives an anisotropic heavy-tailed
      cloud toward isotropic Gaussian: condition number of the covariance
      falls toward 1 with eigenvalues near sigma^2 = 1, and projection excess
      kurtosis goes to ~0.

Uses the SIGReg module copied verbatim from the authors' cifar.py.
Deterministic (seeded), CPU-only. Exits non-zero on failure.
"""
import csv
import os
import sys

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(__file__))
from plot_style import SERIES, apply_style
from cifar_repro import SIGReg

import matplotlib.pyplot as plt

OUT = os.path.join(os.path.dirname(__file__), "..", "outputs", "claim4")
os.makedirs(OUT, exist_ok=True)
torch.manual_seed(0)

D, B = 512, 1024
sigreg = SIGReg(embedding_dim=D)


def loss_of(x, n_avg=20):
    vals = [float(sigreg(x)) for _ in range(n_avg)]  # avg over random slices
    return float(np.mean(vals)), float(np.std(vals))


def axis_loss(x, n_axes=16):
    """Same CF objective but along the first n_axes coordinate axes instead of
    random dense directions — isolates the CLT effect of random slicing."""
    s = x[:, :n_axes].t()
    t = sigreg.t_grid
    loss = 0.0
    for ti in range(sigreg.num_t):
        tt = t[ti]
        re = torch.cos(tt * s).mean(dim=1)
        im = torch.sin(tt * s).mean(dim=1)
        target = float(torch.exp(-0.5 * tt ** 2))
        loss += ((re - target) ** 2 + (im ** 2)).mean()
    return float(loss / sigreg.num_t)


def check_A():
    cases = {}
    cases["isotropic N(0,I)"] = torch.randn(B, D)
    # rank-8 with the same total variance as N(0, I): cov = A^T A, Tr ~= D
    z = torch.randn(B, 8)
    cases["collapsed rank-8"] = z @ (torch.randn(8, D) / 8 ** 0.5)
    cases["variance-doubled N(0,2I)"] = torch.randn(B, D) * 2 ** 0.5
    cases["mean-shifted N(2,I)"] = torch.randn(B, D) + 2.0
    scales = torch.linspace(0.05, 3.0, D).sqrt()
    cases["anisotropic (same trace)"] = torch.randn(B, D) * scales * (D / (scales ** 2).sum()).sqrt()
    lap = torch.distributions.Laplace(0.0, 1 / 2 ** 0.5).sample((B, D))
    cases["heavy-tailed Laplace"] = lap

    rows = []
    for name, x in cases.items():
        m, s = loss_of(x)
        rows.append(dict(case=name, sigreg_loss_mean=m, sigreg_loss_std=s,
                         axis_cf_loss=axis_loss(x)))
        print(f"[A] {name:<26} SIGReg loss = {m:.4f} ± {s:.4f} "
              f"(axis-aligned CF loss {rows[-1]['axis_cf_loss']:.4f})")
    d = {r["case"]: r for r in rows}
    base = d["isotropic N(0,I)"]["sigreg_loss_mean"]
    # What the paper's use of SIGReg requires: the failure modes RL training
    # exhibits (rank collapse; scale/shift drift) must register strongly.
    ok = (d["collapsed rank-8"]["sigreg_loss_mean"] > 5 * base
          and d["variance-doubled N(0,2I)"]["sigreg_loss_mean"] > 5 * base
          and d["mean-shifted N(2,I)"]["sigreg_loss_mean"] > 5 * base)
    print(f"[A] collapse / variance / mean-shift all >5x the isotropic floor -> "
          f"{'PASS' if ok else 'FAIL'}")
    # Measured sensitivity limits (reported, not asserted): random-slice CF is
    # near-blind per batch to moderate same-trace anisotropy and to
    # coordinate-wise heavy tails (CLT along dense random directions) — the
    # axis-aligned probe shows the Laplace tails ARE detectable along axes.
    print(f"[A] measured limits: anisotropic-same-trace loss "
          f"{d['anisotropic (same trace)']['sigreg_loss_mean']:.4f} ~ floor {base:.4f}; "
          f"Laplace random-slice {d['heavy-tailed Laplace']['sigreg_loss_mean']:.4f} "
          f"~ floor, but axis-aligned {d['heavy-tailed Laplace']['axis_cf_loss']:.4f} "
          f">> isotropic axis value {d['isotropic N(0,I)']['axis_cf_loss']:.4f}")
    return ok, rows


def check_B(steps=1500, lr=0.05):
    scales = torch.linspace(0.05, 3.0, D).sqrt()
    x0 = torch.distributions.Laplace(0.0, 1 / 2 ** 0.5).sample((B, D)) * scales
    x = x0.clone().requires_grad_(True)
    opt = torch.optim.Adam([x], lr=lr)
    hist = []
    for t in range(steps):
        opt.zero_grad()
        loss = sigreg(x)
        loss.backward()
        opt.step()
        if t % 100 == 0 or t == steps - 1:
            with torch.no_grad():
                cov = (x.T @ x / B).double()
                eig = torch.linalg.eigvalsh(cov).clamp_min(1e-12)
                kappa = float(eig[-1] / eig[0])
                z = (x - x.mean(0)) / x.std(0).clamp_min(1e-9)
                kurt = float(((z ** 4).mean() - 3.0))
                hist.append(dict(step=t, loss=float(loss), kappa=kappa,
                                 excess_kurt=kurt, eig_mean=float(eig.mean()),
                                 eig_min=float(eig[0]), eig_max=float(eig[-1])))
                print(f"[B] step {t:>4}: loss={float(loss):.4f} kappa={kappa:.2f} "
                      f"kurt={kurt:.3f} eig in [{float(eig[0]):.3f},{float(eig[-1]):.3f}]")
    first, last = hist[0], hist[-1]
    ok = (last["kappa"] < first["kappa"] / 5 and abs(last["excess_kurt"]) < 0.5
          and 0.5 < last["eig_mean"] < 1.5)
    print(f"[B] SIGReg descent shapes anisotropic Laplace -> isotropic Gaussian "
          f"(kappa {first['kappa']:.1f}->{last['kappa']:.1f}, "
          f"kurt {first['excess_kurt']:.2f}->{last['excess_kurt']:.2f}) -> "
          f"{'PASS' if ok else 'FAIL'}")
    return ok, hist, x0.detach(), x.detach()


def rankme_of(x):
    cov = (x.T @ x / x.shape[0]).double()
    eig = torch.linalg.eigvalsh(cov).clamp_min(0)
    p = eig / eig.sum() + 1e-6
    return float((-(p * p.log()).sum()).exp())


def check_B2(steps=1500, lr=0.05):
    """Rank recovery — the RL-relevant failure mode: descent from a collapsed
    rank-8 cloud must re-expand effective rank."""
    z = torch.randn(B, 8)
    x = (z @ (torch.randn(8, D) / 8 ** 0.5)).clone().requires_grad_(True)
    r0 = rankme_of(x.detach())
    opt = torch.optim.Adam([x], lr=lr)
    for _ in range(steps):
        opt.zero_grad()
        loss = sigreg(x)
        loss.backward()
        opt.step()
    r1 = rankme_of(x.detach())
    ok = r1 > 10 * r0
    print(f"[B2] rank recovery from collapse: RankMe {r0:.1f} -> {r1:.1f} "
          f"(of {D}) -> {'PASS' if ok else 'FAIL'}")
    return ok, r0, r1


def main():
    apply_style()
    okA, rowsA = check_A()
    okB, hist, x0, x1 = check_B()
    okB2, r0, r1 = check_B2()

    fig, axes = plt.subplots(1, 3, figsize=(12, 3.4))
    ax = axes[0]
    names = [r["case"] for r in rowsA]
    vals = [r["sigreg_loss_mean"] for r in rowsA]
    errs = [r["sigreg_loss_std"] for r in rowsA]
    colors = [SERIES[0]] + [SERIES[5]] * (len(names) - 1)
    ax.barh(np.arange(len(names)), vals, xerr=errs, color=colors)
    ax.set_yticks(np.arange(len(names)), names, fontsize=8)
    ax.set_xlabel("SIGReg loss (authors' implementation)")
    ax.set_title("Loss separates isotropic Gaussian from failures")
    ax = axes[1]
    ax.plot([h["step"] for h in hist], [h["kappa"] for h in hist],
            color=SERIES[0], label="κ(Σ)")
    ax.set_yscale("log")
    ax.set_xlabel("descent step")
    ax.set_ylabel("condition number")
    ax2 = ax.twinx()
    ax2.plot([h["step"] for h in hist], [h["excess_kurt"] for h in hist],
             color=SERIES[5], label="excess kurtosis")
    ax2.set_ylabel("excess kurtosis", color=SERIES[5])
    ax2.grid(False)
    ax.set_title("Minimizing SIGReg → isotropic (κ→1) & Gaussian (kurt→0)")
    ax = axes[2]
    with torch.no_grad():
        e0 = torch.linalg.eigvalsh((x0.T @ x0 / B).double()).clamp_min(1e-9)
        e1 = torch.linalg.eigvalsh((x1.T @ x1 / B).double()).clamp_min(1e-9)
    ax.plot(np.sort(e0.numpy())[::-1], color=SERIES[5], label="before")
    ax.plot(np.sort(e1.numpy())[::-1], color=SERIES[0], label="after descent")
    ax.set_yscale("log")
    ax.set_xlabel("eigenvalue index")
    ax.set_ylabel("covariance eigenvalue")
    ax.set_title("Covariance spectrum before/after")
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "claim4_mechanism.png"), bbox_inches="tight")

    with open(os.path.join(OUT, "claim4_results.csv"), "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["section", "key", "value1", "value2"])
        for r in rowsA:
            w.writerow(["A_probe_losses", r["case"], r["sigreg_loss_mean"],
                        r["sigreg_loss_std"]])
        for h in hist:
            w.writerow(["B_descent", h["step"], h["kappa"], h["excess_kurt"]])
        w.writerow(["B2_rank_recovery", "rankme_before_after", r0, r1])

    ok = okA and okB and okB2
    print("CLAIM 4 CPU MECHANISM CHECK:", "PASS" if ok else "FAIL")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()