Analytic Bijections for Smooth and Interpretable Normalizing Flows

TL;DR

We address the challenge of designing expressive yet analytically invertible scalar bijections for normalizing flows by introducing globally smooth, unbounded, closed-form invertible bijections. The three analytic families—cubic rational, sinh, and cubic polynomial—together with a radial flow framework (including angular dependence via Fourier parametrization) provide drop-in replacements for affine or spline transforms and enable problem-specific designs that mitigate mode collapse. Across 1D density estimation, 2D coupling, and a φ^4 lattice field theory application, the methods match or exceed spline baselines while using far fewer parameters in radial architectures and offering interpretable, geometrically meaningful transformations. This work expands the normalizing flows design space, delivering stable training, smooth densities, and interpretable structure that scales to physics-inspired problems and beyond.

Abstract

A key challenge in designing normalizing flows is finding expressive scalar bijections that remain invertible with tractable Jacobians. Existing approaches face trade-offs: affine transformations are smooth and analytically invertible but lack expressivity; monotonic splines offer local control but are only piecewise smooth and act on bounded domains; residual flows achieve smoothness but need numerical inversion. We introduce three families of analytic bijections -- cubic rational, sinh, and cubic polynomial -- that are globally smooth ($C^\infty$), defined on all of $\mathbb{R}$, and analytically invertible in closed form, combining the favorable properties of all prior approaches. These bijections serve as drop-in replacements in coupling flows, matching or exceeding spline performance. Beyond coupling layers, we develop radial flows: a novel architecture using direct parametrization that transforms the radial coordinate while preserving angular direction. Radial flows exhibit exceptional training stability, produce geometrically interpretable transformations, and on targets with radial structure can achieve comparable quality to coupling flows with $1000\times$ fewer parameters. We provide comprehensive evaluation on 1D and 2D benchmarks, and demonstrate applicability to higher-dimensional physics problems through experiments on $φ^4$ lattice field theory, where our bijections outperform affine baselines and enable problem-specific designs that address mode collapse.

Figures (18)

• Figure 1: Three types of bijection behavior. Top row: Bijection functions $h(x)$ (solid) vs. identity (dashed). Bottom row: Resulting density (solid) when transforming a standard normal prior (dashed). Left (cubic rational): Local deformation creates bimodal structure while $h(x) \to x$ for large $|x|$. Thus, "stretching" in some parts is compensated by "compressing" in a nearby region. Middle (sinh): Local "stretching" displaces distant points by a constant offset. Right (affine): Uniform scaling compresses or expands the distribution while preserving its overall shape.
• Figure 2: 1D flow performance vs. number of stacked bijections. Both metrics tend to improve with more bijections, with cubic polynomial performing best overall. Error bars show one standard deviation over 6 random seeds.
• Figure 3: Training dynamics for 1D flows. Each panel shows reverse KL divergence vs. training steps for one bijection type, with different colors indicating the number of stacked bijections. All configurations converge stably, with more bijections generally achieving lower final loss.
• Figure 4: Transformation of points on a circle under a learned angular-dependent radial flow. Concentric circles (inset) transform smoothly into the spiral target structure. The color gradient indicates original radius, highlighting the preservation of radial ordering: paths never cross.
• Figure 5: Source distribution samples and where they are mapped to by the trained coupling flow (left pair) and radial Fourier flow (right pair). Colors indicate position of samples in the target spiral. The radial flow preserves radial structure in the source, while the coupling flow mixes the radial structure.