"""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")))