Conditional generation + calibration: classifier-free guidance, ECE / temperature scaling

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

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

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()

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",

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",
 ]

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()

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)``.
 

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

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