ebmkit

Energy-based models in PyTorch — samplers, training losses, and honest evaluation, as composable objects with tested defaults.

An EBM is an unnormalized density p(x) ∝ exp(-E(x)) defined by a network E: (B, *shape) -> (B,). torch is the only runtime dependency.

Not the Explainable Boosting Machines that also go by "EBM" — this is the deep-learning kind (LeCun et al. 2006; Du & Mordatch 2019; Song & Kingma 2021).

Install

pip install ebmkit          # runtime dependency is just torch>=2.0
pip install "ebmkit[viz]"   # + matplotlib for the plotting helpers

The import name is ebm:

import torch, ebm

energy = ebm.nets.MLPEnergy(dim=2, hidden=(128, 128))
sampler = ebm.LangevinDynamics(step_size=1e-2, steps=60)
loss_fn = ebm.ContrastiveDivergence(sampler, buffer=ebm.ReplayBuffer(8192, (2,)))

trainer = ebm.Trainer(energy, loss_fn, lr=1e-3)
trainer.fit(ebm.datasets.two_moons(8192), steps=3000, batch_size=256)

samples = sampler.sample(energy, torch.randn(2000, 2), steps=500)

The Trainer is optional sugar — the loop underneath is plain PyTorch (each loss returns LossOutput(loss, metrics, x_neg); call out.loss.backward()).

What's in the box

Examples

Runnable scripts in examples/ (python examples/<name>.py):

• train_two_moons.py — the canonical 2D contrastive-divergence smoke test
• train_jem.py / train_mnist_jem.py — classify, generate, and detect OOD with one network
• jem_guidance.py — classifier-free guidance sharpening class-conditional samples
• train_mnist.py — the image-scale IGEBM short-run recipe
• train_composition.py — product of experts / mixture / tempering, without retraining
• train_ising.py / train_potts.py — discrete lattices via (categorical) Gibbs-with-Gradients
• train_ncsn.py — score-based generation: multi-sigma denoising + annealed Langevin
• deterministic_sampling.py — annealed Langevin vs the (deterministic) probability-flow ODE vs predictor-corrector
• train_diffusion.py — a variance-preserving (DDPM) diffusion, energy-parameterized
• exact_likelihood_ode.py — exact bits/dim for a score model via the probability-flow ODE
• train_flow.py — a RealNVP normalizing flow with exact density and sampling
• train_spline_flow.py — a rational-quadratic spline flow: sharper fit than affine at equal depth
• train_cnf.py — a continuous normalizing flow (FFJORD): exact likelihood by ODE
• train_cifar_ood.py — energy-based OOD at color scale (CIFAR-10 vs CIFAR-100)
• train_rbm.py — a Bernoulli RBM on binary bars via CD-1, with the exact log Z
• train_ising_pseudolikelihood.py — recover an Ising coupling with no MCMC in the loop
• train_potts_concrete.py — recover a categorical density with concrete score matching (no MCMC)
• train_energy_discrepancy.py — two-moons trained MCMC-free (energy discrepancy)
• 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
• nuts_sampling.py — the No-U-Turn Sampler on Neal's funnel: trajectory length adapts per draw
• 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

Conventions

• Sign: p ∝ exp(-E) — low energy is high probability. Samplers descend the energy gradient; training pushes data energy down. Never flip this.
• Stop-gradients: MCMC negatives are detached and the energy's parameters are frozen during sampling; score-matching losses instead keep the graph (create_graph=True).
• The CD loss value is not a convergence signal — it hovers near zero at equilibrium; watch metrics["energy_gap"] and energy histograms.

See CONTRIBUTING.md for the rest.

Development

uv run pytest        # tests (CPU-only, seeded)
uv run ruff check .  # lint
uv run mypy          # type-check

Citation

If you use ebmkit in your research, please cite it — see CITATION.cff.

License

MIT — see LICENSE.