Analytic Bijections for Smooth and Interpretable Normalizing Flows

AI summary

Authors: Mathis Gerdes, Miranda C. N. Cheng Category: Machine Learning arXiv: 2601.10774v2

Overview

Research area: Machine learning — specifically normalizing flows, the family of generative models that learn probability distributions by transforming a simple base density (typically Gaussian) through invertible maps.

Technical level: Advanced. The paper assumes familiarity with normalizing flows, coupling layers, Jacobians, and KL divergence, and it uses algebraic tools such as Cardano's formula for solving cubics and the Abel–Ruffini theorem. The core ideas are geometrically intuitive, but the constructions are mathematically dense.

Scope: The paper introduces three families of scalar bijections that are globally smooth, defined on all of the real line, and invertible in closed form, and then builds a new "radial flow" architecture on top of them, evaluating all of it on 1D and 2D benchmarks, image and tabular density estimation, and 20×20 lattice field theory.

What This Paper Is About

Normalizing flows need invertible scalar functions that warp one dimension of the data, and the standard choices each give something up. Affine maps are smooth and easy to invert but can only shift and scale. Monotonic splines give fine local control but are only piecewise smooth and only act on a bounded interval. Residual flows are smooth but require numerical root-finding to invert.

The paper's goal is to construct scalar bijections that have all the good properties at once: infinitely differentiable everywhere, defined on all of the real line, invertible in closed form, with a tractable Jacobian, and expressive enough to make local, nonlinear deformations of the input. The authors then exploit those properties to design a new architecture, radial flows, that transforms only the distance from a center while leaving direction untouched.

Key Contributions

1. Three parametric families of analytic bijections. The cubic rational bijection (derived from rational functions whose inverse reduces to a solvable cubic), the sinh conjugation bijection, and the cubic conjugation bijection. All are C^∞, defined on all of ℝ, analytically invertible, and locally expressive.

2. Two construction principles. Algebraic rational functions of the form h(x) = x + g(x) where clearing denominators yields a cubic that can be solved in closed form; and conjugation with monotonic functions, h(x) = g⁻¹(g(x) + δ), with the requirement that g be superlinear so that h(x) → x at large |x|.

3. Radial flow architecture with direct parametrization. A layer that applies a scalar bijection to the radial coordinate r = ‖x‖ relative to a learnable center while preserving the angular direction, with a closed-form log-Jacobian. The parameters are learned directly rather than output by a neural network conditioner.

4. Demonstrated drop-in compatibility and a physics application. The bijections replace affine or spline maps inside existing coupling architectures at comparable memory and runtime cost, and support a problem-specific zero-mode bijection that prevents mode collapse in φ⁴ lattice field theory.

Main Findings

• Three bijections meet all five desiderata. The paper's Table 1 compares methods on smoothness, global domain, closed-form inverse, and local deformation capability: affine maps fail local expressivity, splines are only C^k and act on a bounded domain, residual flows need numerical inversion, and the three proposed families satisfy all criteria. Bijectivity of the cubic rational form requires −1 < λ < 8 and σ > 0.

• 1D stacking improves with depth. Training 1D flows with stacks of N ∈ {3, 9, 27, 128, 256} scalar bijections, averaged over 6 seeds, all three bijection types improve monotonically with depth. At N = 27, cubic conjugation reaches an effective sample size of approximately 99% and forward KL divergence of approximately 3.5 × 10⁻³.

• Cubic conjugation beats affine and spline baselines in 2D coupling flows. On a 2D spiral with 12-layer coupling flows using ResNet conditioners and stack count N per layer, cubic conjugation at N = 9 reaches D_KL ≈ 0.35 versus affine at D_KL ≈ 0.8 and rational quadratic splines with 8 knots at D_KL ≈ 0.45. Beyond N = 9, added expressivity reduces training stability and increases variance in final performance.

• Radial flows improve with more centers and deeper stacks, then destabilize. Sweeping centers L ∈ {5, 10, 20, 40} and stacked bijections per center N ∈ {1, 2, 8, 16, 32}, forward KL improves with both knobs, but very large stack counts degrade performance due to training instability at fixed hyperparameters.

• Fourier radial flows are parameter-efficient. A single angle-dependent radial layer with N = 9 sinh bijections saturates in angular terms: test NLLs at K ∈ {0, 1, 2, 3} are {−0.09, −0.61, −0.69, −0.74}, and increasing depth to N = 32 at K = 2 reaches −0.79. At K = 3 the model uses 319 parameters (2K+1 = 7 terms per bijection parameter); the constant-term-only model uses 49 parameters.

• Comparable quality with roughly three orders of magnitude fewer parameters. A single-layer Fourier radial flow with N = 32 sinh bijections and K = 2 reaches test NLL −0.79 versus −0.52 for the RealNVP baseline on the spiral, using roughly three orders of magnitude fewer parameters. The paper also states that on targets with radial structure, radial flows can reach comparable quality to coupling flows with 1000× fewer parameters.

• Radial flows avoid axis-aligned folding artifacts. Coupling flows exhibit thin lines of high probability produced by iterative axis-aligned transformations. The radial flow preserves radial organization in the source distribution: inner spiral points map near the origin while outer points sit at larger radii, whereas the affine coupling flow mixes these regions.

• Learnable centers work. The radial flow's center was initialized at (−0.5, −1) and training moved it toward the spiral center, showing the center need not be tuned by hand.

• Multimodal targets favor radial flows too. On a 5-component Gaussian mixture arranged in a circle, with 32 centers and a stack of 12 cubic bijections, the pure angle-independent radial flow (1.6k parameters) achieves the best visual fidelity, while the coupling flow (2,311k parameters) shows spiky artifacts. Angle-dependent radial flows with 5 Fourier terms (K = 2) also show more significant artifacts than the pure radial variant.

• CIFAR10 gains as a drop-in replacement. Replacing the affine bijection in a multi-scale Real NVP-style flow with a stack of 8 analytic bijections improves test bits per dimension from the published 3.49 for RealNVP to 3.36 (cubic rational), 3.37 (sinh conjugation), and 3.37 (cubic conjugation) — roughly 0.12 bpd at matched architecture.

• UCI results are mixed and complementary. In the coupling architecture of Durkan et al., spline+ (sinh conjugation followed by a rational-quadratic spline) bests the reproduced spline baseline on POWER (0.64 ± 0.01), GAS (12.92 ± 0.05), and BSDS300 (157.57 ± 0.24), matching or exceeding published RQ-NSF(C) on POWER and BSDS300. It underperforms on HEPMASS and MINIBOONE, consistent with small-dataset overfitting; on MINIBOONE the pure sinh variant is the strongest of all methods (−9.58 ± 0.23).

• φ⁴ lattice field theory shows persistent gains at 400 dimensions. On a 20×20 lattice with m² = −4 and λ = 4.807 (correlation length ξ ≈ L/4), 12-layer coupling flows with 8 stacked analytic bijections achieve final ESS ordering: cubic rational 39.66% > cubic 38.85% > sinh 38.51% > spline 34.34% > affine 31.85%.

• A zero-mode bijection fixes mode collapse. In the bimodal, ℤ₂-symmetric regime, naive reverse-KL training reaches high ESS (90%) while collapsing to a single mode. Pretraining a scalar bijection on the magnitude of the zero-frequency Fourier mode (a stack of 8 cubic conjugation bijections, with coupling layers frozen) captures both modes and matches MCMC ground truth.

Methodology in Plain English

The authors start by asking what a scalar bijection must do to be useful, and derive a checklist: be infinitely smooth, be defined everywhere on the real line, have a closed-form inverse, have a cheap Jacobian, and be able to bend locally rather than just rescale.

Building the bijections. One route is to write the function as identity plus a correction: h(x) = x + g(x), where g vanishes at large |x| so the map acts like a local bump and inherits the prior's tails. Writing g as a ratio of polynomials forces the denominator to have no real roots and therefore even degree, and forces the numerator to be lower degree than the denominator. Degrees of four or higher give a quintic or worse equation with no closed-form solution, and degree zero gives only affine maps — so degree two is the only nontrivial choice, and clearing denominators yields a cubic solvable with Cardano's formula. That is the cubic rational family. Their chosen specialization is h(x) = x + λ(x − γ) / (1 + (x − γ)²/σ²).

The second route is conjugation: pick a strictly monotonic function g with a known inverse and define h(x) = g⁻¹(g(x) + δ). This is bijective for any δ, with derivative h′(x) = g′(x)/g′(h(x)). To make h(x) → x at large |x|, g must grow faster than linearly. Using g = sinh gives the sinh conjugation family, which adds global shift parameters μ and ν alongside the local deformation parameter δ. Using g(x) = ax + bx³ with a, b > 0 gives the cubic conjugation family.

Constraining the parameters. Because the parameters must live in valid ranges, the authors apply softplus to enforce positivity (σ, a, b) and sigmoid to enforce bounded intervals (λ, with −1 < λ < 8). They also shrink the initialization scale of the conditioner network's final layer, since large initial parameters push bijections toward extreme transformations and destabilize training.

Using them. Expressivity is increased not by adding knots but by stacking — composing N independently parameterized copies into one composite bijection. These stacks drop into autoregressive or coupling layers in place of a spline or affine map.

Radial flows. The second half of the paper decomposes any nonzero point as x = r x̂ and appends a layer that applies a scalar bijection f to the radius while keeping the unit direction fixed: g(x) = (f(‖x‖)/‖x‖) x. The log-determinant has a clean closed form, log|det J_g| = log|f′(r)| + (n − 1) log|f(r)/r|. A learnable center c, optional per-dimension scaling s, and the constraint f(0) = 0 (enforced by subtracting f̃(0)) complete the layer. Because the parameters can be optimized directly rather than predicted by a network, radial flows train stably at learning rates of 10⁻², an order of magnitude higher than coupling flows. Making the bijection angle-dependent requires care at the origin: the unit vector is undefined at the origin, and differentiability requires ∂_r f(0, x̂) to be a direction-independent constant. The authors note they did not enforce this in experiments and it did not improve training for few-layer flows, though for models with many layers it is worth considering; it can even be applied post-training.

Evaluation. The authors measure performance with forward KL (negative log-likelihood when target samples are available) and reverse KL (when only an unnormalized target density is available, as in physics). For low-dimensional benchmarks they numerically estimate the partition function and entropy so they can plot absolute KL divergence against the optimum at D_KL = 0.

Why This Matters

The paper matters because it removes a long-standing either/or in flow design. Previously, you could have smoothness or local expressivity or closed-form inversion, but not all three. Having all three at once means smoother learned densities, faster and more reliable training, and the ability to interpret what a layer is actually doing geometrically.

Impact on research:

• Provides a general recipe (perturbation-plus-cubic, or conjugation with a superlinear monotonic function) rather than a single fixed bijection, so the construction can be extended.
• Connects normalizing flows in Cartesian coordinates to polar decomposition and to spherical coordinates, and notes that angle-dependent radial flows are a special case of coupling flows in spherical coordinates with the radius as the only active coordinate — a connection the authors flag for future work.
• Demonstrates that bijection choice interacts with architecture: radial flows make a 400-dimensional physics problem tractable with a fraction of the parameters, and an expressive bijection enabled a symmetry-preserving fix for mode collapse.

Real-world applications supported by this class of work:

• Scientific simulation, particularly accelerating Monte Carlo sampling in lattice field theory.
• Density estimation for high-dimensional tabular and image data, where a better scalar bijection translates into a measurable gain at no architectural cost.
• Anomaly detection and likelihood-based scoring, where smooth densities avoid spurious high-probability regions.
• Generative modeling on constrained hardware, since the Fourier radial variant captures structure with 50–300 parameters rather than millions.

Industry relevance: Anywhere a generative model must be trained reliably and its failures diagnosed, this work is relevant. The training stability at learning rates an order of magnitude higher reduces tuning effort. The interpretability of radial layers — you can inspect a single angle-dependent bijection f(r, φ) rather than a stack of black-box transformations — is useful for debugging. And the parameter efficiency matters for deployment on edge devices and for scientific workflows where each training run is expensive.

Future Directions

• Radial flows need to scale beyond low dimensions. The authors state that while the φ⁴ experiments validate coupling flows, radial flows remain limited to low dimensions. A single radial layer preserves rays, so probability mass can be redistributed along rays from the center but not between them; stacking layers with different learned centers mitigates this, but the number of layers required may scale adversely with dimension. The paper suggests radial flows may be combined with normalizing flows on the sphere.

• Extending angle dependence beyond 2D. Angle dependence is parametrized in 2D via a truncated Fourier series in the polar angle. The paper states this could be extended to higher dimensions using spherical harmonics.

• Full spherical-coordinate coupling flows. The authors observe that angle-dependent radial flows are a special case of coupling flows in spherical coordinates with the radius as the only active coordinate, and suggest future work updating subsets of angle coordinates conditioned on the radius and frozen angles.

• Completing hyperparameter characterization. The available text notes that the authors did not extensively tune hyperparameters, and that the radial flow instabilities at large stack counts may be addressable by tuning learning rate and other hyperparameters that were held fixed across configurations.

• Smooth conditioners for coupling layers. The conclusion notes that for coupling layers, the conditioner network must also be smooth to guarantee overall smoothness of the flow, and must not use activations such as ReLU — an implementation requirement worth developing guidance around.

Target Audience

This paper is for machine learning researchers working on generative models, especially those who care about the design of individual flow layers and about the tradeoffs between expressivity, smoothness, and inversion cost. It also speaks directly to physicists and computational scientists using normalizing flows as samplers for lattice field theory and other high-dimensional scientific targets, who will find the φ⁴ experiments and the zero-mode bijection most relevant. Practitioners in density estimation who want a drop-in upgrade for affine or spline layers will find the CIFAR10 and UCI tables useful, though those results are mixed rather than uniformly better. Readers without a background in normalizing flows, coupling layers, or basic calculus notation will find the middle sections difficult, even though the opening framing is accessible.

Authors’ abstract

A key challenge in normalizing flows is finding expressive invertible scalar bijections. 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 that are globally smooth ($C^\infty$), defined on all of $\mathbb{R}$, and analytically invertible in closed form, combining the favorable properties of prior approaches. Beyond serving as drop-in replacements in coupling flows, where they match or exceed spline performance, 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.