diff --git a/CHANGELOG.md b/CHANGELOG.md index 5130c4b..e64ae77 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. +- **`nets.NeuralSplineCouplingFlow`** — a rational-quadratic neural spline flow + (Durkan et al. 2019): monotonic-spline coupling layers dropped into the + affine-flow scaffold, strictly more expressive per layer, so it fits sharp + multi-modal 2D densities at a lower NLL than the affine flow for the same + depth. Same self-normalized contract (exact `log_prob`, one-pass sampling, + `forward = -log_prob`). Validated: exact invertibility incl. the linear tails + and analytic `log|det|` vs the autograd Jacobian (to 1e-8 in double precision); + fits a Gaussian and two-moons. Example `train_spline_flow.py`. - **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 diff --git a/README.md b/README.md index e2fde61..b355e3c 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,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` | +| **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) | | **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 | @@ -67,6 +67,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/ - `train_diffusion.py` — a variance-preserving (DDPM) diffusion, energy-parameterized - `exact_likelihood_ode.py` — exact bits/dim for a score model via the probability-flow ODE - `train_flow.py` — a RealNVP normalizing flow with exact density and sampling +- `train_spline_flow.py` — a rational-quadratic spline flow: sharper fit than affine at equal depth - `train_cnf.py` — a continuous normalizing flow (FFJORD): exact likelihood by ODE - `train_cifar_ood.py` — energy-based OOD at color scale (CIFAR-10 vs CIFAR-100) - `train_rbm.py` — a Bernoulli RBM on binary bars via CD-1, with the exact `log Z` diff --git a/examples/spline_flow_result.png b/examples/spline_flow_result.png new file mode 100644 index 0000000..ebdb610 Binary files /dev/null and b/examples/spline_flow_result.png differ diff --git a/examples/train_spline_flow.py b/examples/train_spline_flow.py new file mode 100644 index 0000000..b0e9e58 --- /dev/null +++ b/examples/train_spline_flow.py @@ -0,0 +1,74 @@ +"""Fit two-moons with a rational-quadratic neural spline flow (Durkan et al. 2019). + +`NeuralSplineCouplingFlow` is a drop-in for `AffineCouplingFlow` whose per-layer +transform is a monotonic rational-quadratic spline instead of an affine map — +strictly more expressive, so it carves the sharp two-moons crescents with fewer +layers and a lower held-out NLL. Same self-normalized contract: exact `log_prob` +(no MCMC), one-pass sampling, and `forward(x) = -log_prob(x)` is a valid energy. +We train an affine flow and a spline flow with the *same* layer budget and print +both held-out NLLs so the gap is visible. + +Run: python examples/train_spline_flow.py +Outputs spline_flow_result.png next to this script (requires the [viz] extra). +""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +import ebm + + +def _train(flow: torch.nn.Module, data: torch.Tensor, steps: int = 4000) -> torch.nn.Module: + opt = torch.optim.Adam(flow.parameters(), lr=3e-3) + for _ in range(steps): + batch = data[torch.randint(0, len(data), (256,))] + loss = -flow.log_prob(batch).mean() # exact negative log-likelihood + opt.zero_grad() + loss.backward() + opt.step() + return flow + + +def main() -> None: + torch.manual_seed(0) + data = ebm.datasets.two_moons(8192) + test = ebm.datasets.two_moons(4000, generator=torch.Generator().manual_seed(1)) + + affine = _train(ebm.nets.AffineCouplingFlow(dim=2, n_layers=6, hidden=64), data) + spline = _train( + ebm.nets.NeuralSplineCouplingFlow(dim=2, n_layers=6, num_bins=8, bound=3.0), data + ) + + for name, flow in (("affine", affine), ("spline", spline)): + print(f"{name:7s} held-out NLL = {-flow.log_prob(test).mean().item():.3f} nats") + + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + lin = torch.linspace(-2.5, 2.5, 200) + gy, gx = torch.meshgrid(lin, lin, indexing="ij") + grid = torch.stack([gx.reshape(-1), gy.reshape(-1)], dim=1) + + fig, axes = plt.subplots(1, 3, figsize=(15, 5)) + ebm.viz.plot_samples(data[:2000], ax=axes[0]) + axes[0].set_title("data") + for ax, (name, flow) in zip(axes[1:], (("affine", affine), ("spline", spline)), strict=True): + with torch.no_grad(): + dens = flow.log_prob(grid).exp().reshape(200, 200) + ax.imshow( + dens, extent=(-2.5, 2.5, -2.5, 2.5), origin="lower", vmax=float(dens.quantile(0.98)) + ) + ax.set_title(f"{name} flow density (exact)") + fig.suptitle("Affine vs rational-quadratic spline coupling — same 6-layer budget") + out = Path(__file__).parent / "spline_flow_result.png" + fig.savefig(out, dpi=120, bbox_inches="tight") + print(f"saved {out}") + + +if __name__ == "__main__": + main() diff --git a/src/ebm/nets/__init__.py b/src/ebm/nets/__init__.py index 0dd8491..095ad01 100644 --- a/src/ebm/nets/__init__.py +++ b/src/ebm/nets/__init__.py @@ -22,6 +22,7 @@ from ebm.nets.noise_conditional import ( _GaussianFourierFeatures as _GaussianFourierFeatures, # re-export for tests ) +from ebm.nets.spline_flow import NeuralSplineCouplingFlow from ebm.nets.targets import BananaEnergy, FunnelEnergy, GaussianMixtureEnergy __all__ = [ @@ -35,6 +36,7 @@ "GaussianMixtureEnergy", "IsingEnergy", "MLPEnergy", + "NeuralSplineCouplingFlow", "NoiseConditionalConvEnergy", "NoiseConditionalMLPEnergy", "PottsEnergy", diff --git a/src/ebm/nets/spline_flow.py b/src/ebm/nets/spline_flow.py new file mode 100644 index 0000000..c8d0bd5 --- /dev/null +++ b/src/ebm/nets/spline_flow.py @@ -0,0 +1,209 @@ +"""Rational-quadratic neural spline flow (Durkan et al. 2019).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from ebm._functional import standard_normal_logprob + + +def _searchsorted(sorted_seq: Tensor, values: Tensor) -> Tensor: + """Per-row bin index ``k`` with ``sorted_seq[k] <= values < sorted_seq[k+1]``.""" + return torch.searchsorted(sorted_seq, values[..., None], right=True)[..., 0] - 1 + + +def _monotonic_rq_spline( + inputs: Tensor, + unnorm_widths: Tensor, + unnorm_heights: Tensor, + unnorm_derivs: Tensor, + *, + inverse: bool, + bound: float, + min_bin: float, + min_deriv: float, +) -> tuple[Tensor, Tensor]: + """Elementwise rational-quadratic transform on ``[-bound, bound]`` (identity tails). + + ``inputs`` is ``(N,)``; ``unnorm_widths``/``unnorm_heights`` are ``(N, K)`` and + ``unnorm_derivs`` is ``(N, K-1)``. Returns ``(outputs, log|dy/dx|)``. Outside the + interval the map is the identity (slope-1 linear tails, boundary derivatives 1), + so this is the *unconstrained* variant (Durkan et al. 2019, eqs. 4–5). + """ + k = unnorm_widths.shape[-1] + inside = (inputs >= -bound) & (inputs <= bound) + outputs = torch.where(inside, torch.zeros_like(inputs), inputs) # identity tails + logabsdet = torch.zeros_like(inputs) + if not bool(inside.any()): + return outputs, logabsdet + + inp = inputs[inside] + uw, uh, ud = unnorm_widths[inside], unnorm_heights[inside], unnorm_derivs[inside] + + widths = min_bin + (1 - min_bin * k) * F.softmax(uw, dim=-1) + cumwidths = F.pad(torch.cumsum(2 * bound * widths, dim=-1), (1, 0)) + cumwidths = cumwidths - bound + cumwidths[..., 0], cumwidths[..., -1] = -bound, bound + widths = cumwidths[..., 1:] - cumwidths[..., :-1] + + heights = min_bin + (1 - min_bin * k) * F.softmax(uh, dim=-1) + cumheights = F.pad(torch.cumsum(2 * bound * heights, dim=-1), (1, 0)) + cumheights = cumheights - bound + cumheights[..., 0], cumheights[..., -1] = -bound, bound + heights = cumheights[..., 1:] - cumheights[..., :-1] + + derivs = F.pad(min_deriv + F.softplus(ud), (1, 1), value=1.0) # slope-1 tails + delta = heights / widths # per-bin secant slope s_k + + knots = cumheights if inverse else cumwidths + idx = _searchsorted(knots, inp).clamp(0, k - 1)[..., None] + x_k = cumwidths.gather(-1, idx)[..., 0] + w_k = widths.gather(-1, idx)[..., 0] + y_k = cumheights.gather(-1, idx)[..., 0] + h_k = heights.gather(-1, idx)[..., 0] + s_k = delta.gather(-1, idx)[..., 0] + d_k = derivs.gather(-1, idx)[..., 0] + d_k1 = derivs.gather(-1, idx + 1)[..., 0] + + if inverse: + dy = inp - y_k + a = dy * (d_k1 + d_k - 2 * s_k) + h_k * (s_k - d_k) + b = h_k * d_k - dy * (d_k1 + d_k - 2 * s_k) + c = -s_k * dy + disc = b.pow(2) - 4 * a * c + theta = 2 * c / (-b - torch.sqrt(disc)) # numerically stable root + out = theta * w_k + x_k + else: + theta = (inp - x_k) / w_k + + one_m = 1 - theta + denom = s_k + (d_k1 + d_k - 2 * s_k) * theta * one_m + if not inverse: + out = y_k + h_k * (s_k * theta.pow(2) + d_k * theta * one_m) / denom + deriv_num = s_k.pow(2) * (d_k1 * theta.pow(2) + 2 * s_k * theta * one_m + d_k * one_m.pow(2)) + log_deriv = torch.log(deriv_num) - 2 * torch.log(denom) + + outputs[inside] = out + logabsdet[inside] = -log_deriv if inverse else log_deriv + return outputs, logabsdet + + +class _SplineCouplingLayer(nn.Module): + """One RQ-spline coupling: splines the ``1-mask`` half conditioned on the ``mask`` half.""" + + mask: Tensor + + def __init__(self, dim: int, hidden: int, mask: Tensor, num_bins: int, bound: float): + super().__init__() + self.register_buffer("mask", mask) + self.dim = dim + self.num_bins = num_bins + self.bound = bound + self.min_bin = 1e-3 + self.min_deriv = 1e-3 + head = nn.Linear(hidden, dim * (3 * num_bins - 1)) + nn.init.zeros_(head.weight) # zero last layer → near-identity spline at init + nn.init.zeros_(head.bias) + self.net = nn.Sequential( + nn.Linear(dim, hidden), + nn.SiLU(), + nn.Linear(hidden, hidden), + nn.SiLU(), + head, + ) + + def _params(self, conditioned: Tensor) -> tuple[Tensor, Tensor, Tensor]: + p = self.net(conditioned).reshape(-1, self.dim, 3 * self.num_bins - 1) + k = self.num_bins + return p[..., :k], p[..., k : 2 * k], p[..., 2 * k :] + + def _apply_spline(self, x: Tensor, inverse: bool) -> tuple[Tensor, Tensor]: + uw, uh, ud = self._params(x * self.mask) + out, lad = _monotonic_rq_spline( + x.reshape(-1), + uw.reshape(-1, self.num_bins), + uh.reshape(-1, self.num_bins), + ud.reshape(-1, self.num_bins - 1), + inverse=inverse, + bound=self.bound, + min_bin=self.min_bin, + min_deriv=self.min_deriv, + ) + out = out.reshape_as(x) + active = 1 - self.mask + y = self.mask * x + active * out + logdet = (lad.reshape_as(x) * active).sum(dim=-1) + return y, logdet + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor]: + return self._apply_spline(x, inverse=False) + + def inverse(self, y: Tensor) -> Tensor: + return self._apply_spline(y, inverse=True)[0] + + +class NeuralSplineCouplingFlow(nn.Module): + """Rational-quadratic neural spline flow, exposed as an exact-likelihood energy. + + Coupling layers whose transform is a monotonic **rational-quadratic spline** + (Durkan et al. 2019) rather than an affine map — strictly more expressive per + layer than `AffineCouplingFlow`, so it fits sharp, multi-modal 2D densities + with fewer layers. Same contract as the affine flow: exact + ``log p(x) = log N(f(x); 0, I) + log|det ∂f/∂x|`` (no partition function), so + ``forward(x) = -log_prob(x)`` is a valid self-normalized `EnergyFn` with + ``log Z = 0``. Operates on vector data ``(B, dim)``. + + Each spline acts on ``[-bound, bound]`` with ``num_bins`` bins and slope-1 + linear tails outside it; the conditioner MLP is zero-initialized so training + starts from ≈identity. The spline only reshapes mass *inside* ``bound``, so + scale the data to sit within it (or widen ``bound``). Args mirror + `AffineCouplingFlow` plus ``num_bins`` / ``bound``. + """ + + def __init__( + self, + dim: int, + n_layers: int = 6, + hidden: int = 64, + *, + num_bins: int = 8, + bound: float = 3.0, + ): + super().__init__() + if dim < 2: + raise ValueError("NeuralSplineCouplingFlow needs dim >= 2 to split coordinates") + self.dim = dim + layers = [] + for i in range(n_layers): + mask = torch.zeros(dim) + mask[i % 2 :: 2] = 1.0 # alternate which half is conditioned + layers.append(_SplineCouplingLayer(dim, hidden, mask, num_bins, bound)) + self.layers = nn.ModuleList(layers) + + def transform(self, x: Tensor) -> tuple[Tensor, Tensor]: + """Map data ``x`` to the base space, returning ``(z, log|det ∂z/∂x|)``.""" + logdet = x.new_zeros(x.shape[0]) + for layer in self.layers: + x, ld = layer(x) + logdet = logdet + ld + return x, logdet + + def inverse(self, z: Tensor) -> Tensor: + """Map base samples ``z`` back to data space.""" + for layer in reversed(self.layers): + assert isinstance(layer, _SplineCouplingLayer) + z = layer.inverse(z) + return z + + def log_prob(self, x: Tensor) -> Tensor: + z, logdet = self.transform(x) + return standard_normal_logprob(z) + logdet + + def sample(self, n: int) -> Tensor: + z = torch.randn(n, self.dim, device=next(self.parameters()).device) + return self.inverse(z) + + def forward(self, x: Tensor) -> Tensor: + return -self.log_prob(x) # energy = -log p, so log Z = 0 exactly diff --git a/tests/test_nets.py b/tests/test_nets.py index dae495b..4cfffb5 100644 --- a/tests/test_nets.py +++ b/tests/test_nets.py @@ -189,3 +189,81 @@ def test_affine_coupling_flow_fits_a_gaussian(): samples = flow.sample(8000) assert (samples.mean(0) - mu).abs().max().item() < 0.2 assert (torch.cov(samples.T) - cov).abs().max().item() < 0.3 + + +def test_spline_flow_invertible_and_logdet_across_domain(): + # Run in double precision: the rational-quadratic map is *exactly* invertible, + # so this pins the math (float32 loses a few digits per layer, as all RQ-spline + # flows do). Perturb the zero-init net so the spline is genuinely nonlinear. + flow = ebm.nets.NeuralSplineCouplingFlow(dim=3, n_layers=4, num_bins=8, bound=3.0).double() + with torch.no_grad(): + for p in flow.parameters(): + p.add_(0.4 * torch.randn_like(p)) + x = (6 * torch.randn(6, 3)).double() # spans well past the [-3, 3] bound into the tails + z, logdet = flow.transform(x) + assert torch.allclose(flow.inverse(z), x, atol=1e-8) # exact roundtrip incl. tails + for i in range(len(x)): + jac = torch.autograd.functional.jacobian( + lambda v: flow.transform(v.unsqueeze(0))[0].squeeze(0), x[i] + ) + assert abs(logdet[i].item() - torch.linalg.slogdet(jac)[1].item()) < 1e-8 + assert torch.allclose(flow(x), -flow.log_prob(x)) + with pytest.raises(ValueError): + ebm.nets.NeuralSplineCouplingFlow(dim=1) + + +def test_spline_flow_fits_a_gaussian(): + import math + + d = 2 + # The spline reshapes only within [-bound, bound] (identity tails), so the + # target's mass must sit inside it — here a zero-centred Gaussian well within ±3. + mu = torch.zeros(d) + cov = torch.tensor([[0.7, 0.3], [0.3, 0.6]]) + chol = torch.linalg.cholesky(cov) + data = torch.randn(8000, d) @ chol.t() + mu + + flow = ebm.nets.NeuralSplineCouplingFlow(dim=d, n_layers=6, num_bins=8, bound=3.0) + opt = torch.optim.Adam(flow.parameters(), lr=5e-3) + for _ in range(2000): + batch = data[torch.randint(0, len(data), (256,))] + loss = -flow.log_prob(batch).mean() + opt.zero_grad() + loss.backward() + opt.step() + + test = torch.randn(4000, d) @ chol.t() + mu + prec = torch.linalg.inv(cov) + c = test - mu + analytic = ( + -0.5 * ((c @ prec) * c).sum(dim=1) + - 0.5 * d * math.log(2 * math.pi) + - 0.5 * torch.linalg.slogdet(cov)[1] + ) + assert (flow.log_prob(test) - analytic).abs().mean().item() < 0.2 + # It is practically invertible in float32 after training. + z, _ = flow.transform(test) + assert torch.allclose(flow.inverse(z), test, atol=1e-2) + + +def test_spline_flow_fits_two_moons(): + data = ebm.datasets.two_moons(8000) + flow = ebm.nets.NeuralSplineCouplingFlow(dim=2, n_layers=6, num_bins=8, bound=3.0) + opt = torch.optim.Adam(flow.parameters(), lr=5e-3) + for _ in range(2000): + batch = data[torch.randint(0, len(data), (256,))] + loss = -flow.log_prob(batch).mean() + opt.zero_grad() + loss.backward() + opt.step() + + samples = flow.sample(8000) + # Samples land on the data manifold (both crescents), far closer than a Gaussian. + to_data = ebm.eval.mmd(samples, data) + to_normal = ebm.eval.mmd(torch.randn(8000, 2), data) + assert to_data < 0.01 + assert to_data < 0.2 * to_normal + # Both arcs are covered (the moons split above/below the x-axis near the centre). + upper = ((samples[:, 0].abs() < 0.5) & (samples[:, 1] > 0.3)).float().mean() + lower = ((samples[:, 0].abs() < 0.5) & (samples[:, 1] < -0.0)).float().mean() + assert upper > 0.02 and lower > 0.02