MCMC diagnostics + ensemble uncertainty: plots, EnsembleEnergy, disagreement OOD

CHANGELOG.md

@@ -5,6 +5,14 @@
 Internal cleanup plus a diffusion track, exact-likelihood flows, and
 conditional-generation / calibration tools.
 
+- **Diagnostics & ensemble uncertainty.** `viz.autocorrelation_plot` /
+  `rank_plot` / `trace_plot` and a numeric `eval.autocorrelation` for judging MCMC
+  mixing; `EnsembleEnergy` (a deep-ensemble *mean* energy — the geometric-mean
+  density, distinct from `MixtureEnergy`'s logsumexp) with `member_energies` /
+  `disagreement`, plus `eval.ensemble_disagreement` as an epistemic OOD score.
+  Validated: the ACF estimator matches an AR(1)'s `ρ^t`; the ensemble reproduces
+  the closed-form combined Gaussian; disagreement separates in-distribution from
+  OOD at AUROC ≈ 1. Example `ensemble_ood.py`.
 - **`nets.NeuralSplineCouplingFlow`** — a rational-quadratic neural spline flow
   (Durkan et al. 2019): monotonic-spline coupling layers dropped into the
   affine-flow scaffold, strictly more expressive per layer, so it fits sharp

README.md

@@ -47,10 +47,10 @@ returns `LossOutput(loss, metrics, x_neg)`; call `out.loss.backward()`).
 | **Energies** | any callable `(B, *shape) -> (B,)`; `nets.MLPEnergy` / `ConvEnergy` / `ResNetEnergy` / `ConvClassifier` (SiLU, optional spectral norm, no batch norm), `nets.RBM` (Bernoulli RBM with exact `log_z`), `GuidedEnergy` (classifier-free guidance), `nets.IsingEnergy` / `PottsEnergy` (discrete lattices), `nets.FunnelEnergy` / `GaussianMixtureEnergy` / `BananaEnergy` (closed-form targets), `nets.AffineCouplingFlow` (RealNVP) / `nets.NeuralSplineCouplingFlow` (rational-quadratic spline) / `nets.ContinuousNormalizingFlow` (FFJORD — trainable exact-likelihood flows / self-normalized energies), noise-conditional variants for NCSN; `EnergyModel`, `ebm.score` |
 | **Samplers** | `LangevinDynamics` (ULA/SGLD), `MALA`, `AdaptiveMALA` (dual-averaging step-size warmup + diagonal metric), `HMC`, `UnderdampedLangevin` (SGHMC), `PreconditionedLangevin`, `ParallelTempering` (replica exchange), `TemperedTransitions`, `SVGD` (Stein variational), `GibbsSampler` (block Gibbs), `GibbsWithGradients` + `CategoricalGibbsWithGradients`, `AnnealedLangevinDynamics`, `ProbabilityFlowODE` / `PredictorCorrector` (score-SDE), `DDPMAncestralSampler` (VP diffusion) |
 | **Losses** | `ContrastiveDivergence` (CD-k / persistent CD), `DiffusionRecoveryLikelihood` + `drl_sample`, `DenoisingScoreMatching` / `MultiSigmaDenoisingScoreMatching` (NCSN), `VPDenoisingScoreMatching` (DDPM), `SlicedScoreMatching`, `ExactScoreMatching`, `EnergyDiscrepancy` (MCMC-free), `PseudoLikelihood` / `RatioMatching` / `ConcreteScoreMatching` (MCMC-free, discrete), `NoiseContrastiveEstimation`, `JEMLoss` |
-| **Composition** | `SumEnergy` (product of experts), `MixtureEnergy`, `TemperedEnergy` — energies compose like densities and nest |
+| **Composition** | `SumEnergy` (product of experts), `MixtureEnergy`, `EnsembleEnergy` (deep-ensemble mean energy + member disagreement), `TemperedEnergy` — energies compose like densities and nest |
 | **Training** | thin `Trainer` (device, EMA, supervised batches, `save`/`load` checkpointing), `ReplayBuffer`, `EMA` |
-| **Eval** | `ais_log_z` / `reverse_ais_log_z` (bracket `log Z`), `pf_ode_log_likelihood` (exact likelihood via the probability-flow ODE), `bits_per_dim`, `frechet_distance` (FID), `mmd`, `precision_recall`, `inception_score`, `kernel_stein_discrepancy` / `classifier_two_sample_test` / `fisher_divergence` (goodness-of-fit), `mutual_information` (MINE), `expected_calibration_error` / `reliability_curve` / `temperature_scale` (calibration), `ood_auroc`, `effective_sample_size` / `split_rhat` (MCMC diagnostics) |
-| **Data & viz** | 2D toys (`two_moons`, `eight_gaussians`, `checkerboard`, `rings`, `spirals`) and torchvision-free image loaders (`mnist`, `fashion_mnist`, `cifar10`, `cifar100`); `viz.energy_contour` / `plot_samples` / `energy_histogram` / `show_images` |
+| **Eval** | `ais_log_z` / `reverse_ais_log_z` (bracket `log Z`), `pf_ode_log_likelihood` (exact likelihood via the probability-flow ODE), `bits_per_dim`, `frechet_distance` (FID), `mmd`, `precision_recall`, `inception_score`, `kernel_stein_discrepancy` / `classifier_two_sample_test` / `fisher_divergence` (goodness-of-fit), `mutual_information` (MINE), `expected_calibration_error` / `reliability_curve` / `temperature_scale` (calibration), `ood_auroc`, `ensemble_disagreement` (epistemic OOD), `effective_sample_size` / `split_rhat` / `autocorrelation` (MCMC diagnostics) |
+| **Data & viz** | 2D toys (`two_moons`, `eight_gaussians`, `checkerboard`, `rings`, `spirals`) and torchvision-free image loaders (`mnist`, `fashion_mnist`, `cifar10`, `cifar100`); `viz.energy_contour` / `plot_samples` / `energy_histogram` / `show_images` / `autocorrelation_plot` / `rank_plot` / `trace_plot` |
 
 ## Examples
 
@@ -77,6 +77,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/
 - `sampling_hard_targets.py` — parallel tempering escapes a trapped mode; ESS / R̂ diagnostics
 - `adaptive_mala.py` — a self-tuning MALA: dual-averaging step size + a learned diagonal metric
 - `goodness_of_fit.py` — KSD for model selection; classifier two-sample test
+- `ensemble_ood.py` — a deep-ensemble EBM whose member disagreement flags OOD
 - `mine_mutual_information.py` — estimate mutual information with MINE vs the Gaussian closed form
 - `benchmark_samplers.py` — rank samplers on the banana against exact draws (ESS, R̂, MMD)
 - `checkpoint_resume.py` — save a run and resume it in a fresh process

examples/ensemble_ood.py

@@ -0,0 +1,76 @@
+"""Deep-ensemble EBM: member disagreement as an epistemic OOD signal.
+
+A single energy can be confidently wrong off-distribution. An ensemble of energy
+networks trained on the same data agrees where it saw data and *disagrees* where
+it didn't — so the variance of the member energies is an epistemic-uncertainty
+score that flags OOD inputs. `EnsembleEnergy` pools the members into a mean energy
+(a geometric-mean density) for sampling/scoring, and `ensemble_disagreement`
+returns the per-sample variance; here it separates two-moons (in-distribution)
+from a Gaussian blob (OOD) at AUROC ≈ 1.
+
+Run:  python examples/ensemble_ood.py
+Outputs ensemble_ood_result.png next to this script (needs the [viz] extra).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+import ebm
+
+
+def _train_member(data: torch.Tensor, seed: int) -> torch.nn.Module:
+    torch.manual_seed(seed)
+    net = ebm.nets.MLPEnergy(dim=2, hidden=(128, 128))
+    sampler = ebm.LangevinDynamics(step_size=0.01, steps=60)
+    loss_fn = ebm.ContrastiveDivergence(
+        sampler, buffer=ebm.ReplayBuffer(8192, (2,)), energy_reg=0.1
+    )
+    trainer = ebm.Trainer(net, loss_fn, lr=1e-3, ema_decay=0.999, device="cpu")
+    trainer.fit(data, steps=3000, batch_size=256)
+    return net
+
+
+def main() -> None:
+    torch.manual_seed(0)
+    data = ebm.datasets.two_moons(8192)
+    ensemble = ebm.EnsembleEnergy(*[_train_member(data, seed) for seed in range(3)])
+
+    x_in = ebm.datasets.two_moons(2000)
+    x_out = torch.randn(2000, 2) * 0.6 + torch.tensor([0.0, 3.0])  # off-manifold blob
+    d_in = ebm.eval.ensemble_disagreement(ensemble, x_in)
+    d_out = ebm.eval.ensemble_disagreement(ensemble, x_out)
+    auroc = ebm.eval.ood_auroc(lambda z: ebm.eval.ensemble_disagreement(ensemble, z), x_in, x_out)
+    print(f"mean disagreement  in-dist {d_in.mean():.3f}   OOD {d_out.mean():.3f}")
+    print(f"OOD AUROC (disagreement) = {auroc:.3f}")
+
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
+    grid = torch.linspace(-3, 4, 200)
+    gy, gx = torch.meshgrid(grid, grid, indexing="ij")
+    pts = torch.stack([gx.reshape(-1), gy.reshape(-1)], dim=1)
+    dis = ebm.eval.ensemble_disagreement(ensemble, pts).reshape(200, 200)
+    im = axes[0].imshow(dis, extent=(-3, 4, -3, 4), origin="lower", vmax=float(dis.quantile(0.98)))
+    axes[0].scatter(x_in[:, 0], x_in[:, 1], s=3, color="white", alpha=0.3)
+    axes[0].set_title("member disagreement (bright = uncertain / OOD)")
+    fig.colorbar(im, ax=axes[0], fraction=0.046)
+
+    axes[1].hist(d_in.numpy(), bins=50, alpha=0.6, label="two-moons (in)", density=True)
+    axes[1].hist(d_out.numpy(), bins=50, alpha=0.6, label="Gaussian blob (OOD)", density=True)
+    axes[1].set_xlabel("ensemble disagreement")
+    axes[1].set_title(f"OOD AUROC = {auroc:.3f}")
+    axes[1].legend()
+
+    out = Path(__file__).parent / "ensemble_ood_result.png"
+    fig.savefig(out, dpi=120, bbox_inches="tight")
+    print(f"saved {out}")
+
+
+if __name__ == "__main__":
+    main()

src/ebm/__init__.py

@@ -7,7 +7,7 @@
 from ebm import datasets, eval, nets, viz
 from ebm.ais import AISResult, ais_log_z, log_likelihood, reverse_ais_log_z
 from ebm.buffer import ReplayBuffer
-from ebm.compose import MixtureEnergy, SumEnergy, TemperedEnergy
+from ebm.compose import EnsembleEnergy, MixtureEnergy, SumEnergy, TemperedEnergy
 from ebm.diffusion import DDPMAncestralSampler, VPDenoisingScoreMatching, VPSchedule
 from ebm.energy import ConditionalEnergyFn, EnergyFn, EnergyModel, score
 from ebm.jem import ClassifierEnergy, ConditionalEnergy, GuidedEnergy, JEMLoss
@@ -70,6 +70,7 @@
     "EnergyDiscrepancy",
     "EnergyFn",
     "EnergyModel",
+    "EnsembleEnergy",
     "ExactScoreMatching",
     "GibbsSampler",
     "GibbsWithGradients",

src/ebm/compose.py

@@ -82,6 +82,36 @@ def forward(self, x: Tensor) -> Tensor:
         return -torch.logsumexp(stacked, dim=0)
 
 
+class EnsembleEnergy(_Composite):
+    """Deep ensemble as a **mean energy**: ``E(x) = (1/N) Σ_i E_i(x)``.
+
+    Averaging energies is the *geometric* mean of the members' densities
+    (``p ∝ Π p_i^{1/N}``) — the variance-reduced log-density estimate you want
+    from a deep ensemble, and equivalent to ``SumEnergy`` with weights ``1/N``.
+    This is deliberately distinct from `MixtureEnergy`, which mixes densities
+    with a ``logsumexp`` (an *arithmetic* mean).
+
+    Beyond the pooled energy, the spread of the members carries epistemic
+    uncertainty: `member_energies(x)` returns the ``(B, N)`` per-member energies
+    and `disagreement(x)` their per-sample variance — low where the members agree
+    (seen data), high off-distribution. See `eval.ensemble_disagreement`.
+    """
+
+    def __init__(self, *energies: EnergyFn):
+        super().__init__(energies)
+
+    def member_energies(self, x: Tensor) -> Tensor:
+        """Per-member energies, stacked as ``(B, N)``."""
+        return torch.stack([e(x) for e in self.energies], dim=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.member_energies(x).mean(dim=1)
+
+    def disagreement(self, x: Tensor) -> Tensor:
+        """Per-sample variance across members ``(B,)`` — an epistemic OOD signal."""
+        return self.member_energies(x).var(dim=1, unbiased=False)
+
+
 class TemperedEnergy(nn.Module):
     """``E(x) / T``: temperature ``T > 1`` flattens ``p``, ``T < 1`` sharpens it.
 

src/ebm/eval/__init__.py

@@ -12,8 +12,10 @@
     temperature_scale,
 )
 from ebm.eval.diagnostics import (
+    autocorrelation,
     effective_sample_size,
     energies,
+    ensemble_disagreement,
     ood_auroc,
     split_rhat,
 )
@@ -34,10 +36,12 @@
 __all__ = [
     "AISResult",
     "ais_log_z",
+    "autocorrelation",
     "bits_per_dim",
     "classifier_two_sample_test",
     "effective_sample_size",
     "energies",
+    "ensemble_disagreement",
     "expected_calibration_error",
     "fisher_divergence",
     "frechet_distance",

src/ebm/eval/diagnostics.py

@@ -2,11 +2,16 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
 from torch import Tensor
 
 from ebm.energy import EnergyFn
 
+if TYPE_CHECKING:
+    from ebm.compose import EnsembleEnergy
+
 
 def _as_chains(samples: Tensor) -> Tensor:
     """Coerce MCMC output to ``(n_chains, n_samples, dim)`` in float64."""
@@ -104,13 +109,49 @@ def effective_sample_size(samples: Tensor) -> Tensor:
     return ess
 
 
+@torch.no_grad()
+def autocorrelation(samples: Tensor, max_lag: int = 40) -> Tensor:
+    """Mean per-chain autocorrelation ``ρ̂_t`` for lags ``0..max_lag``, per dimension.
+
+    Input is ``(n_chains, n_samples[, dim])``. Returns ``(max_lag+1, dim)`` (lag 0
+    is ``1``). Computed by FFT (zero-padded past ``2N``) and averaged over chains —
+    the same autocovariance used by `effective_sample_size`. A chain that mixes
+    well decays toward zero within a few lags; slow decay flags autocorrelation.
+    """
+    x = _as_chains(samples)
+    _, n, _ = x.shape
+    max_lag = min(max_lag, n - 1)
+    centered = x - x.mean(dim=1, keepdim=True)
+    n_fft = 1
+    while n_fft < 2 * n:
+        n_fft <<= 1
+    spec = torch.fft.rfft(centered, n=n_fft, dim=1)
+    acov = torch.fft.irfft(spec.abs().pow(2), n=n_fft, dim=1)[:, :n, :]  # (M, N, d)
+    acf = acov / acov[:, :1, :]  # normalize by lag-0 per chain/dim
+    return acf.mean(dim=0)[: max_lag + 1]
+
+
 @torch.no_grad()
 def energies(energy: EnergyFn, x: Tensor, batch_size: int = 1024) -> Tensor:
     """Energies of ``x`` computed in batches, returned on CPU."""
     out = [energy(chunk).detach().cpu() for chunk in x.split(batch_size)]
     return torch.cat(out)
 
 
+@torch.no_grad()
+def ensemble_disagreement(ensemble: EnsembleEnergy, x: Tensor, batch_size: int = 1024) -> Tensor:
+    """Per-sample ensemble disagreement (variance of member energies), batched.
+
+    The spread of an ensemble's member energies is an *epistemic* uncertainty
+    signal: members trained on the same data agree where they saw data and
+    diverge off-distribution, so `ensemble_disagreement` tends to be much larger
+    on OOD inputs — feed it to `ood_auroc` as a detector. Returns a ``(B,)`` tensor
+    on CPU.
+    """
+    out = [ensemble.disagreement(chunk).detach().cpu() for chunk in x.split(batch_size)]
+    return torch.cat(out)
+
+
 @torch.no_grad()
 def ood_auroc(energy: EnergyFn, x_in: Tensor, x_out: Tensor) -> float:
     """AUROC for separating in-distribution from OOD data by energy.

src/ebm/viz.py

@@ -9,7 +9,15 @@
 
 from ebm.energy import EnergyFn
 
-__all__ = ["energy_contour", "energy_histogram", "plot_samples", "show_images"]
+__all__ = [
+    "autocorrelation_plot",
+    "energy_contour",
+    "energy_histogram",
+    "plot_samples",
+    "rank_plot",
+    "show_images",
+    "trace_plot",
+]
 
 
 def _require_matplotlib():
@@ -114,6 +122,86 @@ def show_images(
     return ax
 
 
+@torch.no_grad()
+def _chain_ranks(samples: Tensor, dim: int = 0) -> Tensor:
+    """Rank each draw among the pooled draws of all chains; returns ``(n_chains, n_samples)``."""
+    from ebm.eval.diagnostics import _as_chains
+
+    x = _as_chains(samples)[..., dim]  # (M, N)
+    m, n = x.shape
+    order = x.reshape(-1).argsort()
+    ranks = torch.empty(m * n, dtype=torch.float64)
+    ranks[order] = torch.arange(1, m * n + 1, dtype=torch.float64)
+    return ranks.reshape(m, n)
+
+
+@torch.no_grad()
+def autocorrelation_plot(samples: Tensor, max_lag: int = 40, dim: int = 0, ax=None):
+    """Autocorrelation ``ρ̂_t`` vs lag for MCMC output, with a ``±1.96/√N`` band. Returns the axes.
+
+    Input is ``(n_chains, n_samples[, dim])``. Bars decaying to inside the band
+    within a few lags indicate good mixing; a slow decay flags autocorrelation.
+    """
+    from ebm.eval.diagnostics import autocorrelation
+
+    plt = _require_matplotlib()
+    acf = autocorrelation(samples, max_lag)[:, dim]
+    n = samples.shape[1]
+    if ax is None:
+        _, ax = plt.subplots(figsize=(6, 4))
+    ax.bar(range(len(acf)), acf.numpy(), width=0.8, color="#5c50c9")
+    band = 1.96 / (n**0.5)
+    for y in (band, -band):
+        ax.axhline(y, ls="--", color="gray", lw=1)
+    ax.axhline(0.0, color="black", lw=0.8)
+    ax.set_xlabel("lag")
+    ax.set_ylabel("autocorrelation")
+    return ax
+
+
+@torch.no_grad()
+def rank_plot(samples: Tensor, dim: int = 0, bins: int = 20, ax=None):
+    """Vehtari rank histogram per chain (uniform under good mixing). Returns the axes.
+
+    Ranks every draw among the pooled draws of all chains, then histograms each
+    chain's ranks. Well-mixed chains give flat, overlapping histograms near the
+    dashed uniform line; a chain offset from the others shows a sloped or skewed
+    histogram. Reference: Vehtari et al. (2021).
+    """
+    plt = _require_matplotlib()
+    ranks = _chain_ranks(samples, dim)
+    m, n = ranks.shape
+    edges = torch.linspace(1, m * n, bins + 1).numpy()
+    if ax is None:
+        _, ax = plt.subplots(figsize=(6, 4))
+    for i in range(m):
+        ax.hist(ranks[i].numpy(), bins=edges, histtype="step", alpha=0.7)
+    ax.axhline(n / bins, ls="--", color="gray", lw=1)  # uniform expectation per bin
+    ax.set_xlabel("rank")
+    ax.set_ylabel("count")
+    return ax
+
+
+@torch.no_grad()
+def trace_plot(samples: Tensor, dim: int = 0, max_chains: int = 8, ax=None):
+    """Trace of each chain's value over iterations (mixing at a glance). Returns the axes.
+
+    Input is ``(n_chains, n_samples[, dim])``. Chains that overlap and wander
+    across the same range have mixed; a chain stuck at a different level has not.
+    """
+    from ebm.eval.diagnostics import _as_chains
+
+    plt = _require_matplotlib()
+    x = _as_chains(samples)[..., dim]
+    if ax is None:
+        _, ax = plt.subplots(figsize=(8, 4))
+    for i in range(min(x.shape[0], max_chains)):
+        ax.plot(x[i].numpy(), alpha=0.7, lw=0.8)
+    ax.set_xlabel("iteration")
+    ax.set_ylabel(f"dim {dim}")
+    return ax
+
+
 def energy_histogram(energy: EnergyFn, batches: dict[str, Tensor], ax=None, bins: int = 60):
     """Overlaid energy histograms, e.g. ``{"data": x, "samples": x_neg}``.
 

tests/test_compose.py

@@ -77,3 +77,32 @@ def test_compositions_nest():
         ebm.SumEnergy(quadratic_energy, ebm.MixtureEnergy(quadratic_energy)), 2.0
     )
     assert nested(torch.randn(8, 2)).shape == (8,)
+
+
+def test_ensemble_energy_is_geometric_mean_gaussian():
+    # Members N(0, σ_i²) → mean energy is N(0, 1/precision), precision = mean(1/σ_i²).
+    sigmas = [1.0, 0.5, 2.0]
+    members = [gaussian_energy(0.0, s**2) for s in sigmas]
+    ens = ebm.EnsembleEnergy(*members)
+    precision = sum(1 / s**2 for s in sigmas) / len(sigmas)
+    samples = ebm.MALA(step_size=0.1, steps=800).sample(ens, torch.randn(6000, 2))
+    assert (samples.var(0) - 1 / precision).abs().max().item() < 0.05
+    # member_energies stacks per-member; forward is their mean.
+    x = torch.randn(16, 2)
+    me = ens.member_energies(x)
+    assert me.shape == (16, 3)
+    assert torch.allclose(ens(x), me.mean(dim=1))
+
+
+def test_ensemble_disagreement_flags_ood():
+    sigmas = [1.0, 0.5, 2.0]
+    ens = ebm.EnsembleEnergy(*[gaussian_energy(0.0, s**2) for s in sigmas])
+    x_in = torch.randn(2000, 2)
+    x_out = 6.0 + torch.randn(2000, 2)
+    d_in = ebm.eval.ensemble_disagreement(ens, x_in)
+    d_out = ebm.eval.ensemble_disagreement(ens, x_out)
+    assert d_in.shape == (2000,)
+    assert d_out.mean() > 50 * d_in.mean()  # members diverge off-distribution
+    # As an OOD score it separates the two sets almost perfectly.
+    auroc = ebm.eval.ood_auroc(lambda z: ebm.eval.ensemble_disagreement(ens, z), x_in, x_out)
+    assert auroc > 0.99