Latent-variable EBM: LatentEBM (joint E(x,z), block-Gibbs sampled)

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Changelog
 
+## 0.16.0 — unreleased
+
+- **`LatentEBM`** — a latent-variable EBM: a joint `E(x, z) = E_prior(z) +
+  E(x | z)` coupling a prior over a latent `z` (default standard normal) with a
+  decoder energy. The data marginal is intractable, so `sample_joint` runs block
+  Gibbs — alternating an MCMC update of `z` under its posterior `E(z | x)` with an
+  update of `x` under `E(x | z)`; `posterior_energy` / `conditional_energy` expose
+  those blocks as plain `EnergyFn`s. Validated on the linear-Gaussian conjugate
+  case: block Gibbs recovers the exact marginal `N(0, WWᵀ + σ²I)` and the Gaussian
+  posterior `z | x`. Example `latent_ebm.py`.
+
 ## 0.15.0 — 2026-08-20
 
 The largest release yet: a variance-preserving diffusion track and trainable

README.md

@@ -47,7 +47,7 @@ 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`, `EnsembleEnergy` (deep-ensemble mean energy + member disagreement), `TemperedEnergy` — energies compose like densities and nest |
+| **Composition** | `SumEnergy` (product of experts), `MixtureEnergy`, `EnsembleEnergy` (deep-ensemble mean energy + member disagreement), `TemperedEnergy`, `LatentEBM` (joint `E(x, z)` with a prior + decoder, block-Gibbs sampled) — 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`, `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` |
@@ -78,6 +78,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/
 - `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
+- `latent_ebm.py` — a latent-variable EBM: block-Gibbs on a joint `E(x, z)` matches ancestral sampling
 - `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/latent_ebm.py

@@ -0,0 +1,83 @@
+"""Latent-variable EBM: a joint E(x, z) sampled by block Gibbs.
+
+`LatentEBM` couples a prior over a latent `z` with a decoder energy `E(x | z)`
+into a joint `p(x, z) ∝ exp(-E(x, z))`. The data marginal `p(x)` is intractable,
+so you sample the joint: block Gibbs alternates an MCMC update of `z` under its
+posterior with an update of `x` under the decoder. Here a fixed nonlinear decoder
+bends a 1-D Gaussian latent into a curved 2-D manifold; because the decoder is
+Gaussian we also have a cheap *ancestral* reference (draw z ~ N(0,1), then
+x = g(z) + σε), and the block-Gibbs marginal should match it — an MCMC-vs-exact
+check on a nontrivial manifold.
+
+Run:  python examples/latent_ebm.py
+Outputs latent_ebm_result.png next to this script (needs the [viz] extra).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import nn
+
+import ebm
+
+SIGMA = 0.15
+
+
+class CurveDecoder(nn.Module):
+    """Fixed 1-D → 2-D generator g(z); E(x | z) = ‖x − g(z)‖² / (2σ²)."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        torch.manual_seed(0)
+        self.g = nn.Sequential(nn.Linear(1, 64), nn.Tanh(), nn.Linear(64, 2))
+        for p in self.g.parameters():
+            p.requires_grad_(False)
+
+    def mean(self, z: torch.Tensor) -> torch.Tensor:
+        return self.g(z)
+
+    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        return 0.5 * ((x - self.g(z)) ** 2).sum(dim=1) / SIGMA**2
+
+
+def main() -> None:
+    torch.manual_seed(0)
+    decoder = CurveDecoder()
+    model = ebm.LatentEBM(decoder, latent_dim=1)  # standard-normal prior over z
+
+    # Ancestral reference: z ~ N(0, 1), x = g(z) + σε (exact for this Gaussian decoder).
+    z_anc = torch.randn(5000, 1)
+    x_anc = decoder.mean(z_anc) + SIGMA * torch.randn(5000, 2)
+
+    # Block-Gibbs samples of the same joint (MCMC in both blocks).
+    x_gibbs, _ = model.sample_joint(
+        ebm.MALA(step_size=0.02, steps=5), torch.randn(5000, 2), steps=300
+    )
+
+    mmd = ebm.eval.mmd(x_gibbs, x_anc)
+    mmd_ref = ebm.eval.mmd(torch.randn(5000, 2), x_anc)
+    print(f"MMD(block-Gibbs, ancestral) = {mmd:.4f}   (vs a Gaussian: {mmd_ref:.4f})")
+
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, axes = plt.subplots(1, 2, figsize=(11, 5), sharex=True, sharey=True)
+    for ax, s, title in (
+        (axes[0], x_anc, "ancestral (z→x)"),
+        (axes[1], x_gibbs, "block-Gibbs on E(x, z)"),
+    ):
+        ax.scatter(s[:, 0], s[:, 1], s=4, alpha=0.3, color="#5c50c9")
+        ax.set_title(title)
+        ax.set_aspect("equal")
+    fig.suptitle("Latent EBM: block-Gibbs sampling matches the ancestral marginal")
+    out = Path(__file__).parent / "latent_ebm_result.png"
+    fig.savefig(out, dpi=120, bbox_inches="tight")
+    print(f"saved {out}")
+
+
+if __name__ == "__main__":
+    main()

src/ebm/__init__.py

@@ -11,6 +11,7 @@
 from ebm.diffusion import DDPMAncestralSampler, VPDenoisingScoreMatching, VPSchedule
 from ebm.energy import ConditionalEnergyFn, EnergyFn, EnergyModel, score
 from ebm.jem import ClassifierEnergy, ConditionalEnergy, GuidedEnergy, JEMLoss
+from ebm.latent import LatentEBM
 from ebm.losses import (
     ConcreteScoreMatching,
     ContrastiveDivergence,
@@ -48,7 +49,7 @@
 from ebm.training import Trainer
 from ebm.utils import EMA
 
-__version__ = "0.15.0"
+__version__ = "0.16.0"
 
 __all__ = [
     "EMA",
@@ -77,6 +78,7 @@
     "GuidedEnergy",
     "JEMLoss",
     "LangevinDynamics",
+    "LatentEBM",
     "LossOutput",
     "MixtureEnergy",
     "MultiSigmaDenoisingScoreMatching",

src/ebm/latent.py

@@ -0,0 +1,134 @@
+"""Latent-variable energy-based models: a joint ``E(x, z)`` sampled by block Gibbs."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+from torch import Tensor, nn
+
+from ebm._functional import flat_sum
+from ebm.energy import EnergyFn
+from ebm.samplers.base import Sampler
+
+DecoderEnergyFn = Callable[[Tensor, Tensor], Tensor]
+"""A conditional energy ``(x: (B, *shape), z: (B, latent_dim)) -> (B,)`` — ``E(x | z)``."""
+
+
+def _standard_normal_energy(z: Tensor) -> Tensor:
+    return 0.5 * flat_sum(z.pow(2))
+
+
+class LatentEBM(nn.Module):
+    """A latent-variable EBM: joint ``E(x, z) = E_prior(z) + E_dec(x, z)``.
+
+    Couples a **prior** energy over a latent ``z`` (default the standard normal
+    ``½‖z‖²``) with a **decoder** energy ``E(x | z)`` to define a joint density
+    ``p(x, z) ∝ exp(-E(x, z))``. The data marginal ``p(x) ∝ ∫ exp(-E(x, z)) dz``
+    is generally intractable — as with all latent EBMs, you don't evaluate it, you
+    *sample the joint*. `sample_joint` runs block Gibbs, alternating an MCMC update
+    of ``z`` under its posterior ``E(z | x)`` with an update of ``x`` under
+    ``E(x | z)``; the ``x`` marginal of the joint chain is the model's ``p(x)``.
+
+    A linear-Gaussian instance is the conjugate sanity check: prior ``z ~ N(0, I)``
+    and decoder ``½‖x − Wz‖²/σ²`` (i.e. ``x | z ~ N(Wz, σ²I)``) give the exact
+    marginal ``x ~ N(0, WWᵀ + σ²I)`` and Gaussian posterior ``z | x`` — both
+    recovered by the block-Gibbs chain.
+
+    Args:
+        decoder: the conditional energy ``E(x | z)`` as a callable ``(x, z) -> (B,)``
+            (an ``nn.Module`` is registered so its parameters train and freeze).
+        latent_dim: dimensionality of ``z``.
+        prior: energy over ``z`` ``(z) -> (B,)``; defaults to the standard normal.
+    """
+
+    def __init__(
+        self,
+        decoder: DecoderEnergyFn,
+        latent_dim: int,
+        prior: EnergyFn | None = None,
+    ):
+        super().__init__()
+        if latent_dim < 1:
+            raise ValueError("latent_dim must be >= 1")
+        self.latent_dim = latent_dim
+        self._decoder = decoder
+        self._prior = prior if prior is not None else _standard_normal_energy
+        # Register any nn.Module components so params train and freeze during sampling.
+        self._components = nn.ModuleList([m for m in (decoder, prior) if isinstance(m, nn.Module)])
+
+    def prior_energy(self, z: Tensor) -> Tensor:
+        """Prior energy ``E_prior(z)``."""
+        return self._prior(z)
+
+    def decoder_energy(self, x: Tensor, z: Tensor) -> Tensor:
+        """Decoder (conditional) energy ``E(x | z)``."""
+        return self._decoder(x, z)
+
+    def joint_energy(self, x: Tensor, z: Tensor) -> Tensor:
+        """Joint energy ``E(x, z) = E_prior(z) + E(x | z)``."""
+        return self._prior(z) + self._decoder(x, z)
+
+    def conditional_energy(self, z: Tensor) -> EnergyFn:
+        """The energy ``E(x | z)`` as an `EnergyFn` over ``x`` for fixed ``z``."""
+        return lambda x: self._decoder(x, z)
+
+    def posterior_energy(self, x: Tensor) -> EnergyFn:
+        """The posterior ``E(z | x) = E_prior(z) + E(x | z)`` (up to a constant), over ``z``.
+
+        This is a valid `EnergyFn` in ``z``: the missing ``x``-only normalizer is
+        constant in ``z``, so any sampler targets the exact posterior ``p(z | x)``.
+        """
+        return lambda z: self._prior(z) + self._decoder(x, z)
+
+    @torch.no_grad()
+    def sample_joint(
+        self,
+        sampler: Sampler,
+        x_init: Tensor,
+        *,
+        z_init: Tensor | None = None,
+        steps: int = 100,
+        inner_steps: int | None = None,
+        latent_sampler: Sampler | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        """Block-Gibbs sample the joint; returns ``(x, z)``.
+
+        Each outer step updates ``z`` under `posterior_energy` then ``x`` under
+        `conditional_energy`, running ``inner_steps`` transitions of ``sampler``
+        per block (``sampler.steps`` if ``None``). ``latent_sampler`` overrides the
+        sampler used for the ``z`` block. ``z`` starts from the prior sample space
+        (standard normal) unless ``z_init`` is given.
+        """
+        x = x_init.detach().clone()
+        z = (
+            z_init.detach().clone()
+            if z_init is not None
+            else torch.randn(x.shape[0], self.latent_dim, device=x.device, dtype=x.dtype)
+        )
+        z_sampler = latent_sampler if latent_sampler is not None else sampler
+        for _ in range(steps):
+            z = z_sampler.sample(self.posterior_energy(x), z, steps=inner_steps)
+            x = sampler.sample(self.conditional_energy(z), x, steps=inner_steps)
+        return x, z
+
+    def sample(
+        self,
+        sampler: Sampler,
+        x_init: Tensor,
+        *,
+        z_init: Tensor | None = None,
+        steps: int = 100,
+        inner_steps: int | None = None,
+        latent_sampler: Sampler | None = None,
+    ) -> Tensor:
+        """Block-Gibbs sample and return the data marginal ``x`` (see `sample_joint`)."""
+        x, _ = self.sample_joint(
+            sampler,
+            x_init,
+            z_init=z_init,
+            steps=steps,
+            inner_steps=inner_steps,
+            latent_sampler=latent_sampler,
+        )
+        return x

tests/test_latent.py

@@ -0,0 +1,78 @@
+"""Latent-variable EBM checked against the linear-Gaussian conjugate identities."""
+
+import pytest
+import torch
+from torch import nn
+
+import ebm
+
+# Linear-Gaussian model: z ~ N(0, I), x | z ~ N(Wz, σ²I).
+_W = torch.tensor([[1.2, 0.4], [0.3, 0.9]])
+_SIG = 0.5
+
+
+def _decoder(x, z):
+    return 0.5 * ((x - z @ _W.t()) ** 2).sum(dim=1) / _SIG**2
+
+
+def test_latent_ebm_recovers_the_gaussian_marginal():
+    # Block Gibbs on the joint → x marginal is N(0, WWᵀ + σ²I).
+    model = ebm.LatentEBM(_decoder, latent_dim=2)
+    x, z = model.sample_joint(ebm.MALA(step_size=0.05, steps=5), torch.randn(5000, 2), steps=250)
+    assert x.shape == (5000, 2) and z.shape == (5000, 2)
+    assert not x.requires_grad
+    marginal_cov = _W @ _W.t() + _SIG**2 * torch.eye(2)
+    assert (torch.cov(x.T) - marginal_cov).abs().max().item() < 0.1
+
+
+def test_latent_ebm_posterior_is_conjugate_gaussian():
+    model = ebm.LatentEBM(_decoder, latent_dim=2)
+    x0 = torch.tensor([[1.0, -0.5]]).repeat(4000, 1)
+    # posterior_energy(x0) is a valid EnergyFn in z; sample it directly.
+    z = torch.randn(4000, 2)
+    sampler = ebm.MALA(step_size=0.03, steps=3)
+    for _ in range(400):
+        z = sampler.sample(model.posterior_energy(x0), z)
+    sig_post = torch.linalg.inv(torch.eye(2) + _W.t() @ _W / _SIG**2)
+    mu_post = (x0[:1] @ (_W / _SIG**2)) @ sig_post.t()
+    assert (torch.cov(z.T) - sig_post).abs().max().item() < 0.05
+    assert (z.mean(0) - mu_post[0]).abs().max().item() < 0.05
+
+
+def test_latent_ebm_joint_custom_prior_and_registration():
+    # joint_energy = prior + decoder; a custom prior is honored.
+    prior = ebm.nets.GaussianMixtureEnergy(torch.tensor([[-2.0, 0.0], [2.0, 0.0]]), std=0.5)
+    model = ebm.LatentEBM(_decoder, latent_dim=2, prior=prior)
+    x, z = torch.randn(8, 2), torch.randn(8, 2)
+    assert torch.allclose(model.joint_energy(x, z), prior(z) + _decoder(x, z))
+    assert torch.allclose(model.prior_energy(z), prior(z))
+    assert torch.allclose(model.decoder_energy(x, z), _decoder(x, z))
+    assert torch.allclose(model.conditional_energy(z)(x), _decoder(x, z))
+    assert torch.allclose(model.posterior_energy(x)(z), prior(z) + _decoder(x, z))
+
+    # An nn.Module decoder is registered so its parameters train/freeze.
+    class LinearDecoder(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.w = nn.Linear(2, 2, bias=False)
+
+        def forward(self, x, z):
+            return 0.5 * ((x - self.w(z)) ** 2).sum(dim=1)
+
+    dec = LinearDecoder()
+    m2 = ebm.LatentEBM(dec, latent_dim=2)
+    assert any(p is dec.w.weight for p in m2.parameters())  # registered
+
+    with pytest.raises(ValueError):
+        ebm.LatentEBM(_decoder, latent_dim=0)
+
+
+def test_latent_ebm_sample_returns_marginal_and_takes_z_init():
+    model = ebm.LatentEBM(_decoder, latent_dim=2)
+    x = model.sample(
+        ebm.MALA(step_size=0.05, steps=3),
+        torch.randn(64, 2),
+        z_init=torch.zeros(64, 2),
+        steps=5,
+    )
+    assert x.shape == (64, 2)