ebmkit

A small, reliable PyTorch toolkit for training and using energy-based models — the MCMC samplers, training losses, replay buffers, and diagnostics every EBM project otherwise rebuilds from scratch, 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

Docs

• Training methods — choosing a loss; CD/PCD, NCSN, DRL, NCE, JEM
• Samplers — Langevin/MALA/HMC, Gibbs-with-Gradients, annealed
• Evaluation — log-Z bracketing, FID, MMD, OOD, bits/dim
• Composition — products, mixtures, tempering
• Benchmarks — every loss family scored on the eval stack

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
• 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
• train_cifar_ood.py — energy-based OOD at color scale (CIFAR-10 vs CIFAR-100)
• 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 docs/ and 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.