NUTS: the No-U-Turn Sampler (self-tuning trajectory length)

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 ## 0.16.0 — unreleased
 
+- **`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
+  step size (dual-averaging warmup to `target_accept`), so it needs neither
+  `leapfrog_steps` nor a step-size sweep. All chains are evolved in lockstep with a
+  per-chain freeze mask, so each chain's draw is identical to independent
+  single-chain NUTS. Exposes `last_tree_depth` and a `divergences` count. Validated:
+  recovers a standard and a correlated Gaussian's covariance, tunes to 0.8
+  acceptance, enters Neal's funnel, and terminates by U-turn (small tree depths,
+  zero divergences on the Gaussian). Example `nuts_sampling.py`.
 - **`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

README.md

@@ -45,7 +45,7 @@ returns `LossOutput(loss, metrics, x_neg)`; call `out.loss.backward()`).
 | Piece | Contents |
 |---|---|
 | **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) |
+| **Samplers** | `LangevinDynamics` (ULA/SGLD), `MALA`, `AdaptiveMALA` (dual-averaging step-size warmup + diagonal metric), `HMC`, `NUTS` (No-U-Turn Sampler, self-tuning trajectory length), `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`, `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` |
@@ -76,6 +76,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/
 - `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

examples/nuts_sampling.py

@@ -0,0 +1,70 @@
+"""NUTS: HMC that picks its own trajectory length, on Neal's funnel.
+
+The No-U-Turn Sampler removes HMC's two hand-tuned knobs: it doubles each
+trajectory until the path starts to double back (a U-turn), and it tunes the step
+size to a target acceptance by dual averaging during a warmup. No `leapfrog_steps`,
+no step-size sweep. Neal's funnel — a Gaussian whose width is itself a Gaussian
+latent `v` — is the classic stress test: the neck is sharp where `v` is negative,
+so the sampler must *lengthen* its trajectories there. We plot the samples and the
+distribution of tree depths (how far NUTS doubled each draw), and print the tuned
+step size, mean acceptance, and divergence count.
+
+Run:  python examples/nuts_sampling.py
+Outputs nuts_sampling_result.png next to this script (needs the [viz] extra).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+import ebm
+
+
+def main() -> None:
+    torch.manual_seed(0)
+    energy = ebm.nets.FunnelEnergy(dim=2, v_scale=3.0)  # v ~ N(0, 9), x | v ~ N(0, e^v)
+
+    sampler = ebm.NUTS(step_size=0.3, steps=300, warmup=300, target_accept=0.8)
+    x = sampler.sample(energy, torch.randn(3000, 2))
+    print(f"tuned step size = {sampler.step_size:.3f}")
+    print(f"mean acceptance = {sampler.last_accept_rate:.3f}  (target 0.8)")
+    print(
+        f"v marginal std  = {x[:, 0].std():.2f}  (true 3.0; identity metric under-samples the neck)"
+    )
+    print(f"divergences     = {sampler.divergences}")
+
+    # Collect tree depths over a batch of post-warmup draws (step size now frozen).
+    depths = []
+    xd = x
+    for _ in range(40):
+        xd = sampler.step(energy, xd)
+        depths.append(sampler.last_tree_depth)
+    depths = torch.stack(depths).reshape(-1)
+
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
+    axes[0].scatter(x[:, 1], x[:, 0], s=5, alpha=0.3, color="#5c50c9")
+    axes[0].set_xlabel("x  (neck coordinate)")
+    axes[0].set_ylabel("v  (log-scale latent)")
+    axes[0].set_title("NUTS samples of Neal's funnel")
+    axes[0].set_xlim(-15, 15)
+
+    hi = int(depths.max().item())
+    axes[1].hist(depths.numpy(), bins=range(hi + 2), align="left", rwidth=0.8, color="#5c50c9")
+    axes[1].set_xlabel("tree depth reached")
+    axes[1].set_ylabel("draws")
+    axes[1].set_title("Trajectory length adapts per draw (U-turn, not a fixed length)")
+
+    out = Path(__file__).parent / "nuts_sampling_result.png"
+    fig.savefig(out, dpi=120, bbox_inches="tight")
+    print(f"saved {out}")
+
+
+if __name__ == "__main__":
+    main()

src/ebm/__init__.py

@@ -31,6 +31,7 @@
 from ebm.samplers import (
     HMC,
     MALA,
+    NUTS,
     SVGD,
     AdaptiveMALA,
     AnnealedLangevinDynamics,
@@ -55,6 +56,7 @@
     "EMA",
     "HMC",
     "MALA",
+    "NUTS",
     "SVGD",
     "AISResult",
     "AdaptiveMALA",

src/ebm/samplers/__init__.py

@@ -10,13 +10,15 @@
     PreconditionedLangevin,
     UnderdampedLangevin,
 )
+from ebm.samplers.nuts import NUTS
 from ebm.samplers.score_sde import PredictorCorrector, ProbabilityFlowODE
 from ebm.samplers.svgd import SVGD
 from ebm.samplers.tempering import ParallelTempering, TemperedTransitions
 
 __all__ = [
     "HMC",
     "MALA",
+    "NUTS",
     "SVGD",
     "AdaptiveMALA",
     "AnnealedLangevinDynamics",