Neural spline flows: NeuralSplineCouplingFlow (rational-quadratic)

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

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`

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

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

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