reproduce_mfvi_variance.py

#!/usr/bin/env python3
"""Exact identities and randomized BLR checks for MFVI predictive-variance overestimation."""

from __future__ import annotations

import hashlib
import json
import warnings
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize


ROOT = Path(__file__).resolve().parent
OUT = ROOT / "outputs"
OUT.mkdir(exist_ok=True)
warnings.filterwarnings("ignore", message=r".*encountered in matmul", category=RuntimeWarning)


def random_rotation(rng, dimension):
    q, r = np.linalg.qr(rng.normal(size=(dimension, dimension)))
    q *= np.sign(np.diag(r))[None, :]
    return q


def one_problem(dimension, seed):
    rng = np.random.default_rng(2500000 + 1000 * dimension + seed)
    samples = max(8 * dimension, 160)
    condition = 10.0 ** rng.uniform(0.5, 4.0)
    rotation = random_rotation(rng, dimension)
    data_covariance = rotation @ np.diag(np.geomspace(condition, 1.0, dimension)) @ rotation.T
    x = rng.multivariate_normal(np.zeros(dimension), data_covariance, size=samples)
    x -= x.mean(axis=0)
    alpha = 10.0 ** rng.uniform(-3.0, 3.0)
    noise_variance = 10.0 ** rng.uniform(-3.0, 3.0)
    gram = x.T @ x
    precision = gram / noise_variance + alpha * np.eye(dimension)
    posterior = np.linalg.inv(precision)
    mfvi = np.diag(1.0 / np.diag(precision))

    eigenvalues, eigenvectors = np.linalg.eigh(posterior)
    v_min, v_max = eigenvectors[:, 0], eigenvectors[:, -1]
    trace_margin = float(np.trace(posterior) - np.trace(mfvi))
    min_direction_margin = float(v_min @ (mfvi - posterior) @ v_min)
    max_direction_control = float(v_max @ (mfvi - posterior) @ v_max)

    empirical_data_covariance = gram / samples
    training_distribution_margin = float(np.trace(empirical_data_covariance @ (posterior - mfvi)))
    identity_rhs = float(-noise_variance * alpha / samples * trace_margin)

    gram_eigenvalues, gram_eigenvectors = np.linalg.eigh(gram)
    first_pc = gram_eigenvectors[:, -1]
    last_pc = gram_eigenvectors[:, 0]
    first_pc_margin = float(first_pc @ (mfvi - posterior) @ first_pc)
    last_pc_control = float(last_pc @ (mfvi - posterior) @ last_pc)
    alignment = float(abs(first_pc @ v_min))

    competing_scale = float(np.linalg.eigvalsh(gram / noise_variance).max() + alpha)
    nonspherical_prior = np.diag(np.geomspace(competing_scale * 1e3, competing_scale * 1e6, dimension))
    posterior_ns = np.linalg.inv(gram / noise_variance + nonspherical_prior)
    _, eigenvectors_ns = np.linalg.eigh(posterior_ns)
    nonspherical_alignment = float(abs(first_pc @ eigenvectors_ns[:, 0]))

    # Reverse-KL numerical recovery of the diagonal optimum for a subset of problems.
    optimizer_relative_error = None
    if dimension <= 8 and seed < 6:
        def objective(log_variance):
            variance = np.exp(log_variance)
            return 0.5 * (float(np.diag(precision) @ variance) - float(np.sum(log_variance)))

        result = minimize(objective, np.zeros(dimension), method="BFGS", options={"gtol": 1e-11, "maxiter": 1000})
        recovered = np.exp(result.x)
        optimizer_relative_error = float(np.linalg.norm(recovered - np.diag(mfvi)) / np.linalg.norm(np.diag(mfvi)))

    return {
        "dimension": dimension,
        "seed": seed,
        "samples": samples,
        "data_covariance_condition": condition,
        "alpha": alpha,
        "noise_variance": noise_variance,
        "posterior_condition": float(np.linalg.cond(posterior)),
        "trace_exact_posterior": float(np.trace(posterior)),
        "trace_mfvi": float(np.trace(mfvi)),
        "trace_margin": trace_margin,
        "minimum_posterior_direction_margin": min_direction_margin,
        "maximum_posterior_direction_control": max_direction_control,
        "training_distribution_expected_exact_minus_mfvi": training_distribution_margin,
        "training_distribution_identity_rhs": identity_rhs,
        "identity_absolute_error": abs(training_distribution_margin - identity_rhs),
        "first_pc_min_posterior_alignment": alignment,
        "first_pc_mfvi_minus_exact_margin": first_pc_margin,
        "last_pc_control_margin": last_pc_control,
        "nonspherical_prior_first_pc_alignment_control": nonspherical_alignment,
        "kl_optimizer_relative_error": optimizer_relative_error,
        "gates": {
            "posterior_trace_exceeds_mfvi": trace_margin >= -1e-10,
            "mfvi_overestimates_min_direction": min_direction_margin >= -1e-10,
            "max_direction_mutation_reverses": max_direction_control < 0,
            "training_distribution_mfvi_exceeds_exact": training_distribution_margin <= 1e-10,
            "training_identity": abs(training_distribution_margin - identity_rhs) <= 1e-8 * max(1.0, abs(identity_rhs)),
            "first_pc_is_min_posterior_direction": alignment >= 1 - 1e-8,
            "first_pc_overestimate": first_pc_margin >= -1e-10,
            "last_pc_mutation_reverses": last_pc_control < 0,
            "nonspherical_prior_breaks_alignment": nonspherical_alignment < 0.999,
        },
    }


def make_figures(records):
    fig, axes = plt.subplots(1, 3, figsize=(14, 4.3))
    dims = np.array([r["dimension"] for r in records])
    axes[0].scatter(dims - 0.08, [r["minimum_posterior_direction_margin"] for r in records], s=14, alpha=0.6, label="min-eigenvector")
    axes[0].scatter(dims + 0.08, [r["maximum_posterior_direction_control"] for r in records], s=14, alpha=0.6, label="max control")
    axes[0].axhline(0, color="black", linestyle="--")
    axes[0].set(xlabel="dimension", ylabel="MFVI minus exact variance", title="Directional variance inversion")
    axes[0].set_yscale("symlog", linthresh=1e-10)
    axes[0].legend()
    x = np.array([r["training_distribution_identity_rhs"] for r in records])
    y = np.array([r["training_distribution_expected_exact_minus_mfvi"] for r in records])
    axes[1].scatter(x, y, s=16, alpha=0.65)
    lo, hi = min(x.min(), y.min()), max(x.max(), y.max())
    axes[1].plot([lo, hi], [lo, hi], "k--")
    axes[1].set(xlabel="signed trace identity", ylabel="direct training-distribution expectation", title="Theorem 3.7 identity")
    axes[1].set_xscale("symlog", linthresh=1e-12)
    axes[1].set_yscale("symlog", linthresh=1e-12)
    axes[2].scatter(dims - 0.08, [r["first_pc_mfvi_minus_exact_margin"] for r in records], s=14, alpha=0.6, label="first PC")
    axes[2].scatter(dims + 0.08, [r["last_pc_control_margin"] for r in records], s=14, alpha=0.6, label="last PC control")
    axes[2].axhline(0, color="black", linestyle="--")
    axes[2].set(xlabel="dimension", ylabel="MFVI minus exact variance", title="Overestimation follows data concentration")
    axes[2].set_yscale("symlog", linthresh=1e-10)
    axes[2].legend()
    fig.tight_layout()
    fig.savefig(OUT / "mfvi_variance_evidence.png", dpi=180)
    plt.close(fig)


def main():
    records = [one_problem(d, seed) for d in (2, 3, 5, 8, 16, 32, 64) for seed in range(20)]
    gate_names = list(records[0]["gates"])
    gate_counts = {name: sum(r["gates"][name] for r in records) for name in gate_names}
    optimizer_errors = [r["kl_optimizer_relative_error"] for r in records if r["kl_optimizer_relative_error"] is not None]
    summary = {
        "problems": len(records),
        "dimensions": [2, 3, 5, 8, 16, 32, 64],
        "alpha_range": [min(r["alpha"] for r in records), max(r["alpha"] for r in records)],
        "noise_variance_range": [min(r["noise_variance"] for r in records), max(r["noise_variance"] for r in records)],
        "maximum_posterior_condition": max(r["posterior_condition"] for r in records),
        "gate_counts": gate_counts,
        "claim_1": {
            "minimum_trace_margin": min(r["trace_margin"] for r in records),
            "minimum_overestimate_margin": min(r["minimum_posterior_direction_margin"] for r in records),
            "maximum_direction_mutation_failures": sum(r["maximum_posterior_direction_control"] < 0 for r in records),
            "kl_optimizer_checks": len(optimizer_errors),
            "kl_optimizer_worst_relative_error": max(optimizer_errors),
        },
        "claim_2": {
            "largest_expected_exact_minus_mfvi": max(r["training_distribution_expected_exact_minus_mfvi"] for r in records),
            "worst_identity_absolute_error": max(r["identity_absolute_error"] for r in records),
            "all_training_distribution_expectations_nonpositive": all(r["training_distribution_expected_exact_minus_mfvi"] <= 1e-10 for r in records),
        },
        "claim_3": {
            "worst_first_pc_alignment_error": max(1 - r["first_pc_min_posterior_alignment"] for r in records),
            "minimum_first_pc_overestimate": min(r["first_pc_mfvi_minus_exact_margin"] for r in records),
            "last_pc_mutation_failures": sum(r["last_pc_control_margin"] < 0 for r in records),
            "nonspherical_prior_alignment_breaks": sum(r["nonspherical_prior_first_pc_alignment_control"] < 0.999 for r in records),
        },
        "general_derivation": {
            "mfvi_optimum": "S*=diag(1/diag(Sigma^-1))",
            "trace_proof": "For A=Sigma^-1, A_ii <= lambda_max(A), and Schur complements give Sigma_ii >= 1/A_ii; summing yields Tr(Sigma)>=Tr(S*).",
            "direction_proof": "min_i 1/A_ii >= 1/lambda_max(A)=lambda_min(Sigma), hence S* dominates Sigma along Sigma's minimum-eigenvalue direction.",
            "training_distribution_identity": "tr[(X'X/N)(Sigma-S*)]=-(sigma^2 alpha/N)[Tr(Sigma)-Tr(S*)] <= 0.",
            "first_pc_proof": "Sigma^-1=X'X/sigma^2+alpha I commutes with X'X; its largest data eigenvalue maps to Sigma's smallest eigenvalue.",
        },
    }
    results = {
        "paper": {
            "title": "Gaussian Mean Field Variational Inference can Overestimate Predictive Variance",
            "openreview": "RG7maF4bGu",
        },
        "protocol": {"precision": "float64", "device": "CPU", "external_compute_cost": 0},
        "summary": summary,
        "records": records,
    }
    path = OUT / "results.json"
    path.write_text(json.dumps(results, indent=2))
    make_figures(records)
    lines = []
    for item in sorted(OUT.glob("*")):
        if item.is_file() and item.name != "SHA256SUMS.txt":
            lines.append(f"{hashlib.sha256(item.read_bytes()).hexdigest()}  {item.name}")
    (OUT / "SHA256SUMS.txt").write_text("\n".join(lines) + "\n")
    print(json.dumps(summary, indent=2))
    print(f"RESULTS_SHA256={hashlib.sha256(path.read_bytes()).hexdigest()}")


if __name__ == "__main__":
    main()