diff --git a/CHANGELOG.md b/CHANGELOG.md index 37e0ef4..91ef7f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,9 @@ # Changelog -## 0.16.0 — unreleased +## 0.16.0 — 2026-08-21 + +The No-U-Turn Sampler and a latent-variable EBM — the last two items on the +roadmap. Both closed-form / distributionally validated and torch-only. - **`NUTS`** — the No-U-Turn Sampler (Hoffman & Gelman 2014, multinomial variant): HMC that tunes its own trajectory length (doubling until a whole-span U-turn) and diff --git a/ebmkit-blog-post.txt b/ebmkit-blog-post.txt deleted file mode 100644 index c25da8c..0000000 --- a/ebmkit-blog-post.txt +++ /dev/null @@ -1,315 +0,0 @@ -# ebmkit: an energy-based-models library you can check against the math - -*Draft blog post. Plain text / light markdown so it ports anywhere. Figures are -referenced inline by filename — they live in the repo's `examples/` folder.* - ---- - -`pip install ebmkit`, `import ebm`. It's a small PyTorch library for energy-based -models — samplers, training losses, and evaluation, as composable objects with -tested defaults. The one thing I want it to be known for: **almost every piece is -validated against closed-form ground truth.** Not "looks reasonable on a plot" — -checked against a number you can derive with a pen. - -This post is about why I built it that way, and what "validated against the math" -actually buys you. - -## Energy-based models, and why they're a minefield - -An energy-based model is about the simplest generative model you can write down. -You define a scalar energy `E(x)` and declare that the probability of a point is - - p(x) ∝ exp(-E(x)) - -Low energy, high probability. That's it — no normalized output layer, no -invertibility constraint, no fixed latent shape. The energy is just a network -that maps a batch of data to a batch of scalars. The flexibility is the whole -appeal: the same object can be an image model, a lattice model, a classifier, or -a density you compose out of other densities. - -The catch is that almost everything downstream of that definition is a place to -get subtly, silently wrong: - -- **The sign.** `p ∝ exp(-E)` means samplers *descend* the energy and training - pushes data energy *down*. Flip it once — in a loss, in a sampler, in a metric — - and everything still runs, still produces plots, and is completely wrong. -- **The negatives.** Contrastive methods pull data energy down and "negative" - (model) samples up. Those negatives come from MCMC, and if you forget to detach - them, or forget to freeze the network while sampling, your gradient quietly - includes terms that shouldn't be there. -- **The sampler.** Langevin dynamics has a "run it correctly" regime and a "short - run practitioner" regime that look almost identical in code and behave nothing - alike. MCMC that hasn't mixed produces samples that look fine and are biased. -- **The loss value.** For contrastive divergence, the loss hovers near zero at - equilibrium — it is *not* a convergence signal. Watch it instead of the energy - gap and you'll declare victory on a model that hasn't learned anything. - -None of these throw an exception. They produce a training curve that goes down -and samples that look plausible. That's what makes EBMs a minefield: the failure -mode is a confident wrong answer. - -And the ecosystem doesn't help. EBMs are **paper-rich and tool-poor.** When I -surveyed the landscape, I found essentially one maintained general-purpose -deep-learning EBM library, and around it a graveyard of frozen paper code — JEM, -ebm-anatomy, igebm-pytorch, the original release repos — each implementing one -method from one paper, un-maintained, mutually incompatible. Every EBM project -starts by re-implementing the same samplers, buffers, and diagnostics from -scratch, and re-introducing the same sign bugs. - -(An aside the name has to address: search "EBM" and you'll mostly find -Explainable Boosting Machines, the glassbox GAM from interpretml. This is the -other EBM — the deep-learning kind, LeCun et al. 2006, Du & Mordatch 2019, Song & -Kingma 2021. The PyPI package is `ebmkit`; the import stays `ebm`.) - -The bet behind this library is simple: **narrow and reliable beats broad.** Ship -the pieces every EBM project rebuilds, get the conventions right once, and — the -part nobody does — check them against ground truth. - -## Design: a few decisions, held firmly - -The whole library follows from a handful of choices, each one aimed at a specific -failure above. - -**An energy is a callable, not a base class.** Anything with the signature -`(B, *event_shape) -> (B,)` works everywhere — a plain function, a lambda, an -`nn.Module`. There's no `Energy` superclass to inherit, no registration. You pass -your model *to* the algorithm (the pattern torchdiffeq uses for ODEs). - -```python -import torch, ebm - -# an energy is just a function to a scalar per sample -def gaussian_energy(x): - return 0.5 * x.pow(2).sum(dim=1) # p(x) = N(0, I) - -sampler = ebm.MALA(step_size=0.2, steps=200) -samples = sampler.sample(gaussian_energy, torch.randn(1000, 2)) -``` - -**One sign convention, enforced and tested.** `p ∝ exp(-E)` everywhere, and it's -pinned by tests that recover known Gaussians — if someone flips a sign, a -distributional test fails, not a code review. - -**Losses are `nn.Module`s with one contract.** Every loss is called -`loss_fn(energy, x)` and returns a `LossOutput(loss, metrics, x_neg)`. Same shape -whether it's contrastive divergence, score matching, or noise-contrastive -estimation. (This fixes a real papercut in the one existing library, where losses -returned different types and you had to remember which.) - -```python -loss_fn = ebm.ContrastiveDivergence( - ebm.LangevinDynamics(step_size=1e-2, steps=60), - buffer=ebm.ReplayBuffer(8192, (2,)), -) -out = loss_fn(energy, x) # -> LossOutput -out.loss.backward() # metrics like out.metrics["energy_gap"] for free -``` - -**The two Langevin regimes are one keyword apart.** With `noise_scale=None` you -get the mathematically correct `sqrt(2ε)` noise and a chain that targets -`exp(-E)`. Pass the decoupled cold noise plus gradient clipping and you get the -IGEBM short-run recipe practitioners actually use for images. Same class, one -argument — and the docs say which is which and why. - -**The Trainer is thin and optional.** It handles device, EMA, and checkpointing, -but the loop underneath is six lines of plain PyTorch. No Lightning-style -inheritance, no framework lock-in. You can throw the Trainer away and keep the -library. - -**No BatchNorm in energy nets; spectral norm is a flag; SiLU activations.** -BatchNorm breaks per-sample energies and MCMC — so it's simply not there. - -## The part I actually care about: checking against the math - -Here is the thing that makes this library different. For a large fraction of the -surface, there is a *closed form* you can compare against — a known density, an -analytic partition function, an exact log-likelihood — and the test suite -compares against it. Not a golden file from a previous run. The real number. - -Here's a sample, with the **actual measured error** on each (not the test -tolerance — the number you get when you run it): - -| What | Ground truth | Measured error | -|------------------------------|--------------------------------------|----------------| -| RBM free energy | brute-force enumerated joint marginal | 1.5e-8 | -| ESS diagnostic | AR(1) closed form N(1−ρ)/(1+ρ) | 0.17% (rel.) | -| AIS log-partition | log Z = log(2π) for N(0, I) in 2D | 3e-8 | -| PF-ODE exact likelihood | analytic log N(x; 0, I) | 0.061 nats (mean) | -| Kernel Stein discrepancy | 0 for the true model, >0 otherwise | 0.007 vs 0.32 | -| Fisher divergence | two-Gaussian closed form | 0.004 | -| Mutual information (MINE) | −½ log(1 − ρ²) | 0.013 | -| Banana exact-sample variance | b²·2σ₀⁴ + σ₁² = 1.5 | 0.045 | - -Read a couple of those out loud. The Bernoulli RBM's free energy — the object it -exposes as its energy — matches the brute-force marginal of the full joint to -**1.5e-8**. Annealed importance sampling recovers the log-partition of a Gaussian, -whose true value is `log(2π)`, to **3e-8**. The effective-sample-size diagnostic, -which people usually eyeball, matches the exact autocorrelation formula for an -AR(1) chain to a fraction of a percent. - -These aren't cherry-picked demos; they're the regression tests. When I added the -probability-flow ODE likelihood (an exact-likelihood estimator for score models, -via the FFJORD change of variables), the test that gates it asserts it recovers -`log N(x; 0, I)` — and it does, to 0.06 nats. When I added concrete score matching -for categorical data, the test trains it on a fully-enumerable 3-colour lattice -and checks the recovered probability mass function against the exact pmf. - -The reason this matters: an EBM that's subtly wrong looks exactly like one that's -right. The only way to know is to build a case where you know the answer, and -check. Doing that for every primitive is the whole product. - -*(Figure: `goodness_of_fit_result.png` — the kernel Stein discrepancy, given a -family of candidate Gaussians, is minimized exactly at the true parameter; the -classifier two-sample test sits at chance for a good fit. Model selection with no -samples from the model and no partition function.)* - -*(Figure: `exact_likelihood_result.png` — exact bits-per-dimension from the -probability-flow ODE cleanly separates in-distribution from out-of-distribution -points, and its average agrees with an independent AIS estimate.)* - -## Benchmarks: methods on the same honest targets - -Closed-form validation gives you something most benchmarks lack: a real ground -truth to score against, instead of one model's samples versus another's. - -**Samplers, on the banana.** The banana (a twisted Gaussian) ships an -`exact_sample` method, so I can score a sampler's output against exact i.i.d. -draws with MMD, alongside the usual convergence diagnostics. Five samplers, same -target, 2000 chains: - -| Sampler | ESS | split-R̂ | accept | MMD² to exact ↓ | time (s) | -|-----------------|-----|---------|--------|-----------------|----------| -| ULA | 57 | 1.26 | — | 0.00205 | 0.1 | -| MALA | 35 | 1.42 | 0.77 | 0.00002 | 0.1 | -| HMC | 477 | 1.02 | 0.99 | 0.00007 | 0.6 | -| Underdamped | 30 | 1.41 | — | -0.00015 | 0.1 | -| Preconditioned | 19 | 2.04 | — | 0.00146 | 0.1 | - -The honest reading is the interesting part. HMC dominates on mixing (ESS 477, -R̂ ≈ 1.02). The diagonal **preconditioner actively hurts** here (R̂ ≈ 2.0) — -because the banana's correlation isn't axis-aligned, so a diagonal rescaling -fights the geometry. That's not a bug to hide; it's the kind of result a -benchmark exists to surface. - -*(Figure: `benchmark_samplers_result.png` — each sampler's output overlaid on the -exact draws.)* - -**Training methods, on two-moons.** Six training methods — short-run and -persistent contrastive divergence, noise-contrastive estimation, sliced score -matching, multi-sigma denoising (NCSN), and diffusion recovery likelihood — -trained on the same data and scored with Fréchet distance, MMD, an AIS-bracketed -test log-likelihood, and wall time: - -| Method | FD ↓ | MMD² ↓ | Test log-lik (nats) | Train time | -|---------------------|-------|--------|---------------------|------------| -| CD (short-run) | 0.002 | 0.0010 | [−1.16, −0.80] | 100 s | -| Persistent CD | 0.008 | 0.0013 | [−1.12, −1.11] | 104 s | -| NCE | 0.002 | 0.0016 | [−1.04, −0.92] | 13 s | -| Sliced SM | 0.038 | 0.0072 | [−1.24, −1.16] | 11 s | -| NCSN (multi-σ DSM) | 0.000 | 0.0183 | [−1.77, −1.73] | 8 s | -| DRL | 0.003 | 0.0157 | [−1.91, −1.88] | 76 s | - -The scores tell a more nuanced story than any single column. Noise-contrastive -estimation is the quiet winner here — fastest of the accurate methods and the -best test likelihood — while contrastive divergence gets tight samples but pays -for its MCMC in wall time. The instructive row is NCSN: it posts the *best* -Fréchet distance (0.000) and the *worst* MMD (0.018) of the six. Fréchet distance -only sees the first two moments, so a slightly blurred ring slips right past it; -MMD at a structure-scale bandwidth does not. Trust FD alone and you'd crown the -blurriest model. Sliced score matching has the weakest fit on this sharp 2-D -density — score matching earns its keep on smoother, higher-dimensional problems, -not two-moons. None of that is visible from a training curve; it takes a metric -with a ground truth behind it. - -*(Figure: `benchmark_result.png` — samples from each method, model in purple over -the data in grey.)* - -Every number in both tables is computed at run time by scripts in the repo -(`examples/benchmark_samplers.py`, `examples/benchmark_losses.py`) — nothing is -hard-coded, so you can reproduce or extend them. - -## What's actually in the box - -Once the conventions are right, breadth is cheap, because everything composes -through the same callable-energy interface. As of 0.14.0: - -- **~13 samplers** — Langevin (ULA/SGLD), MALA, HMC, underdamped Langevin - (SGHMC), preconditioned Langevin, parallel tempering, tempered transitions, - SVGD, block Gibbs and Gibbs-with-gradients (binary and categorical), annealed - Langevin, and the score-SDE pair: the deterministic probability-flow ODE and a - predictor-corrector. -- **~12 losses** — contrastive divergence (CD-k / persistent), denoising and - sliced and *exact* score matching, energy discrepancy (MCMC-free), noise- - contrastive estimation, diffusion recovery likelihood, the joint energy-based - model (JEM) loss, and three MCMC-free discrete losses: pseudo-likelihood, ratio - matching, and concrete score matching. -- **~14 evaluation metrics** — the AIS log-Z bracket and the probability-flow-ODE - exact likelihood; FID, MMD, precision/recall, inception score; kernel Stein - discrepancy, the classifier two-sample test, Fisher divergence; MINE for mutual - information; OOD AUROC; and the MCMC convergence diagnostics (effective sample - size, split-R̂). -- **Energies** — MLP, conv, and residual (IGEBM) nets; the Bernoulli RBM with an - exact `log_z`; Ising and Potts lattices; a RealNVP flow that doubles as a - self-normalized energy; and closed-form test targets (Neal's funnel, a Gaussian - mixture, the banana). -- **Composition** — energies add (product of experts), mix, and temper, and the - results are themselves energies, so they nest. `SumEnergy` is intersection, - `MixtureEnergy` is union, `TemperedEnergy` flattens or sharpens — no retraining. -- **Data & viz** — the 2D toys plus torchvision-free MNIST / Fashion-MNIST / - CIFAR loaders (parsed straight from the raw binaries), and plotting helpers. - -*(Figure: `composition_result.png` — two stripe experts combined as a product, a -mixture, and a tempered density, all without retraining.)* - -## See it run - -The toy case is six lines: - -```python -import torch, ebm - -data = ebm.datasets.two_moons(8192) -energy = ebm.nets.MLPEnergy(dim=2, hidden=(128, 128)) -loss_fn = ebm.ContrastiveDivergence( - ebm.LangevinDynamics(step_size=1e-2, steps=60), - buffer=ebm.ReplayBuffer(8192, (2,)), -) -ebm.Trainer(energy, loss_fn, lr=1e-3).fit(data, steps=3000, batch_size=256) - -samples = ebm.LangevinDynamics(step_size=1e-2, steps=500).sample( - energy, torch.randn(2000, 2) -) -``` - -It scales past toys with the same objects — the IGEBM short-run recipe generates -MNIST digits in a few minutes on Apple Silicon, and the same JEM network -classifies, generates, and does OOD detection at once. Every one of these is a -runnable script in `examples/`, each producing a figure. - -## What it deliberately doesn't do - -Narrowness is a feature, so it's worth naming the non-goals. ebmkit is not trying -to be a general generative-modeling framework. It doesn't chase flow matching or -Schrödinger bridges (the direction the one other library drifted). `torch` is the -only runtime dependency — no numpy, scipy, or sklearn in the core. It won't win a -FID leaderboard: the point is correct, composable, inspectable primitives, not a -state-of-the-art image model. - -And when a benchmark is unflattering, it stays. The CIFAR-10-vs-CIFAR-100 OOD -example reports ~0.57 AUROC — a genuinely hard near-OOD task where natural images -sit close together — and the docstring says so, rather than quietly swapping in -far-OOD SVHN to post a prettier number. A library whose whole pitch is honesty -about correctness doesn't get to cherry-pick its demos. - -## Where it's going - -ebmkit is on PyPI (`pip install ebmkit`, `import ebm`), MIT-licensed, torch-only, -tested across Python 3.10–3.13. The research directions I'm eyeing next both build -on what's there: a variance-preserving / DDPM diffusion track to sit alongside -the variance-exploding score stack, and a trainable neural-ODE flow that learns -through the exact-likelihood machinery already shipped. - -If you work with energy-based models and you've ever burned a day on a sign flip -or an un-mixed chain, that's exactly the pain this is meant to remove. The repo, -issues, and contribution guide are on GitHub. Try to break a closed-form test — -that's the most useful bug report there is.