diff --git a/CHANGELOG.md b/CHANGELOG.md index b789a8c..322d4d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,13 @@ ## 0.15.0 — unreleased -Internal cleanup plus the start of a diffusion track. +Internal cleanup plus a diffusion track, exact-likelihood flows, and +conditional-generation / calibration tools. + +- **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` / + `temperature_scale` for trustworthy JEM classifiers. Example `jem_guidance.py`. - **Housekeeping (no API change).** Deduplicated shared helpers into a private `ebm/_functional.py`, and split the two largest modules into packages diff --git a/README.md b/README.md index 56ed535..6c3a74e 100644 --- a/README.md +++ b/README.md @@ -44,12 +44,12 @@ 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`), `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` | +| **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) | | **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` | -| **Eval** | `ais_log_z` / `reverse_ais_log_z` (bracket `log Z`), `pf_ode_log_likelihood` (exact likelihood via the probability-flow ODE), `bits_per_dim`, `frechet_distance` (FID), `mmd`, `precision_recall`, `inception_score`, `kernel_stein_discrepancy` / `classifier_two_sample_test` / `fisher_divergence` (goodness-of-fit), `mutual_information` (MINE), `ood_auroc`, `effective_sample_size` / `split_rhat` (MCMC diagnostics) | +| **Eval** | `ais_log_z` / `reverse_ais_log_z` (bracket `log Z`), `pf_ode_log_likelihood` (exact likelihood via the probability-flow ODE), `bits_per_dim`, `frechet_distance` (FID), `mmd`, `precision_recall`, `inception_score`, `kernel_stein_discrepancy` / `classifier_two_sample_test` / `fisher_divergence` (goodness-of-fit), `mutual_information` (MINE), `expected_calibration_error` / `reliability_curve` / `temperature_scale` (calibration), `ood_auroc`, `effective_sample_size` / `split_rhat` (MCMC diagnostics) | | **Data & viz** | 2D toys (`two_moons`, `eight_gaussians`, `checkerboard`, `rings`, `spirals`) and torchvision-free image loaders (`mnist`, `fashion_mnist`, `cifar10`, `cifar100`); `viz.energy_contour` / `plot_samples` / `energy_histogram` / `show_images` | ## Examples @@ -58,6 +58,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/ - `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 diff --git a/examples/jem_guidance.py b/examples/jem_guidance.py new file mode 100644 index 0000000..3b0bed7 --- /dev/null +++ b/examples/jem_guidance.py @@ -0,0 +1,62 @@ +"""Classifier-free guidance on an energy classifier — sharpening class selection. + +`ebm.ClassifierEnergy` turns any classifier into an energy with a marginal `E(x)` +and per-class conditionals `E(x|y)`. Classifier-free guidance combines them, +`Ẽ_w(x|y) = (1+w)E(x|y) − w E(x)`, to concentrate samples on the target class. +The effect is only visible when the classes *overlap* (a confident, well-separated +classifier needs no guidance), so we use two overlapping Gaussian classes and +watch guidance suppress the wrong mode as `w` grows. The identical +`energy.guide(y, w)` call works on a trained JEM (`ebm.JEMLoss`). + +Run: python examples/jem_guidance.py +Outputs jem_guidance_result.png next to this script (needs the [viz] extra). +""" + +from __future__ import annotations + +from pathlib import Path + +import torch +from torch import nn + +import ebm + +MU = torch.tensor([[-1.5, 0.0], [1.5, 0.0]]) # two overlapping class centres + + +class GaussianLogits(nn.Module): + """Logits[:, k] = -½‖x - μ_k‖², so E(x|k) = ½‖x - μ_k‖² and the marginal is the mixture.""" + + def forward(self, x): + return torch.stack([-0.5 * ((x - MU[k]) ** 2).sum(1) for k in range(2)], dim=1) + + +def main() -> None: + energy = ebm.ClassifierEnergy(GaussianLogits()) + + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharex=True, sharey=True) + for ax, w in zip(axes, (0.0, 2.0, 6.0), strict=True): + torch.manual_seed(1) + samples = ebm.MALA(step_size=0.05, steps=400).sample( + energy.guide(0, w), 2 * torch.randn(3000, 2) + ) + frac = (samples[:, 0] < 0).float().mean().item() # class-0 mode sits at x0 < 0 + ax.scatter(samples[:, 0], samples[:, 1], s=5, alpha=0.35, color="#5c50c9") + ax.scatter(*MU.t(), s=120, marker="x", color="#c95c50") + ax.set_title(f"guidance w = {w:.0f} (correct-mode frac {frac:.3f})") + ax.set_xlim(-4, 4) + ax.set_ylim(-3, 3) + print(f"w={w:.0f} P(correct mode) = {frac:.3f}") + fig.suptitle("Classifier-free guidance suppresses the wrong class as w grows") + out = Path(__file__).parent / "jem_guidance_result.png" + fig.savefig(out, dpi=120, bbox_inches="tight") + print(f"saved {out}") + + +if __name__ == "__main__": + main() diff --git a/examples/jem_guidance_result.png b/examples/jem_guidance_result.png new file mode 100644 index 0000000..b857b87 Binary files /dev/null and b/examples/jem_guidance_result.png differ diff --git a/src/ebm/__init__.py b/src/ebm/__init__.py index df8eb04..684bb5c 100644 --- a/src/ebm/__init__.py +++ b/src/ebm/__init__.py @@ -10,7 +10,7 @@ from ebm.compose import MixtureEnergy, SumEnergy, TemperedEnergy from ebm.diffusion import DDPMAncestralSampler, VPDenoisingScoreMatching, VPSchedule from ebm.energy import ConditionalEnergyFn, EnergyFn, EnergyModel, score -from ebm.jem import ClassifierEnergy, ConditionalEnergy, JEMLoss +from ebm.jem import ClassifierEnergy, ConditionalEnergy, GuidedEnergy, JEMLoss from ebm.losses import ( ConcreteScoreMatching, ContrastiveDivergence, @@ -71,6 +71,7 @@ "ExactScoreMatching", "GibbsSampler", "GibbsWithGradients", + "GuidedEnergy", "JEMLoss", "LangevinDynamics", "LossOutput", diff --git a/src/ebm/eval/__init__.py b/src/ebm/eval/__init__.py index 0e71b7f..5ae16d0 100644 --- a/src/ebm/eval/__init__.py +++ b/src/ebm/eval/__init__.py @@ -6,6 +6,11 @@ """ from ebm.ais import AISResult, ais_log_z, log_likelihood, reverse_ais_log_z +from ebm.eval.calibration import ( + expected_calibration_error, + reliability_curve, + temperature_scale, +) from ebm.eval.diagnostics import ( effective_sample_size, energies, @@ -33,6 +38,7 @@ "classifier_two_sample_test", "effective_sample_size", "energies", + "expected_calibration_error", "fisher_divergence", "frechet_distance", "inception_score", @@ -43,6 +49,8 @@ "ood_auroc", "pf_ode_log_likelihood", "precision_recall", + "reliability_curve", "reverse_ais_log_z", "split_rhat", + "temperature_scale", ] diff --git a/src/ebm/eval/calibration.py b/src/ebm/eval/calibration.py new file mode 100644 index 0000000..61cd930 --- /dev/null +++ b/src/ebm/eval/calibration.py @@ -0,0 +1,75 @@ +"""Classifier calibration metrics (for JEM / any softmax classifier).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import Tensor + + +def _bin_edges(n_bins: int, device) -> Tensor: + return torch.linspace(0, 1, n_bins + 1, device=device) + + +def _bin_mask(conf: Tensor, lo: Tensor, hi: Tensor, first: bool) -> Tensor: + # Left-open / right-closed bins, with the first bin including its left edge. + return (conf >= lo) & (conf <= hi) if first else (conf > lo) & (conf <= hi) + + +@torch.no_grad() +def expected_calibration_error(probs: Tensor, labels: Tensor, n_bins: int = 15) -> float: + """Top-label expected calibration error (Guo et al. 2017). + + Bins predictions by their confidence ``max_k p(k|x)`` and averages the gap + between bin accuracy and bin confidence, weighted by bin population. ``0`` for + a perfectly calibrated classifier; larger means over- or under-confident. + """ + conf, pred = probs.max(dim=1) + acc = (pred == labels).to(probs.dtype) + edges = _bin_edges(n_bins, probs.device) + ece = probs.new_zeros(()) + for i in range(n_bins): + m = _bin_mask(conf, edges[i], edges[i + 1], first=(i == 0)) + if m.any(): + ece += m.float().mean() * (acc[m].mean() - conf[m].mean()).abs() + return float(ece) + + +@torch.no_grad() +def reliability_curve( + probs: Tensor, labels: Tensor, n_bins: int = 15 +) -> tuple[Tensor, Tensor, Tensor]: + """Per-bin ``(mean confidence, accuracy, count)`` for a reliability diagram.""" + conf, pred = probs.max(dim=1) + acc = (pred == labels).to(probs.dtype) + edges = _bin_edges(n_bins, probs.device) + confs, accs, counts = [], [], [] + for i in range(n_bins): + m = _bin_mask(conf, edges[i], edges[i + 1], first=(i == 0)) + confs.append(conf[m].mean() if m.any() else (edges[i] + edges[i + 1]) / 2) + accs.append(acc[m].mean() if m.any() else probs.new_zeros(())) + counts.append(m.sum()) + return torch.stack(confs), torch.stack(accs), torch.stack(counts) + + +def temperature_scale( + logits: Tensor, labels: Tensor, *, lr: float = 0.05, max_iter: int = 200 +) -> Tensor: + """Fit a scalar temperature ``T > 0`` by NLL (Guo et al. 2017); returns ``T``. + + Divide logits by ``T`` at inference (``(logits/T).softmax(1)``) to recalibrate + ``p(y|x)`` without changing the argmax. Fit on a **held-out** split. Optimizes + ``log T`` with LBFGS so ``T`` stays positive without a constraint. + """ + logits = logits.detach() # only T is optimized, never the classifier + log_t = torch.zeros(1, device=logits.device, requires_grad=True) + opt = torch.optim.LBFGS([log_t], lr=lr, max_iter=max_iter) + + def closure() -> Tensor: + opt.zero_grad() + loss = F.cross_entropy(logits / log_t.exp(), labels) + loss.backward() + return loss + + opt.step(closure) # type: ignore[arg-type] + return log_t.exp().detach() diff --git a/src/ebm/jem.py b/src/ebm/jem.py index 2beaff0..2c6344f 100644 --- a/src/ebm/jem.py +++ b/src/ebm/jem.py @@ -49,6 +49,10 @@ def condition(self, y: Tensor | int) -> ConditionalEnergy: """Energy function of ``x`` with the class fixed, for conditional sampling.""" return ConditionalEnergy(self, y) + def guide(self, y: Tensor | int, weight: float = 1.0) -> GuidedEnergy: + """Classifier-free-guided energy of ``x`` for class ``y`` (weight ``w``).""" + return GuidedEnergy(self, y, weight) + class ConditionalEnergy(nn.Module): """``E(x, y)`` with ``y`` held fixed — an energy function of ``x`` alone. @@ -66,6 +70,39 @@ def forward(self, x: Tensor) -> Tensor: return self.energy.conditional(x, self.y) +class GuidedEnergy(nn.Module): + """Classifier-free-guided energy ``Ẽ_w(x|y) = (1+w)·E(x|y) − w·E(x)``. + + Sharpens class selection by amplifying the class-conditional energy relative to + the marginal — the energy of ``p̃(x|y) ∝ p(x|y)^{1+w} / p(x)^w`` (Ho & Salimans, + 2022). ``weight=0`` is the plain conditional; larger ``weight`` pushes samples + toward regions where class ``y`` dominates. Both branches come from the *same* + `ClassifierEnergy` (no label-dropout needed), and ``logits`` is computed once. + An ``nn.Module`` (not a closure) so samplers freeze the classifier during + sampling. Note: for large ``weight`` the guided energy can become + non-normalizable — keep ``weight`` modest and prefer MALA/HMC (which reject + divergent proposals). + """ + + def __init__(self, energy: ClassifierEnergy, y: Tensor | int, weight: float = 1.0): + super().__init__() + if weight < 0: + raise ValueError("guidance weight must be >= 0") + self.energy = energy + self.y = y + self.weight = weight + + def forward(self, x: Tensor) -> Tensor: + logits = self.energy.logits(x) + if isinstance(self.y, int): + e_cond = -logits[:, self.y] + else: + y = self.y.to(logits.device) + e_cond = -logits.gather(1, y.reshape(-1, 1)).squeeze(1) + e_marg = -logits.logsumexp(dim=1) + return (1 + self.weight) * e_cond - self.weight * e_marg + + class JEMLoss(nn.Module): """Joint classifier + EBM objective: ``CE(logits, y) + cd_weight * CD(E)``. diff --git a/tests/test_datasets_eval.py b/tests/test_datasets_eval.py index 4f3db38..e3625e6 100644 --- a/tests/test_datasets_eval.py +++ b/tests/test_datasets_eval.py @@ -437,3 +437,37 @@ def test_classifier_two_sample_test_validation(): ebm.eval.classifier_two_sample_test(torch.randn(100, 2), torch.randn(100, 2), test_frac=0.0) with pytest.raises(ValueError): ebm.eval.classifier_two_sample_test(torch.randn(1, 2), torch.randn(1, 2)) + + +def test_expected_calibration_error_hand_binned(): + # Two confidence bins: conf .6 (acc .5) and conf .9 (acc 1.0). + # ECE = .5·|.5-.6| + .5·|1-.9| = 0.1. + probs = torch.tensor([[0.6, 0.4], [0.6, 0.4], [0.9, 0.1], [0.9, 0.1]]) + labels = torch.tensor([0, 1, 0, 0]) + assert abs(ebm.eval.expected_calibration_error(probs, labels, n_bins=10) - 0.1) < 1e-6 + + # A perfectly-calibrated stream (confidence == accuracy) → ECE → 0. + n = 200000 + conf = 0.5 + 0.5 * torch.rand(n) + correct = torch.rand(n) < conf + p = torch.stack([torch.where(correct, conf, 1 - conf), torch.where(correct, 1 - conf, conf)], 1) + assert ebm.eval.expected_calibration_error(p, torch.zeros(n, dtype=torch.long)) < 0.01 + + +def test_temperature_scaling_recovers_and_calibrates(): + import torch.nn.functional as tnf + + k, n = 5, 8000 + true = torch.randn(n, k) + labels = torch.distributions.Categorical(logits=true).sample() + logits = true * 3.0 # over-confident relative to the label-generating softmax + t = ebm.eval.temperature_scale(logits, labels) + assert abs(t.item() - 3.0) < 0.5 # recovers the true temperature + assert tnf.cross_entropy(logits / t, labels) < tnf.cross_entropy(logits, labels) + ece0 = ebm.eval.expected_calibration_error(logits.softmax(1), labels) + ece_t = ebm.eval.expected_calibration_error((logits / t).softmax(1), labels) + assert ece_t < ece0 + + conf, acc, count = ebm.eval.reliability_curve(logits.softmax(1), labels, n_bins=10) + assert conf.shape == acc.shape == count.shape == (10,) + assert int(count.sum()) == n diff --git a/tests/test_jem.py b/tests/test_jem.py index f1a63a5..39359c5 100644 --- a/tests/test_jem.py +++ b/tests/test_jem.py @@ -1,3 +1,4 @@ +import pytest import torch from torch import nn @@ -141,3 +142,44 @@ def test_jem_joint_training_two_gaussians(): s1 = sampler.sample(energy.condition(1), torch.randn(256, 2)) assert s0[:, 0].mean() < 0 assert s1[:, 0].mean() > 0 + + +class _GaussianLogits(nn.Module): + """Logits[:, k] = -½‖x - μ_k‖² so E(x|k)=½‖x-μ_k‖² and the marginal is a mixture.""" + + def __init__(self, mus): + super().__init__() + self.register_buffer("mus", mus) + + def forward(self, x): + return torch.stack( + [-0.5 * ((x - self.mus[k]) ** 2).sum(1) for k in range(len(self.mus))], 1 + ) + + +def test_guided_energy_formula_and_validation(): + ce = ebm.ClassifierEnergy(nn.Linear(2, 3)) + x = torch.randn(50, 2) + w = 2.5 + manual = (1 + w) * ce.conditional(x, 1) - w * ce(x) + assert torch.allclose(ce.guide(1, w)(x), manual) # exact CFG formula + with pytest.raises(ValueError): + ebm.GuidedEnergy(ce, 0, weight=-1.0) + + +def test_guided_energy_recovers_conditional_mean(): + mu = torch.tensor([[2.0, -1.0], [-3.0, 3.0]]) + ce = ebm.ClassifierEnergy(_GaussianLogits(mu)) + samples = ebm.MALA(step_size=0.05, steps=400).sample(ce.guide(0, 3.0), torch.randn(4000, 2)) + assert (samples.mean(0) - mu[0]).abs().max().item() < 0.1 + + +def test_guided_energy_suppresses_the_wrong_mode(): + ce = ebm.ClassifierEnergy(_GaussianLogits(torch.tensor([[-2.0, 0.0], [2.0, 0.0]]))) + fracs = [] + for w in (0.0, 2.0, 6.0): + torch.manual_seed(1) + s = ebm.MALA(step_size=0.05, steps=400).sample(ce.guide(0, w), 2 * torch.randn(3000, 2)) + fracs.append((s[:, 0] < 0).float().mean().item()) # class-0 mode is at x0<0 + assert fracs[0] < fracs[1] < fracs[2] # guidance monotonically suppresses the wrong mode + assert fracs[2] > 0.99