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

[davidkhjo]

LatentEBM — the roadmap's final item (E7): a latent-variable energy-based model.

What

• Couples a prior energy over a latent z (default standard normal) with a decoder energy E(x | z) into a joint E(x, z) = E_prior(z) + E(x | z), p(x, z) ∝ exp(-E).
• The data marginal p(x) ∝ ∫ exp(-E(x,z)) dz is intractable — as with all latent EBMs you sample the joint. sample_joint runs block Gibbs: alternate an MCMC update of z under its posterior E(z | x) with an update of x under E(x | z). posterior_energy(x) and conditional_energy(z) expose those blocks as plain EnergyFns usable with any sampler.
• nn.Module decoder/prior components are registered, so their parameters train and freeze during sampling like everything else.

Validation (closed-form / distributional)

• Linear-Gaussian conjugate case (z ~ N(0,I), x|z ~ N(Wz, σ²I)): block Gibbs recovers the exact marginal N(0, WWᵀ + σ²I), and posterior_energy(x) sampled directly matches the conjugate Gaussian posterior z|x (covariance and mean).
• Example latent_ebm.py: a fixed nonlinear decoder bends a Gaussian latent into a curved 2-D manifold; the block-Gibbs marginal matches the ancestral marginal at MMD ≈ 0 (vs 0.14 for a Gaussian).

Notes

• Opens the 0.16.0 cycle (0.15.0 is now published to PyPI). torch-only; public API preserved.
• With this, the full E1–E7 expansion roadmap is complete. NUTS remains the one deliberate deferral (batched per-chain termination is easy to bias silently).

🤖 Generated with Claude Code