Source-linked AI summary

Neural Spline Flows

Conor Durkan, Artur Bekasov, Iain Murray, George Papamakarios

arXiv:1906.04032v2stat.MLcs.LG

TL;DR

Normalizing flows need flexible elementwise transformations that remain analytically invertible, but affine coupling layers can be less flexible than autoregressive alternatives. The paper introduces monotonic rational-quadratic spline modules as drop-in replacements, retaining analytic inversion while improving flow flexibility. Across density estimation, variational inference, and image modeling, neural spline flows provide stronger performance within this flexibility–practicality trade-off, with added implementation and training cost.

  • Problem

    Affine or additive transformations are easy to invert but can lack the flexibility needed to model complex, multi-modal, or discontinuous densities.

  • Method

    The paper replaces affine or additive transformations in coupling and autoregressive flows with fully differentiable monotonic rational-quadratic splines that have analytic inverses.

  • Results

    Neural spline flows improve density estimation and generative modeling while bringing coupling-layer performance on par with leading autoregressive flows in some settings.

  • Takeaways & Limitations

    Spline transforms provide a practical middle ground between flexibility and invertibility, including exact one-pass sampling for coupling-based models.

  • Takeaways & Limitations

    Rational-quadratic splines added approximately 30–40% to wallclock time per training update in the reported experiments, and their implementation is more involved.

Abstract

from arXiv · show

A normalizing flow models a complex probability density as an invertible transformation of a simple base density. Flows based on either coupling or autoregressive transforms both offer exact density evaluation and sampling, but rely on the parameterization of an easily invertible elementwise transformation, whose choice determines the flexibility of these models. Building upon recent work, we propose a fully-differentiable module based on monotonic rational-quadratic splines, which enhances the flexibility of both coupling and autoregressive transforms while retaining analytic invertibility. We demonstrate that neural spline flows improve density estimation, variational inference, and generative modeling of images.

1 Introduction

Normalizing flows provide invertible transformations for exact density evaluation and sampling, but existing coupling and autoregressive designs trade flexibility against inversion speed or analytic invertibility. The paper proposes rational-quadratic spline modules to improve this trade-off while retaining analytic inversion.

  • Flow-based models: Normalizing flows transform simple noise into data through an invertible differentiable function, supporting exact density evaluation and, in some cases, one-pass sampling.The density is computed through a change of variables involving the transformation’s Jacobian determinant.
  • Practical requirements: Flows require fast evaluation of the Jacobian determinant and inverse for density estimation, plus an analytic forward map for efficient sampling.Ideally, both directions require only a single neural-network pass.
  • Existing trade-offs: Autoregressive flows can be D times slower to invert than to evaluate, while more flexible variants may require numerical inversion.Coupling flows offer an analytic one-pass inverse but are often less flexible than autoregressive flows.
  • Proposed approach: The proposed fully differentiable rational-quadratic spline module replaces affine or additive transformations in coupling and autoregressive layers while retaining an analytic inverse.The module is designed to enhance flexibility without abandoning tractable inversion.
  • Proposed approach: The module is intended to enhance both coupling and autoregressive flows, potentially bringing coupling performance to the level of leading autoregressive flows.This addresses the flexibility–invertibility trade-off motivating the work.

2 Background

Coupling transforms provide tractable Jacobians and exact one-pass inversion, but their usual affine or additive elementwise functions are inflexible. The background develops spline-based alternatives and invertible linear transformations to increase expressiveness while preserving computational tractability.

  • Coupling transforms: Coupling transforms leave part of the input unchanged and transform the remaining dimensions in parallel using parameters produced by a neural network.The transformed dimensions therefore share conditioning on the unchanged subset.
  • Coupling transforms: Their Jacobian is lower triangular, making the determinant tractable, and their inverse can be computed exactly in a single pass.Multiple coupling layers can be composed to increase flexibility.
  • Invertible elementwise transformations: Affine and additive elementwise transformations are easy to invert but lack flexibility for modeling multi-modal or discontinuous densities.The affine form is gθi(xi) = αixi + βi, with additive transformations as the αi = 1 special case.
  • Spline alternatives: Polynomial and cubic splines increase expressiveness through monotonic piecewise polynomial segments, but earlier cubic-spline flows required bounded domains and numerically delicate cubic inversion.The cubic-spline construction used sigmoid and logit transformations to handle unconstrained inputs.
  • Invertible linear transformations: Invertible linear transformations can mix dimensions while preserving efficient determinant and inverse computation through LU decomposition.The determinant costs O(D), while triangular solves cost O(D2M), or W^-1 can be cached after an O(D3) setup.

3 Method

The method replaces affine or additive elementwise transforms with fully differentiable monotonic rational-quadratic splines that remain analytically invertible. These splines support both coupling and autoregressive flows while improving flexibility on complex densities.

  • 3.1 Monotonic rational-quadratic transforms: Monotonic rational-quadratic splines provide differentiable, analytically invertible elementwise transformations with linear tails outside [−B, B].Each bin uses a monotonic rational-quadratic function, and the overall spline maps [−B, B] to [−B, B].
  • 3.1 Monotonic rational-quadratic transforms: The spline uses K rational-quadratic functions separated by K+1 monotonically increasing knots spanning (−B, −B) to (B, B).Internal knot derivatives are parameterized positively, while boundary derivatives equal 1 to match the linear tails.
  • 3.1 Monotonic rational-quadratic transforms: A neural network outputs parameters that define positive bin widths, heights, and derivatives for each transformed dimension.Softmax outputs determine bin widths and heights over [−B, B], while softplus outputs parameterize derivatives.
  • 3.1 Monotonic rational-quadratic transforms: The transformation has a closed-form Jacobian determinant and an analytic inverse obtained by solving a quadratic equation.Elementwise monotonicity identifies the correct inverse root, preserving exact invertibility.
  • 3.1 Monotonic rational-quadratic transforms: The same spline parameterization adapts to autoregressive transforms, while coupling layers can transform all elements of x simultaneously.The paper defines the resulting coupling operation as yi = gθi(xi) for i = 1, . . . , D.
  • 3.2 Neural spline flows: With two coupling layers and K = 128 bins, rational-quadratic splines fit complex discontinuous synthetic densities that affine coupling struggles to model.The qualitative comparison concerns two-dimensional synthetic datasets.
  • 3.2 Neural spline flows: RQ-NSF combines spline transforms with alternating invertible linear transformations in coupling and autoregressive variants.RQ-NSF (C) uses coupling layers, while RQ-NSF (AR) uses autoregressive layers and retains an exact analytic inverse.

4 Related Work

Related work increases the flexibility of invertible transformations through linear layers, alternative elementwise functions, autoregressive networks, and continuous-time dynamics. These approaches differ in whether inversion is analytic, numerical, or approximate.

  • Invertible linear transformations: Glow replaces permutations with LU-decomposed invertible linear transformations, improving image modeling and inspiring related audio-flow architectures.WaveGlow and FloWaveNet adapt Glow-style invertible convolutions for audio generation.
  • Invertible elementwise transformations: Flow++ uses a mixture-of-logistics CDF for flexible coupling transforms but requires bisection because its inverse lacks a closed form.This contrasts with methods retaining analytic inversion.
  • Invertible elementwise transformations: NAF and Block-NAF increase autoregressive flexibility with monotonic neural networks, but their resulting models are not analytically invertible.Both methods replace simpler affine parameterizations with more flexible monotonic neural transformations.
  • Continuous-time flows: Continuous-time flows describe transformations with ordinary differential equations and may compute Jacobians approximately through truncated series and trace estimation.Deep diffeomorphic flow and Neural ODEs exemplify this alternative to discrete flow steps.

5 Experiments

Experiments evaluate neural spline flows for density estimation, variational autoencoders, and image generation. Results show strong density-estimation and image-modeling gains, while VAE improvements over other flow variants are limited.

  • 5.1 Density estimation of tabular data: RQ-NSF (C) and RQ-NSF (AR) achieve state-of-the-art normalizing-flow results on Power, Gas, and Hepmass.They tie Q-NSF variants on Power, and RQ-NSF (C) is competitive with the best autoregressive models.
  • 5.1 Density estimation of tabular data: RQ-NSF (C) closes the performance gap between coupling and autoregressive flows while retaining one-pass sampling.The results indicate that density-estimation performance need not require sacrificing one-pass sampling.
  • 5.2 Improving the variational autoencoder: All VAE models improve significantly over the standard baseline, but otherwise perform similarly with mostly overlapping error bars.The experiments use dynamically binarized MNIST and EMNIST and evaluate ELBO and importance-weighted log likelihood.
  • 5.3 Generative modeling of images: RQ-NSF (C) improves over the affine image-modeling baseline in three out of four tasks.The largest improvement occurs on 8-bit ImageNet64, while scores remain competitive with the original Glow model.
  • 5.3 Generative modeling of images: RQ-NSF (C) can reduce parameter counts relative to Glow by almost an order of magnitude in some image-generation settings.Figure 3 shows diverse, globally coherent samples that closely resemble real data.

6 Discussion

Neural spline flows improve flexibility and performance while retaining analytic invertibility and one-pass sampling. Their benefits are strongest when data availability is high, although the transforms add implementation and computation costs.

  • Interpretation: Differentiable spline-based coupling layers represent complex-shaped high-dimensional dependencies more quickly than affine-based flows when enough data are available.The discussion connects this advantage to the greater flexibility of spline transformations.
  • Empirical scope: RQ-NSF models excel on Power, Gas, and Hepmass, while remaining competitive but not superior when dimensionality rises without more data.The strongest tabular results occur on datasets with the highest data-point-to-dimensionality ratios.
  • Main conclusions: Monotonic spline transforms let coupling flows reach density-estimation performance on par with the best autoregressive flows while retaining exact one-pass sampling.This combines increased flexibility with the practical sampling advantage of coupling architectures.
  • Costs and implementation: Rational-quadratic splines add approximately 30–40% to wall-clock time per training update under the reported linear search implementation.Binary search can reduce binning to O(log2 K) time for K bins, while the method may require fewer flow steps.
  • Spline mechanics: The rational-quadratic transform is monotonic and elementwise, so its log absolute Jacobian determinant is computed as a sum across transformed inputs.Its inverse is obtained by solving a quadratic equation within the relevant spline bin.

B.1 Tabular density estimation

The tabular density-estimation appendix documents training, architecture, and hyperparameter choices for coupling and autoregressive flows. It uses validation-based model selection and dataset-specific adjustments to control optimization and overfitting.

  • Architecture choices: Higher-dimensional Hepmass and BSDS300 experiments use more coupling layers, whereas Miniboone does not because overfitting occurs with fewer data points.The number of coupling layers is therefore adjusted according to dataset characteristics.
  • Reported configurations: Coupling and autoregressive flows use distinct hyperparameter tables, with table 4 covering coupling layers and table 5 covering autoregressive layers.Table 6 reports validation log likelihoods in nats for the UCI datasets and BSDS300.
  • VAE configuration: VAE experiments use 32 latent features, 10 flow steps, and 8 bins for all neural spline flow models.Residual networks use 64 latent features for coupling layers and 128 for autoregressive layers.
  • Image configuration: Image experiments use Glow-like architectures with rational-quadratic or affine coupling transforms and train for 100,000 steps at 5-bit depth or 200,000 steps at 8-bit depth.Training uses cosine learning-rate annealing and a 1% development split.

C.1 Affine coupling transforms for 2D datasets

The appendix shows densities fitted by a model with two affine coupling layers on synthetic two-dimensional datasets.

  • Displayed results: Densities fitted by a model with two affine coupling layers on synthetic two-dimensional datasets are shown in figure 4.The passage identifies the model and dataset setting but does not report a quantitative comparison.

C.2 Samples

The appendix collects qualitative samples from the VAE and generative image-modeling experiments, including synthetic two-dimensional results and additional image outputs.

  • Synthetic datasets: Figure 4 presents qualitative results for two-dimensional synthetic datasets using two affine coupling layers.The caption identifies the architecture and evaluation setting without stating an outcome.
  • VAE samples: Figure 5 compares training data with RQ-NSF coupling and autoregressive VAE samples.The caption specifies the top-to-bottom ordering of the displayed rows.
  • Additional samples: Figure 6 contains additional image samples for the generative image-modeling experiments.The passage does not specify an ordering or quantitative result for these samples.
Loading 1906.04032v2…