diff --git a/CHANGELOG.md b/CHANGELOG.md index a86f234..f7a850e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index 5e925e9..649db4c 100644 --- a/README.md +++ b/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 diff --git a/examples/latent_ebm.py b/examples/latent_ebm.py new file mode 100644 index 0000000..f04f211 --- /dev/null +++ b/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() diff --git a/examples/latent_ebm_result.png b/examples/latent_ebm_result.png new file mode 100644 index 0000000..b3df267 Binary files /dev/null and b/examples/latent_ebm_result.png differ diff --git a/src/ebm/__init__.py b/src/ebm/__init__.py index af45d6f..816cb77 100644 --- a/src/ebm/__init__.py +++ b/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", diff --git a/src/ebm/latent.py b/src/ebm/latent.py new file mode 100644 index 0000000..8c5371b --- /dev/null +++ b/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 diff --git a/tests/test_latent.py b/tests/test_latent.py new file mode 100644 index 0000000..f2263a8 --- /dev/null +++ b/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)