Neural spline flows: NeuralSplineCouplingFlow (rational-quadratic)

[davidkhjo]

nets.NeuralSplineCouplingFlow — a rational-quadratic neural spline flow (Durkan et al. 2019).

What

• Monotonic rational-quadratic spline coupling layers dropped into the existing AffineCouplingFlow scaffold (same masks / transform / inverse / log_prob / sample / forward). Strictly more expressive per layer than affine, so it fits sharp, multi-modal 2D densities at a lower NLL for the same depth.
• Same self-normalized contract: exact log p(x) = log N(f(x);0,I) + log|det| (no partition function), so forward(x) = -log_prob(x) is a valid EnergyFn with log Z = 0.
• Slope-1 linear tails outside [-bound, bound], numerically stable inverse root 2c/(-b-√disc), zero-init conditioner head (≈identity at init). Defaults num_bins=8, bound=3, n_layers=6.

Validation (closed-form / distributional)

• Exact invertibility including the linear tails, and analytic log|det| == autograd Jacobian to 1e-8 (double precision — the RQ map is exactly invertible; float32 loses a few digits per layer as all spline flows do).
• Fits a Gaussian (log-density vs analytic) and the two-moons manifold (MMD, both crescents).

Notes

• The spline only reshapes mass inside bound (identity tails) — scale data to fit, or widen bound. Documented on the class and in the example.
• Exposed as ebm.nets.NeuralSplineCouplingFlow (alongside AffineCouplingFlow). torch-only; public API preserved; accumulates under the unreleased 0.15.0.

🤖 Generated with Claude Code