audit_theorem31.py

"""Numerical audit of Theorem 3.1 (Redundancy-Constrained Information Maximization).

Theorem 3.1 asserts that minimising  L_total = L_rec + lambda * L_div  is *formally
equivalent* to  max_theta  I(X;F') - lambda * TC(F').

The proof (App. A.1) decomposes into two independent claims, which we audit separately
to double precision, each with a control that relaxes the stated condition:

  (A) Sufficiency.  Under a Gaussian decoder p(X|F') = N(D(F'), sigma^2 I),
      minimising the MSE maximises the Barber-Agakov variational lower bound
      I_LB(X;F') = H(X) + E[log p(x|f')]  on I(X;F').
      -> check  I_LB  is an exact affine, strictly decreasing function of the MSE,
         and that I_LB <= I(X;F') for a case where the true MI is known in closed form.
      CONTROL: a Laplace decoder, where the affine MSE<->ELBO identity must fail.

  (B) Independence.  For Gaussian F',
      TC(F') = 1/2 sum_i log(Sigma_ii) - 1/2 log det(Sigma).
      The paper's L_div = -1/2 log det(I + d/(k eps^2) F'^T F') is claimed to be
      "geometrically equivalent" to minimising TC.
      -> check the Gaussian TC identity itself to double precision (exact),
         and then measure how tightly L_div tracks TC.
      CONTROL: non-Gaussian (uniform / heavy-tailed) latents, where the Gaussian
      entropy formula and hence the TC identity must break.

This is a numerical audit, not a proof replacement; it needs no GPU.
"""

import json

import numpy as np

RNG = np.random.default_rng(0)
LOG2PI = np.log(2.0 * np.pi)
report = {}


# ---------------------------------------------------------------- (A) sufficiency
def gaussian_elbo(x, xhat, sigma2):
    """E[log p(x|f')] for p = N(xhat, sigma^2 I), in nats, per sample."""
    d = x.shape[1]
    sq = ((x - xhat) ** 2).sum(1)
    return float((-0.5 * d * (LOG2PI + np.log(sigma2)) - sq / (2 * sigma2)).mean())


def part_a():
    n, d, sigma2 = 20000, 6, 0.7
    x = RNG.normal(size=(n, d))
    out = {}

    # A1: log-likelihood is exactly affine in the MSE, with negative slope
    mses, elbos = [], []
    for scale in np.linspace(0.0, 2.0, 25):
        xhat = x + scale * RNG.normal(size=(n, d))
        mse = float(((x - xhat) ** 2).mean())          # per-element MSE = L_rec
        mses.append(mse)
        elbos.append(gaussian_elbo(x, xhat, sigma2))
    mses, elbos = np.array(mses), np.array(elbos)
    slope = -d / (2 * sigma2) / d                       # d(ELBO)/d(mse_per_element)
    pred = elbos[0] + slope * d * (mses - mses[0])
    out["affine_max_abs_dev_nats"] = float(np.abs(pred - elbos).max())
    out["slope_sign_negative"] = bool(slope < 0)
    out["spearman_mse_vs_elbo"] = float(
        np.corrcoef(np.argsort(np.argsort(mses)), np.argsort(np.argsort(elbos)))[0, 1])

    # A2: I_LB <= I(X;F') on a jointly Gaussian pair with closed-form MI.
    # X = F' + noise, F' ~ N(0, I), noise ~ N(0, s^2 I): I = d/2 log(1 + 1/s^2).
    s2 = 0.5
    f = RNG.normal(size=(n, d))
    xg = f + np.sqrt(s2) * RNG.normal(size=(n, d))
    true_mi = 0.5 * d * np.log(1 + 1.0 / s2)
    hx = 0.5 * d * (LOG2PI + np.log(1 + s2) + 1)        # differential entropy of X
    # optimal Gaussian decoder E[X|F'] = F', residual variance s2
    i_lb = hx + gaussian_elbo(xg, f, s2)
    out["true_MI_nats"] = float(true_mi)
    out["I_LB_at_optimal_decoder_nats"] = float(i_lb)
    out["bound_holds_I_LB_le_MI"] = bool(i_lb <= true_mi + 1e-9)
    out["bound_gap_nats"] = float(true_mi - i_lb)
    # a deliberately worse decoder must give a looser bound
    i_lb_bad = hx + gaussian_elbo(xg, 0.5 * f, s2)
    out["I_LB_suboptimal_decoder_nats"] = float(i_lb_bad)
    out["worse_decoder_is_looser"] = bool(i_lb_bad < i_lb)

    # CONTROL: Laplace decoder -> the MSE<->log-likelihood identity must break
    b = 0.6
    lap_ll, lap_mse = [], []
    for scale in np.linspace(0.05, 2.0, 25):
        xhat = x + scale * RNG.normal(size=(n, d))
        lap_mse.append(float(((x - xhat) ** 2).mean()))
        lap_ll.append(float((-d * np.log(2 * b) - np.abs(x - xhat).sum(1) / b).mean()))
    lap_mse, lap_ll = np.array(lap_mse), np.array(lap_ll)
    A = np.vstack([lap_mse, np.ones_like(lap_mse)]).T
    coef, *_ = np.linalg.lstsq(A, lap_ll, rcond=None)
    out["control_laplace_affine_max_abs_dev_nats"] = float(
        np.abs(A @ coef - lap_ll).max())
    return out


# ------------------------------------------------------------- (B) independence
def gaussian_tc_closed_form(cov):
    return float(0.5 * np.log(np.diag(cov)).sum() - 0.5 * np.linalg.slogdet(cov)[1])


def gaussian_tc_from_entropies(cov):
    k = cov.shape[0]
    h_marg = sum(0.5 * (LOG2PI + 1 + np.log(cov[i, i])) for i in range(k))
    h_joint = 0.5 * (k * (LOG2PI + 1) + np.linalg.slogdet(cov)[1])
    return float(h_marg - h_joint)


def tcr_loss(F, eps=0.5):
    """L_div of Eq. 3 for a single sample's token matrix F in R^{k x d}."""
    k, d = F.shape
    Fn = F / np.linalg.norm(F, axis=1, keepdims=True)
    M = np.eye(k) + (d / (k * eps ** 2)) * (Fn @ Fn.T)
    return float(-0.5 * np.linalg.slogdet(M)[1])


def random_cov(k, rho):
    """Equicorrelated covariance with unit marginals: TC grows monotonically with rho."""
    return (1 - rho) * np.eye(k) + rho * np.ones((k, k))


def part_b():
    out = {}
    k = 8

    # B1: the two Gaussian TC expressions (sum H(f_i) - H(F) vs Eq. 11) agree exactly
    devs = []
    for rho in np.linspace(0.0, 0.95, 40):
        cov = random_cov(k, rho)
        devs.append(abs(gaussian_tc_closed_form(cov) - gaussian_tc_from_entropies(cov)))
    out["TC_identity_max_abs_dev_nats"] = float(max(devs))

    # B2: TC = 0 iff the covariance is diagonal (Lemma 3.2 direction)
    out["TC_at_rho_0"] = gaussian_tc_closed_form(random_cov(k, 0.0))
    out["TC_at_rho_0.9"] = gaussian_tc_closed_form(random_cov(k, 0.9))
    out["TC_monotone_in_rho"] = bool(
        np.all(np.diff([gaussian_tc_closed_form(random_cov(k, r))
                        for r in np.linspace(0, 0.95, 40)]) > 0))

    # B3: does the paper's L_div actually track TC?  Sample token sets with a
    # controlled correlation and compare the *ordering* and correlation of
    # L_div against the true Gaussian TC of the token Gram/covariance.
    d = 128
    tcs, divs = [], []
    for rho in np.linspace(0.0, 0.95, 40):
        cov = random_cov(k, rho)
        L = np.linalg.cholesky(cov)
        F = L @ RNG.normal(size=(k, d))          # k tokens, correlation rho
        emp = np.cov(F)                          # empirical k x k covariance
        tcs.append(gaussian_tc_closed_form(emp))
        divs.append(tcr_loss(F))
    tcs, divs = np.array(tcs), np.array(divs)
    out["pearson_Ldiv_vs_TC"] = float(np.corrcoef(divs, tcs)[0, 1])
    # is the relation an *exact* identity, or only a monotone surrogate?
    A = np.vstack([tcs, np.ones_like(tcs)]).T
    coef, *_ = np.linalg.lstsq(A, divs, rcond=None)
    out["affine_fit_Ldiv_vs_TC_slope"] = float(coef[0])
    out["affine_fit_Ldiv_vs_TC_max_abs_resid_nats"] = float(np.abs(A @ coef - divs).max())
    out["Ldiv_range_nats"] = float(divs.max() - divs.min())
    out["spearman_Ldiv_vs_TC"] = float(np.corrcoef(
        np.argsort(np.argsort(divs)), np.argsort(np.argsort(tcs)))[0, 1])
    out["Ldiv_at_rho_0"] = float(divs[0])
    out["Ldiv_at_rho_0.95"] = float(divs[-1])
    out["Ldiv_increases_with_redundancy"] = bool(divs[-1] > divs[0])
    # L_div is minimised by orthogonal tokens (Lemma 3.2)
    Q = np.linalg.qr(RNG.normal(size=(d, k)))[0].T      # k orthonormal rows
    out["Ldiv_orthogonal_tokens"] = tcr_loss(Q)
    dup = np.repeat(Q[:1], k, axis=0) + 1e-6 * RNG.normal(size=(k, d))
    out["Ldiv_collapsed_tokens"] = tcr_loss(dup)
    out["orthogonal_is_the_minimiser"] = bool(
        out["Ldiv_orthogonal_tokens"] < out["Ldiv_collapsed_tokens"])
    rand_devs = [tcr_loss(RNG.normal(size=(k, d))) for _ in range(200)]
    out["Ldiv_random_tokens_mean"] = float(np.mean(rand_devs))
    out["orthogonal_beats_random"] = bool(
        out["Ldiv_orthogonal_tokens"] <= min(rand_devs))

    # B4: Sylvester check - the d x d form in Eq. 3 and the k x k Gram form agree
    F = RNG.normal(size=(k, d))
    Fn = F / np.linalg.norm(F, axis=1, keepdims=True)
    c = d / (k * 0.5 ** 2)
    lhs = np.linalg.slogdet(np.eye(d) + c * (Fn.T @ Fn))[1]
    rhs = np.linalg.slogdet(np.eye(k) + c * (Fn @ Fn.T))[1]
    out["sylvester_abs_dev"] = float(abs(lhs - rhs))

    # CONTROL: non-Gaussian latents -> the Gaussian entropy formula (and hence the
    # claimed TC identity) is no longer valid.  We compare the Gaussian-formula TC
    # against a k-NN (Kozachenko-Leonenko) estimate of the true TC.
    def kl_entropy(samples):
        from scipy.spatial import cKDTree
        from scipy.special import digamma, gammaln
        n, dd = samples.shape
        tree = cKDTree(samples)
        eps_ = tree.query(samples, k=4)[0][:, 3]
        eps_ = np.maximum(eps_, 1e-12)
        log_vol = dd / 2 * np.log(np.pi) - gammaln(dd / 2 + 1)
        return float(digamma(n) - digamma(3) + log_vol + dd * np.mean(np.log(eps_)))

    def true_tc(samples):
        h_joint = kl_entropy(samples)
        h_marg = sum(kl_entropy(samples[:, [i]]) for i in range(samples.shape[1]))
        return h_marg - h_joint

    n, kk, rho = 40000, 3, 0.7
    cov = random_cov(kk, rho)
    L = np.linalg.cholesky(cov)
    zg = RNG.normal(size=(n, kk)) @ L.T
    # matched-covariance uniform (platykurtic) and Student-t (heavy-tailed) latents
    u = (RNG.uniform(-np.sqrt(3), np.sqrt(3), size=(n, kk))) @ L.T
    t = (RNG.standard_t(3, size=(n, kk)) / np.sqrt(3.0)) @ L.T
    for nm, z in (("gaussian", zg), ("uniform", u), ("student_t3", t)):
        emp = np.cov(z.T)
        out[f"control_{nm}_TC_gauss_formula"] = gaussian_tc_closed_form(emp)
        out[f"control_{nm}_TC_knn_estimate"] = float(true_tc(z))
        out[f"control_{nm}_abs_error"] = abs(
            out[f"control_{nm}_TC_gauss_formula"] - out[f"control_{nm}_TC_knn_estimate"])
    return out


if __name__ == "__main__":
    report["A_sufficiency"] = part_a()
    report["B_independence"] = part_b()
    print(json.dumps(report, indent=2))
    with open("results/theorem31_audit.json", "w") as fh:
        json.dump(report, fh, indent=2)

    a, b = report["A_sufficiency"], report["B_independence"]
    print("\n---- verdict ----")
    print(f"A. MSE <-> Gaussian ELBO affine identity: max dev "
          f"{a['affine_max_abs_dev_nats']:.3e} nats  "
          f"(Laplace control: {a['control_laplace_affine_max_abs_dev_nats']:.3f} nats)")
    print(f"A. I_LB <= I(X;F'): {a['bound_holds_I_LB_le_MI']} "
          f"(gap {a['bound_gap_nats']:.3e} nats at the optimal decoder)")
    print(f"B. Gaussian TC identity: max dev {b['TC_identity_max_abs_dev_nats']:.3e} nats")
    print(f"B. Sylvester d x d == k x k: {b['sylvester_abs_dev']:.3e}")
    print(f"B. corr(L_div, TC) = {b['pearson_Ldiv_vs_TC']:.4f} "
          f"(Spearman {b['spearman_Ldiv_vs_TC']:.4f})")
    print(f"B. L_div(orthogonal) = {b['Ldiv_orthogonal_tokens']:.3f} < "
          f"L_div(collapsed) = {b['Ldiv_collapsed_tokens']:.3f}")
    print("B. control (Gaussian formula vs kNN TC): " + ", ".join(
        f"{nm} err {b[f'control_{nm}_abs_error']:.3f}"
        for nm in ("gaussian", "uniform", "student_t3")))