diff --git a/CHANGELOG.md b/CHANGELOG.md index 322d4d7..5130c4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,14 @@ Internal cleanup plus a diffusion track, exact-likelihood flows, and conditional-generation / calibration tools. +- **Self-tuning sampler.** `AdaptiveMALA` — a MALA that tunes its own + `step_size` by Nesterov dual averaging (Hoffman & Gelman 2014) to the + MALA-optimal 0.574 acceptance during a warmup, then freezes and samples + unbiasedly. `precondition=True` estimates a diagonal metric from a first warmup + window and re-tunes under it, so ill-conditioned targets mix at a far larger + step (`x ← x − εM∇E + √(2εM)ξ`, still exact). Validated: acceptance → 0.574 and + covariance recovered on a correlated Gaussian; the metric recovers a 100:1 + condition number. Example `adaptive_mala.py`. - **Classifier-free guidance & calibration.** `GuidedEnergy` (+ `ClassifierEnergy.guide`) — `Ẽ_w(x|y) = (1+w)E(x|y) − w E(x)` to sharpen class selection; and `eval.expected_calibration_error` / `reliability_curve` / diff --git a/README.md b/README.md index 6c3a74e..e2fde61 100644 --- a/README.md +++ b/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.ContinuousNormalizingFlow` (FFJORD — trainable exact-likelihood flows / self-normalized energies), noise-conditional variants for NCSN; `EnergyModel`, `ebm.score` | -| **Samplers** | `LangevinDynamics` (ULA/SGLD), `MALA`, `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`, `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`, `TemperedEnergy` — energies compose like densities and nest | | **Training** | thin `Trainer` (device, EMA, supervised batches, `save`/`load` checkpointing), `ReplayBuffer`, `EMA` | @@ -74,6 +74,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/ - `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 - `goodness_of_fit.py` — KSD for model selection; classifier two-sample test - `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) diff --git a/examples/adaptive_mala.py b/examples/adaptive_mala.py new file mode 100644 index 0000000..bf8fafc --- /dev/null +++ b/examples/adaptive_mala.py @@ -0,0 +1,67 @@ +"""Self-tuning MALA: no step-size sweep, and a metric for ill-conditioned targets. + +Plain MALA needs its ``step_size`` hand-tuned per target: too large and every +proposal is rejected, too small and the chain crawls. ``AdaptiveMALA`` runs a +dual-averaging warmup that drives the acceptance rate to the MALA-optimal 0.574 +on its own, then freezes and samples. On an ill-conditioned target (here a +Gaussian with a 100:1 spread between axes) an isotropic step is capped by the +tight direction; ``precondition=True`` learns a diagonal metric that lets the +usable step grow ~7× (3.0 vs 0.43 here) — the same optimal acceptance, but far +faster mixing per step. Both recover the true per-axis std. + +Run: python examples/adaptive_mala.py +Outputs adaptive_mala_result.png next to this script (needs the [viz] extra). +""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +import ebm + +COV = torch.diag(torch.tensor([25.0, 0.25])) # 100:1 condition number +PRECISION = torch.linalg.inv(COV) + + +def energy(x: torch.Tensor) -> torch.Tensor: + return 0.5 * ((x @ PRECISION) * x).sum(dim=1) + + +def main() -> None: + torch.manual_seed(0) + true_std = COV.diag().sqrt() + x0 = torch.randn(4000, 2) + + plain = ebm.AdaptiveMALA(step_size=0.1, steps=800, warmup=1000) + xp = plain.sample(energy, x0.clone()) + print(f"isotropic: tuned eps={plain.step_size:.3f} accept={plain.last_accept_rate:.3f}") + print(f" per-axis std {xp.std(0).tolist()} (true {true_std.tolist()})") + + pre = ebm.AdaptiveMALA(step_size=0.1, steps=800, warmup=1000, precondition=True) + xq = pre.sample(energy, x0.clone()) + m = pre.preconditioner + print(f"preconditioned: tuned eps={pre.step_size:.3f} accept={pre.last_accept_rate:.3f}") + print(f" learned metric ratio {(m[0] / m[1]).item():.1f} (true 100)") + print(f" per-axis std {xq.std(0).tolist()} (true {true_std.tolist()})") + + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, axes = plt.subplots(1, 2, figsize=(11, 5), sharex=True, sharey=True) + for ax, s, title in ((axes[0], xp, "isotropic"), (axes[1], xq, "preconditioned")): + ax.scatter(s[:, 0], s[:, 1], s=5, alpha=0.3, color="#5c50c9") + ax.set_title(f"AdaptiveMALA — {title}") + ax.set_xlim(-16, 16) + ax.set_ylim(-2, 2) + fig.suptitle("Self-tuning MALA on a 100:1 ill-conditioned Gaussian") + out = Path(__file__).parent / "adaptive_mala_result.png" + fig.savefig(out, dpi=120, bbox_inches="tight") + print(f"saved {out}") + + +if __name__ == "__main__": + main() diff --git a/examples/adaptive_mala_result.png b/examples/adaptive_mala_result.png new file mode 100644 index 0000000..599bf35 Binary files /dev/null and b/examples/adaptive_mala_result.png differ diff --git a/src/ebm/__init__.py b/src/ebm/__init__.py index 684bb5c..c910478 100644 --- a/src/ebm/__init__.py +++ b/src/ebm/__init__.py @@ -31,6 +31,7 @@ HMC, MALA, SVGD, + AdaptiveMALA, AnnealedLangevinDynamics, CategoricalGibbsWithGradients, GibbsSampler, @@ -55,6 +56,7 @@ "MALA", "SVGD", "AISResult", + "AdaptiveMALA", "AnnealedLangevinDynamics", "CategoricalGibbsWithGradients", "ClassifierEnergy", diff --git a/src/ebm/samplers/__init__.py b/src/ebm/samplers/__init__.py index 8684743..85ca709 100644 --- a/src/ebm/samplers/__init__.py +++ b/src/ebm/samplers/__init__.py @@ -1,3 +1,4 @@ +from ebm.samplers.adaptive import AdaptiveMALA from ebm.samplers.annealed import AnnealedLangevinDynamics from ebm.samplers.base import Sampler from ebm.samplers.discrete import CategoricalGibbsWithGradients, GibbsWithGradients @@ -17,6 +18,7 @@ "HMC", "MALA", "SVGD", + "AdaptiveMALA", "AnnealedLangevinDynamics", "CategoricalGibbsWithGradients", "GibbsSampler", diff --git a/src/ebm/samplers/adaptive.py b/src/ebm/samplers/adaptive.py new file mode 100644 index 0000000..f6ee4ff --- /dev/null +++ b/src/ebm/samplers/adaptive.py @@ -0,0 +1,163 @@ +"""Self-tuning MALA: dual-averaging step-size warmup + optional preconditioner.""" + +from __future__ import annotations + +import math + +import torch +from torch import Tensor, nn + +from ebm._functional import flat_sum as _flat_sum +from ebm._functional import mh_accept +from ebm.energy import EnergyFn +from ebm.samplers.langevin import MALA +from ebm.utils import frozen_params + + +class AdaptiveMALA(MALA): + """MALA that tunes its own ``step_size`` (and optionally a diagonal metric). + + A warmup phase adapts ``ε`` with Nesterov dual averaging (Hoffman & Gelman, + 2014) toward a target acceptance rate — ``0.574`` is optimal for MALA — then + freezes the *averaged* ``ε`` and runs an ordinary (unbiased) MALA chain. No + hand-tuning of ``step_size``; the passed value is only the starting guess. + + The recursion drives the smooth **acceptance probability** + ``ᾱ = mean min(1, e^{logα})`` (lower variance than the 0/1 accept indicator) + to ``target_accept``, updating ``logε̄`` as a running average so the frozen + value is stable rather than the last noisy iterate. + + With ``precondition=True`` a diagonal metric ``M`` (per-coordinate, geometric + mean 1) is estimated from a first warmup window's draws, then dual averaging + is **restarted** and ``ε`` re-tuned under that fixed ``M`` — the step and its + noise rescale together (``x ← x − ε M∇E + sqrt(2εM) ξ``), so the chain still + targets exactly ``p ∝ exp(-E)``. This fixes the slow mixing plain MALA + suffers on ill-conditioned targets. + + Args: + step_size: initial ε (a starting guess; overwritten by warmup). + steps: default number of *post-warmup* transitions per ``sample`` call. + warmup: dual-averaging iterations (per window; two windows if preconditioning). + target_accept: δ, the acceptance the warmup targets (0.574 is MALA-optimal). + precondition: estimate and use a diagonal metric ``M``. + gamma, t0, kappa: dual-averaging shrinkage / stabilization / decay constants. + """ + + def __init__( + self, + step_size: float = 0.1, + steps: int = 100, + *, + warmup: int = 1000, + target_accept: float = 0.574, + precondition: bool = False, + gamma: float = 0.05, + t0: float = 10.0, + kappa: float = 0.75, + ): + super().__init__(step_size, steps) + if not 0.0 < target_accept < 1.0: + raise ValueError("target_accept must be in (0, 1)") + if warmup < 0: + raise ValueError("warmup must be >= 0") + self.warmup = warmup + self.target_accept = target_accept + self.precondition = precondition + self.gamma = gamma + self.t0 = t0 + self.kappa = kappa + self._precond: Tensor | None = None + + @property + def preconditioner(self) -> Tensor | None: + """The fitted diagonal metric ``M`` (``None`` until a preconditioned run).""" + return self._precond + + def _transition(self, energy: EnergyFn, x: Tensor) -> tuple[Tensor, Tensor]: + """One MALA transition; returns ``(x_next, mean acceptance probability)``. + + Honors ``self._precond`` (a diagonal metric ``M``; ``None`` = identity) in + both the drift and the noise, keeping the proposal reversible. + """ + eps = self.step_size + m = self._precond + e_x, grad_x = self._energy_grad(energy, x) + mean_fwd = x.detach() - eps * (grad_x if m is None else m * grad_x) + if m is None: + noise = math.sqrt(2 * eps) * torch.randn_like(x) + else: + noise = torch.sqrt(2 * eps * m) * torch.randn_like(x) + proposal = mean_fwd + noise + + e_prop, grad_prop = self._energy_grad(energy, proposal) + mean_bwd = proposal.detach() - eps * (grad_prop if m is None else m * grad_prop) + + inv = 1.0 if m is None else 1.0 / m + log_q_fwd = -_flat_sum((proposal - mean_fwd).pow(2) * inv) / (4 * eps) + log_q_bwd = -_flat_sum((x.detach() - mean_bwd).pow(2) * inv) / (4 * eps) + log_alpha = (e_x - e_prop) + (log_q_bwd - log_q_fwd) + + accept = torch.log(torch.rand_like(log_alpha)) < log_alpha + self._last_accept = accept.float().mean() + accept_prob = torch.exp(log_alpha.clamp(max=0.0)).mean() + return mh_accept(x.detach(), proposal.detach(), accept), accept_prob + + def _dual_average( + self, energy: EnergyFn, x: Tensor, n: int, eps0: float, *, collect: bool + ) -> tuple[Tensor, float, list[Tensor]]: + """Run ``n`` dual-averaging steps from ``eps0``; return ``(x, ε̄, draws)``.""" + if n == 0: + return x, eps0, [] + mu = math.log(10 * eps0) # bias adaptation toward larger steps + log_eps = math.log(eps0) + log_ebar = 0.0 + h_bar = 0.0 + draws: list[Tensor] = [] + for m in range(1, n + 1): + self.step_size = math.exp(log_eps) + x, accept_prob = self._transition(energy, x) + x = x.detach() + gap = self.target_accept - float(accept_prob) + h_bar = (1 - 1 / (m + self.t0)) * h_bar + gap / (m + self.t0) + log_eps = mu - math.sqrt(m) / self.gamma * h_bar + eta = m**-self.kappa + log_ebar = eta * log_eps + (1 - eta) * log_ebar + if collect and m > n // 2: + draws.append(x) + return x, math.exp(log_ebar), draws + + def sample( + self, + energy: EnergyFn, + x_init: Tensor, + *, + steps: int | None = None, + return_trajectory: bool = False, + ) -> Tensor: + """Warm up (tuning ``ε``, optionally ``M``), freeze, then sample.""" + n_steps = self.steps if steps is None else steps + eps0 = self.step_size + x = x_init.detach().clone() + module = energy if isinstance(energy, nn.Module) else None + with frozen_params(module), torch.enable_grad(): + if self.precondition: + self._precond = None # window 1: identity metric, tune ε + collect + x, _, draws = self._dual_average(energy, x, self.warmup, eps0, collect=True) + var = torch.cat(draws, dim=0).var(dim=0).clamp_min(1e-8) + self._precond = var / var.log().mean().exp() # geometric mean 1 + # window 2: fix M, RESTART dual averaging, re-tune ε + x, eps, _ = self._dual_average(energy, x, self.warmup, eps0, collect=False) + else: + self._precond = None + x, eps, _ = self._dual_average(energy, x, self.warmup, eps0, collect=False) + self.step_size = eps # freeze the averaged step size + + trajectory = [x.clone()] if return_trajectory else None + for _ in range(n_steps): + x, _ = self._transition(energy, x) + x = x.detach() + if trajectory is not None: + trajectory.append(x.clone()) + if trajectory is not None: + return torch.stack(trajectory) + return x diff --git a/tests/test_samplers.py b/tests/test_samplers.py index e71fe30..757bc85 100644 --- a/tests/test_samplers.py +++ b/tests/test_samplers.py @@ -196,3 +196,52 @@ def aniso(x): assert abs(std[1].item() - 5.0) < 0.8 with pytest.raises(ValueError): ebm.PreconditionedLangevin(precond=torch.tensor([1.0, -1.0])) + + +def _correlated_gaussian(cov): + precision = torch.linalg.inv(cov) + + def energy(x): + return 0.5 * ((x @ precision) * x).sum(dim=1) + + return energy + + +def test_adaptive_mala_tunes_step_size_to_target_accept(): + # A wildly wrong initial step is driven to the 0.574-optimal acceptance, + # and the frozen chain recovers a correlated Gaussian's covariance. + cov = torch.tensor([[2.0, 1.2], [1.2, 1.5]]) + sampler = ebm.AdaptiveMALA(step_size=0.1, steps=600, warmup=1000) + samples = sampler.sample(_correlated_gaussian(cov), 3 * torch.randn(4000, 2)) + assert abs(sampler.last_accept_rate - 0.574) < 0.05 + assert (torch.cov(samples.T) - cov).abs().max().item() < 0.15 + assert not samples.requires_grad + assert sampler.preconditioner is None + + +def test_adaptive_mala_preconditioner_learns_the_scales(): + # Target N(0, diag(25, 0.25)) — condition number 100. The estimated diagonal + # metric should recover that ~100:1 ratio, and acceptance still hits target. + cov = torch.diag(torch.tensor([25.0, 0.25])) + sampler = ebm.AdaptiveMALA(step_size=0.1, steps=600, warmup=1000, precondition=True) + samples = sampler.sample(_correlated_gaussian(cov), torch.randn(4000, 2)) + m = sampler.preconditioner + assert m is not None + ratio = (m[0] / m[1]).item() + assert 40.0 < ratio < 250.0 # true condition number is 100 + assert abs(m.log().mean().exp().item() - 1.0) < 1e-4 # geometric mean 1 + assert abs(sampler.last_accept_rate - 0.574) < 0.06 + assert abs(samples.std(0)[0].item() - 5.0) < 0.7 # sqrt(25) + assert abs(samples.std(0)[1].item() - 0.5) < 0.1 # sqrt(0.25) + + +def test_adaptive_mala_zero_warmup_keeps_step_size_and_validates(): + sampler = ebm.AdaptiveMALA(step_size=0.2, steps=8, warmup=0) + traj = sampler.sample(quadratic_energy, torch.randn(256, 2), return_trajectory=True) + assert traj.shape == (9, 256, 2) # init + 8 transitions + assert sampler.step_size == 0.2 # nothing to adapt + assert isinstance(sampler.last_accept_rate, float) + with pytest.raises(ValueError): + ebm.AdaptiveMALA(target_accept=1.5) + with pytest.raises(ValueError): + ebm.AdaptiveMALA(warmup=-1)