Neural Spline Flows

  • 2019-06-10 14:43:23
  • Conor Durkan, Artur Bekasov, Iain Murray, George Papamakarios
  • 27

Abstract

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

 

Quick Read (beta)

Neural Spline Flows

Conor Durkan  Artur Bekasov11footnotemark: 1  Iain Murray  George Papamakarios
School of Informatics, University of Edinburgh
{conor.durkan, artur.bekasov, i.murray, g.papamakarios}@ed.ac.uk
Equal contribution
Abstract

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.

plus 2pt minus 2pt

 

Neural Spline Flows


  Conor Durkanthanks: Equal contribution  Artur Bekasov11footnotemark: 1  Iain Murray  George Papamakarios School of Informatics, University of Edinburgh {conor.durkan, artur.bekasov, i.murray, g.papamakarios}@ed.ac.uk

\@float

noticebox[b]Preprint. Under review.\[email protected]

1 Introduction

Models that can reason about the joint distribution of high-dimensional random variables are central to modern unsupervised machine learning. Explicit density evaluation is required in many statistical procedures, while synthesis of novel examples can enable agents to imagine and plan in an environment prior to choosing a action. In recent years, the variational autoencoder [VAE, Kingma and Welling, 2014, Rezende et al., 2014] and generative adversarial network [GAN, Goodfellow et al., 2014] have received particular attention in the generative-modeling community, and both are capable of sampling with a single forward pass of a neural network. However, these models do not offer exact density evaluation, and can be difficult to train. On the other hand, autoregressive density estimators [Uria et al., 2013, Germain et al., 2015, van den Oord et al., 2016c, b, Salimans et al., 2017, van den Oord et al., 2016a] can be trained by maximum likelihood, but sampling requires a sequential loop over the output dimensions.

Flow-based models present an alternative approach to the above methods, and in some cases provide both exact density evaluation and sampling in a single neural-network pass. A normalizing flow models data 𝐱 as the output of an invertible, differentiable transformation 𝐟 of noise 𝐮:

𝐱=𝐟(𝐮)where𝐮π(𝐮). (1)

The probability density of 𝐱 under the flow is obtained by a change of variables:

p(𝐱)=π(𝐟-1(𝐱))|det(𝐟-1𝐱)|. (2)
Figure 1: Monotonic rational-quadratic transforms are drop-in replacements for additive or affine transformations in coupling or autoregressive layers, greatly enhancing their flexibility while retaining exact invertibility. Left: A random monotonic rational-quadratic transform with K=10 bins and linear tails is parameterized by a series of K+1 ‘knot’ points in the plane, and the K-1 derivatives at the internal knots. Right: Derivative of the transform on the left with respect to x. Monotonic rational-quadratic splines naturally induce multi-modality when used to transform random variables.

Intuitively, the function 𝐟 compresses and expands the density of the noise distribution π(𝐮), and this change is quantified by the determinant of the Jacobian of the transformation. The noise distribution π(𝐮) is typically chosen to be simple, such as a standard normal, whereas the transformation 𝐟 and its inverse 𝐟-1 are often implemented by composing a series of invertible neural-network modules. Given a dataset 𝒟={𝐱(n)}n=1N, the flow is trained by maximizing the total log likelihood nlogp(𝐱(n)) with respect to the parameters of the transformation 𝐟. In recent years, normalizing flows have received widespread attention in the machine-learning literature, seeing successful use in density estimation [Dinh et al., 2017, Papamakarios et al., 2017], variational inference [Kingma et al., 2016, Rezende and Mohamed, 2015, Louizos and Welling, 2017, van den Berg et al., 2018], image, audio and video generation [Kingma and Dhariwal, 2018, Kim et al., 2018, Prenger et al., 2018, Kumar et al., 2019], likelihood-free inference [Papamakarios et al., 2019], and learning maximum-entropy distributions [Loaiza-Ganem et al., 2017].

A flow is defined by specifying the bijective function 𝐟 or its inverse 𝐟-1, usually with a neural network. Depending on the flow’s intended use cases, there are practical constraints in addition to formal invertibility:

  • To train a density estimator, we need to be able to evaluate the Jacobian determinant and the inverse function 𝐟-1 quickly. We don’t evaluate 𝐟, so the flow is usually defined by specifying 𝐟-1.

  • If we wish to draw samples using creftype 1, we would like 𝐟 to be available analytically, rather than having to invert 𝐟-1 with iterative or approximate methods.

  • Ideally, we would like both 𝐟 and 𝐟-1 to require only a single pass of a neural network to compute, so that both density evaluation and sampling can be performed quickly.

Autoregressive flows such as inverse autoregressive flow [IAF, Kingma et al., 2016] or masked autoregressive flow [MAF, Papamakarios et al., 2017] are D times slower to invert than to evaluate, where D is the dimensionality of 𝐱. Subsequent work which enhances their flexibility has resulted in models which do not have an analytic inverse, and require numerical optimization to invert [Huang et al., 2018]. Flows based on coupling layers [NICE, RealNVP, Dinh et al., 2015, 2017] have an analytic one-pass inverse, but are often less flexible than their autoregressive counterparts.

In this work, we propose a fully-differentiable module based on monotonic rational-quadratic splines which has an analytic inverse. The module acts as a drop-in replacement for the affine or additive transformations commonly found in coupling and autoregressive transforms. We demonstrate that this module significantly enhances the flexibility of both classes of flows, and in some cases brings the performance of coupling transforms on par with the best-known autoregressive flows. An illustration of our proposed transform is shown in creftype 1.

2 Background

2.1 Coupling transforms

A coupling transform ϕ [Dinh et al., 2015] maps an input 𝐱 to an output 𝐲 in the following way:

  1. 1.

    Split the input 𝐱 into two parts, 𝐱=[𝐱1:d-1,𝐱d:D].

  2. 2.

    Compute parameters 𝜽=NN(𝐱1:d-1), where NN is an arbitrary neural network.

  3. 3.

    Compute yi=g𝜽i(xi) for i=d,,D in parallel, where g𝜽i is an invertible function parameterized by 𝜽i.

  4. 4.

    Set 𝐲1:d-1=𝐱1:d-1, and return 𝐲=[𝐲1:d-1,𝐲d:D].

The Jacobian matrix of a coupling transform is lower triangular, since 𝐲d:D is given by transforming 𝐱d:D elementwise as a function of 𝐱1:d-1, and 𝐲1:d-1 is equal to 𝐱1:d-1. Thus, the Jacobian determinant of the coupling transform ϕ is given by det(ϕ𝐱)=i=dDg𝜽ixi, the product of the diagonal elements of the Jacobian.

Coupling transforms solve two important problems for normalizing flows: they have a tractable Jacobian determinant, and they can be inverted exactly in a single pass. The inverse of a coupling transform can be easily computed by running steps 1–4 above, this time inputting 𝐲, and using g𝜽i-1 to compute 𝐱d:D in step 3. Multiple coupling layers can also be composed in a natural way to construct a normalizing flow with increased flexibility. A coupling transform can also be viewed as a special case of an autoregressive transform where we perform two splits of the input data instead of D, as noted by Papamakarios et al. [2017]. In this way, advances in flows based on coupling transforms can be applied to autoregressive flows, and vice versa.

2.2 Invertible elementwise transformations

Affine/additive

Typically, the function g𝜽i takes the form of an additive [Dinh et al., 2015] or affine [Dinh et al., 2017] transformation for computational ease. The affine transformation is given by:

g𝜽i(xi)=αixi+βi,where𝜽i={αi,βi}, (3)

and αi is usually constrained to be positive. The additive transformation corresponds to the special case αi=1. Both the affine and additive transformations are easy to invert, but they lack flexibility. Recalling that the base distribution of a flow is typically simple, flow-based models may struggle to model multi-modal or discontinuous densities using just affine or additive transformations, since they may find it difficult to compress and expand the density in a suitably nonlinear fashion (for an illustration, see creftype C.1). We aim to choose a more flexible g𝜽i, that is still differentiable and easy to invert.

Polynomial splines

Recently, Müller et al. [2018] proposed a powerful generalization of the above affine transformations, based on monotonic piecewise polynomials. The idea is to restrict the input domain of g𝜽i to the interval [0,1], partition the input domain into K bins, and define g𝜽i to be a simple polynomial segment within each bin. Müller et al. [2018] restrict themselves to monotonically-increasing linear and quadratic polynomial segments, whose coefficients are parameterized by 𝜽i. Moreover, the polynomial segments are restricted to match at the bin boundaries so that g𝜽i is continuous. Functions of this form, which interpolate between data using piecewise polynomials, are known as polynomial splines.

Cubic splines

In a previous iteration of this work [Durkan et al., 2019], we explored the cubic-spline flow, a natural extension to the framework of Müller et al. [2018]. We proposed to implement g𝜽i as a monotonic cubic spline Steffen [1990], where g𝜽i is defined to be a monotonically-increasing cubic polynomial in each bin. By composing coupling layers featuring elementwise monotonic cubic-spline transforms with invertible linear transformations, we found flows of this type to be much more flexible than the standard coupling-layer models in the style of RealNVP [Dinh et al., 2017], achieving similar results to autoregressive models on a suite of density-estimation tasks.

Like Müller et al. [2018], our spline transform and its inverse were defined only on the interval [0,1]. To ensure that the input is always between 0 and 1, we placed a sigmoid transformation before each coupling layer, and a logit transformation after each coupling layer. These transformations allow the spline transform to be composed with linear layers, which have an unconstrained domain. However, the limitations of 32-bit floating point precision mean that in practice the sigmoid saturates for inputs outside the approximate range of [-13,13], which results in numerical difficulties. In addition, computing the inverse of the transform requires inverting a cubic polynomial, which is prone to numerical instability if not carefully treated [Blinn, 2007]. In creftype 3.1 we propose a modified method based on rational-quadratic splines which overcomes these difficulties.

2.3 Invertible linear transformations

To ensure all input variables can interact with each other, it is common to randomly permute the dimensions of intermediate layers in a normalizing flow. Permutation is an invertible linear transformation, with absolute determinant equal to 1. Kingma and Dhariwal [2018] generalize permutations to a more general class of linear transformations, demonstrating improvements on a range of image tasks. In particular, they parameterize a linear transformation with matrix 𝐖 in terms of its LU-decomposition 𝐖=𝐏𝐋𝐔, where 𝐏 is a fixed permutation matrix, 𝐋 is lower triangular with ones on the diagonal, and 𝐔 is upper triangular. By restricting the diagonal elements of 𝐔 to be positive, 𝐖 is guaranteed to be invertible.

By making use of the LU-decomposition, both the determinant and the inverse of the linear transformation can be computed efficiently. First, the determinant of 𝐖 can be calculated in 𝒪(D) time as the product of the diagonal elements of 𝐔. Second, inverting the linear transformation can be done by solving two triangular systems, one for 𝐔 and one for 𝐋, each of which costs 𝒪(D2M) time where M is the batch size. Alternatively, we can pay a one-time cost of 𝒪(D3) to explicitly compute 𝐖-1, which can then be cached for re-use.

3 Method

3.1 Monotonic rational-quadratic transforms

We propose to implement the function g𝜽i using monotonic rational-quadratic splines as a building block, where each bin is defined by a monotonically-increasing rational-quadratic function. A rational-quadratic function takes the form of a quotient of two quadratic polynomials. Rational-quadratic functions are easily differentiable, more flexible than a polynomial in that they have an infinite Taylor-series expansion, and since we consider only monotonic segments of these functions, they are also analytically invertible. In our implementation, we use the method of Gregory and Delbourgo [1982] to parameterize a monotonic rational-quadratic spline. The spline itself maps an interval [-B,B] to [-B,B]. We define the transformation outside this range as the identity, resulting in linear ‘tails’, so that the overall transformation can take unconstrained inputs.

The spline uses K different rational-quadratic functions, with boundaries set by K+1 coordinates {(x(k),y(k))}k=0K known as knots. The knots monotonically increase between (x(0),y(0))=(-B,-B) and (x,(K)y)(K)=(B,B). We give the spline K-1 arbitrary positive values for the derivatives at the internal points, and set the boundary derivatives to 1 to match the linear tails, which we found to be important for stable training. Then, the method constructs a monotonic, continuously-differentiable, rational-quadratic spline which passes through the knots, with the given derivatives at the knots. The transform is illustrated in creftype 1, and creftype A.1 gives full details.

Implementation

The practical implementation of the monotonic rational-quadratic coupling transform is as follows:

  1. 1.

    A neural network NN takes 𝐱1:d-1 as input and outputs an unconstrained parameter vector 𝜽i of length 3K-1 for each i=d,,D.

  2. 2.

    Vector 𝜽i is partitioned as 𝜽i=[𝜽iw,𝜽ih,𝜽id], where 𝜽iw and 𝜽ih have length K, and 𝜽id has length K-1.

  3. 3.

    Vectors 𝜽iw and 𝜽ih are each passed through a softmax and multiplied by 2B; the outputs are interpreted as the widths and heights of the K bins, which must be positive and span the [-B,B] interval. Cumulative sums of the K bin widths and heights, each starting at -B, yield the K+1 knots {(x(k),y(k))}k=0K.

  4. 4.

    The vector 𝜽id is passed through a softplus function and is interpreted as the values of the derivatives at the internal knots.

Evaluating a rational-quadratic spline transform at location x requires finding the bin in which x lies, which can be done efficiently with binary search, since the bins are sorted. The Jacobian determinant can be computed in closed-form as a product of quotient derivatives, while the inverse requires solving a quadratic equation whose coefficients depend on the value to invert; we provide details of these procedures in creftype A.2 and creftype A.3. Unlike the additive and affine transformations, which have limited flexibility, a monotonic spline with sufficiently many bins can approximate any continuous monotonic function on the specified interval [-B,B], yet has a closed-form, tractable Jacobian determinant, and can be inverted analytically. Finally, our parameterization is fully-differentiable, which allows for training by gradient methods.

The above formulation can also easily be adapted for autoregressive transforms; each 𝜽i can be computed as a function of 𝐱1:i-1 using an autoregressive neural network, and then all elements of 𝐱 can be transformed at once. Inspired by this, we also introduce a set of splines for our coupling layers which act elementwise on 𝐱1:d-1 (the typically non-transformed variables), and whose parameters are optimized directly by stochastic gradient descent. This means that our coupling layer transforms

Training data Flow density Flow samples
Figure 2: Qualitative results for two-dimensional synthetic datasets using RQ-NSF with two coupling layers.

all elements of 𝐱 at once as follows:

𝜽1:d-1 =Trainable parameters (4)
𝜽d:D =NN(𝐱1:d-1) (5)
yi =g𝜽i(xi)for i=1,,D. (6)

creftypecap 2 demonstrates the flexibility of our rational-quadratic coupling transform on synthetic two-dimensional datasets. Using just two coupling layers, each with K=128 bins, the monotonic rational-quadratic spline transforms have no issue fitting complex, discontinuous densities with potentially hundreds of modes. In contrast, a coupling layer with affine transformations has significant difficulty with these tasks (see creftype C.1).

3.2 Neural spline flows

The monotonic rational-quadratic spline transforms described in the previous section act as drop-in replacements for affine or additive transformations in both coupling and autoregressive transforms. When combined with alternating invertible linear transformations, we refer to the resulting class of normalizing flows as rational-quadratic neural spline flows (RQ-NSF), which may feature coupling layers, RQ-NSF (C), or autoregressive layers, RQ-NSF (AR). RQ-NSF (C) corresponds to Glow [Kingma and Dhariwal, 2018] with affine or additive transformations replaced with monotonic rational-quadratic transforms, where Glow itself is exactly RealNVP with permutations replaced by invertible linear transformations. RQ-NSF (AR) corresponds to either IAF or MAF, depending on whether the flow parameterizes 𝐟 or 𝐟-1, again with affine transformations replaced by monotonic rational-quadratic transforms, and also with permutations replaced with invertible linear layers. Overall, RQ-NSF resembles a traditional feed-forward neural network architecture, alternating between linear transformations and elementwise non-linearities, while retaining an exact, analytic inverse. In the case of RQ-NSF (C), the inverse is available in a single neural-network pass.

4 Related Work

Invertible linear transformations

Invertible linear transformations have found significant use in normalizing flows. Glow [Kingma and Dhariwal, 2018] replaces the permutation operation of RealNVP with an LU-decomposed linear transformation interpreted as a 1×1 convolution, yielding superior performance for image modeling. WaveGlow [Prenger et al., 2018] and FloWaveNet [Kim et al., 2018] have also successfully adapted Glow for generative modeling of audio. Expanding on the invertible 1×1 convolution presented in Glow, Hoogeboom et al. [2019] propose the emerging convolution, based on composing autoregressive convolutions in a manner analogous to an LU-decomposition, and the periodic convolution, which uses multiplication in the Fourier domain to perform convolution. Hoogeboom et al. [2019] also introduce linear transformations based on the QR-decomposition, where the orthogonal matrix is parameterized by a sequence of Householder transformations [Tomczak and Welling, 2016].

Invertible elementwise transformations

Outside of those discussed in creftype 2.2, there has been much recent work in developing more flexible invertible elementwise transformations for normalizing flows. Flow++ [Ho et al., 2019] uses the CDF of a mixture of logistic distributions as a monotonic transformation in coupling layers, but requires bisection search to compute an inverse, since a closed form is not available. Non-linear squared flow [Ziegler and Rush, 2019] adds an inverse-quadratic perturbation to an affine transformation in an autoregressive flow, which is invertible under certain restrictions of the parameterization. Computing this inverse requires solving a cubic polynomial, and the overall transform is less flexible than a monotonic rational-quadratic spline. Sum-of-squares polynomial flow [SOS, Jaini et al., 2019] parameterizes a monotonic transformation by specifying the coefficients of a polynomial of some chosen degree which can be written as a sum of squares. For low-degree polynomials, an analytic inverse may be available, but the method would require an iterative solution in general.

Neural autoregressive flow [NAF, Huang et al., 2018] replaces the affine transformation in MAF by parameterizing a monotonic neural network for each dimension. This greatly enhances the flexibility of the transformation, but the resulting model is again not analytically invertible. Block neural autoregressive flow [Block-NAF, De Cao et al., 2019] directly fits an autoregressive monotonic neural network end-to-end rather than parameterizing a sequence for each dimension as in NAF, but is also not analytically invertible.

Continuous-time flows

Rather than constructing a normalizing flow as a series of discrete steps, it is also possible to use a continuous-time flow, where the transformation from noise 𝐮 to data 𝐱 is described by an ordinary differential equation. Deep diffeomorphic flow [Salman et al., 2018] is one such instance, where the model is trained by backpropagation through an Euler integrator, and the Jacobian is computed approximately using a truncated power series and the Hutchison trace estimator [Hutchinson, 1990]. Neural ordinary differential equations [Neural ODEs, Chen et al., 2018] define an additional ODE which describes the trajectory of the flow’s gradient, avoiding the need to backpropagate through an ODE solver. A third ODE can be used to track the evolution of the log density, and the entire system can be solved with a suitable integrator. The resulting continuous-time flow is known as FFJORD [Grathwohl et al., 2018]. Like flows based on coupling layers, FFJORD is also invertible in ‘one pass’, but here this term refers to solving a system of ODEs, rather than performing a single neural-network pass.

5 Experiments

Table 1: Test log likelihood (in nats) for UCI datasets and BSDS300, with error bars corresponding to two standard deviations. NAF, Block-NAF, and SOS report error bars across repeated runs rather than across the test set. FFJORD do not report error bars. Superscript indicates results are taken from the existing literature.
Model POWER GAS HEPMASS MINIBOONE BSDS300
FFJORD 0.46 8.59 -14.92 -10.43 157.40
Glow 0.42±0.01 12.24±0.03 -16.99±0.02 -10.55±0.45 156.95±0.28
Q-NSF (C) 0.64±0.01 12.80±0.02 -15.35±0.02 -9.35±0.44 157.65±0.28
RQ-NSF (C) 0.64±0.01 13.09±0.02 -14.75±0.03 -9.67±0.47 157.54±0.28
MAF 0.45±0.01 12.35±0.02 -17.03±0.02 -10.92±0.46 156.95±0.28
Q-NSF (AR) 0.66±0.01 12.91±0.02 -14.67±0.03 -9.72±0.47 157.42±0.28
NAF 0.62±0.01 11.96±0.33 -15.09±0.40 -8.86±0.15 157.73±0.04
Block-NAF 0.61±0.01 12.06±0.09 -14.71±0.38 -8.95±0.07 157.36±0.03
SOS 0.60±0.01 11.99±0.41 -15.15±0.10 -8.90±0.11 157.48±0.41
RQ-NSF (AR) 0.66±0.01 13.09±0.02 -14.01±0.03 -9.22±0.48 157.31±0.28

In our experiments, the neural network NN which computes the parameters of the elementwise transformations is a residual network [He et al., 2016a] with pre-activation residual blocks [He et al., 2016b]. For autoregressive transformations, the layers must be masked so as to preserve autoregressive structure, and so we use the ResMADE architecture outlined by Nash and Durkan [2019]. Preliminary results indicated only minor differences in setting the tail bound B within the range [1,5], and so we fix a value B=3 across experiments, and find this to work robustly. We also fix the number of bins K=8 across our experiments, unless otherwise noted. We implement all invertible linear transformations using the LU-decomposition, where the permutation matrix 𝐏 is fixed at the beginning of training, and the product 𝐋𝐔 is initialized to the identity. For all non-image experiments, we define a flow ‘step’ as the composition of an invertible linear transformation with either a coupling or autoregressive transform, and we use 10 steps per flow in all our experiments, unless otherwise noted. All flows use a standard-normal noise distribution. We use the Adam optimizer [Kingma and Ba, 2014], and anneal the learning rate according to a cosine schedule [Loshchilov and Hutter, 2017]. In some cases, we find applying dropout [Srivastava et al., 2014] in the residual blocks beneficial for regularization. Full experimental details are provided in creftype B. Code is available online at https://github.com/bayesiains/nsf.

5.1 Density estimation of tabular data

We first evaluate our proposed flows using a selection of datasets from the UCI machine-learning repository Dheeru and Karra Taniskidou [2017] and BSDS300 collection of natural images [Martin et al., 2001]. We follow the experimental setup and pre-processing of Papamakarios et al. [2017], who make their data available online [Papamakarios, 2018]. We also update their MAF results using our codebase with ResMADE and invertible linear layers instead of permutations, providing a stronger baseline. For comparison, we modify the quadratic splines of Müller et al. [2018] to match the rational-quadratic transforms, by defining them on the range [-B,B] instead of [0,1], and adding linear tails, also matching the boundary derivatives as in the rational-quadratic case. We denote this model Q-NSF. Our results are shown in creftype 1, where the mid-rule separates flows with one-pass inverse from autoregressive flows.

Both RQ-NSF (C) and RQ-NSF (AR) achieve state-of-the-art results for a normalizing flow on the Power, Gas, and Hepmass datasets, tied with Q-NSF (C) and Q-NSF (AR) on the Power dataset. Moreover, RQ-NSF (C) significantly outperforms both Glow and FFJORD, achieving scores competitive with the best autoregressive models. These results close the gap between autoregressive flows and flows based on coupling layers, and demonstrate that, in some cases, it may not be necessary to sacrifice one-pass sampling for density-estimation performance.

5.2 Improving the variational autoencoder

Next, we examine our proposed flows in the context of the variational autoencoder [VAE, Kingma and Welling, 2014, Rezende et al., 2014], where they can act as both flexible prior and approximate posterior distributions. For our experiments, we use dynamically binarized versions of the MNIST dataset of handwritten digits [LeCun et al., 1998], and the EMNIST dataset variant featuring handwritten letters [Cohen et al., 2017]. We measure the capacity of our flows to improve over the commonly used baseline of a standard-normal prior and diagonal-normal approximate posterior, as well as over either coupling or autoregressive distributions with affine transformations. Quantitative results are shown in creftype 2, and image samples in creftype C.

All models improve significantly over the baseline, but perform very similarly otherwise, with most featuring overlapping error bars. Considering the disparity in density-estimation performance in the previous section, this is likely due to flows with affine transformations being sufficient to model the latent space for these datasets, with little scope for RQ-NSF flows to demonstrate their increased flexibility. Nevertheless, it is worthwhile to highlight that RQ-NSF (C) is the first class of model which can potentially match the flexibility of autoregressive models, and which requires no modification for use as either a prior or approximate posterior, due to its one-pass invertibility.

Table 2: Variational autoencoder test-set results (in nats) for the evidence lower bound (ELBO) and importance-weighted estimate of the log likelihood (computed as by Burda et al. [2016] using 1000 importance samples). Error bars correspond to two standard deviations.
MNIST EMNIST
Posterior/Prior ELBO logp(𝐱) ELBO logp(𝐱)
Baseline -85.61±0.51 -81.31±0.43 -125.89±0.41 -120.88±0.38
Glow -82.25±0.46 -79.72±0.42 -120.04±0.40 -117.54±0.38
RQ-NSF (C) -82.08±0.46 -79.63±0.42 -119.74±0.40 -117.35±0.38
IAF/MAF -82.56±0.48 -79.95±0.43 -119.85±0.40 -117.47±0.38
RQ-NSF (AR) -82.14±0.47 -79.71±0.43 -119.49±0.40 -117.28±0.38

5.3 Generative modeling of images

Finally, we evaluate neural spline flows as generative models of images, measuring their capacity to improve upon baseline models with affine transforms. In this section, we focus solely on flows with a one-pass inverse in the style of RealNVP [Dinh et al., 2017] and Glow [Kingma and Dhariwal, 2018]. We use the CIFAR-10 [Krizhevsky and Hinton, 2009] and downsampled 64×64 ImageNet [van den Oord et al., 2016c, Russakovsky et al., 2015] datasets, with original 8-bit colour depth and with reduced 5-bit colour depth. We use Glow-like architectures with either affine (in the baseline model) or rational-quadratic coupling transforms, and provide full experimental detail in creftype B. Quantitative results are shown in creftype 3, and samples are shown in creftype 3 and creftype C.

RQ-NSF (C) improves upon the affine baseline in three out of four tasks, and the improvement is most significant on the 8-bit version of ImageNet64. At the same time, RQ-NSF (C) achieves scores that are competitive with the original Glow model, while significantly reducing the number of parameters required, in some cases by almost an order of magnitude. creftypecap 3 demonstrates that the model is capable of producing diverse, globally coherent samples which closely resemble real data. There is potential to further improve our results by replacing the uniform dequantization used in Glow with variational dequantization, and using more powerful networks with gating and self-attention mechanisms to parameterize the coupling transforms, both of which are explored by Ho et al. [2019].

Table 3: Test-set bits per dimension (BPD, lower is better) and parameter count for CIFAR-10 and ImageNet64 models. Superscript indicates results are taken from the existing literature.
CIFAR-10 5-bit CIFAR-10 8-bit ImageNet64 5-bit ImageNet64 8-bit
Model BPD Params BPD Params BPD Params BPD Params
Baseline 1.70 05.2M 3.41 11.1M 1.81 014.3M 3.91 014.3M
RQ-NSF (C) 1.70 05.3M 3.38 11.8M 1.77 015.6M 3.82 015.6M
Glow 1.67 44.0M 3.35 44.0M 1.76 110.9M 3.81 110.9M
Figure 3: Samples from image models for 5-bit (top) and 8-bit (bottom) datasets. Left: CIFAR-10. Right: ImageNet64.

6 Discussion

Long-standing probabilistic models such as copulas [Elidan, 2013] and Gaussianization [Chen and Gopinath, 2001] can simply represent complex marginal distributions that would require many layers of transformations in flow-based models like RealNVP and Glow. Differentiable spline-based coupling layers allow these flows, which are powerful ways to represent high-dimensional dependencies, to model distributions with complex shapes more quickly. Our results show that when we have enough data, the extra flexibility of spline-based layers leads to better generalization.

For tabular density estimation, both RQ-NSF (C) and RQ-NSF (AR) excel on Power, Gas, and Hepmass, the datasets with the highest ratio of data points to dimensionality from the five considered. In image experiments, RQ-NSF (C) achieves the best results on the ImageNet dataset, which has over an order of magnitude more data points than CIFAR-10. When the dimension is increased without a corresponding increase in dataset size, RQ-NSF still performs competitively with other approaches, but does not outperform them.

Overall, neural spline flows demonstrate that there is significant performance to be gained by upgrading the commonly-used affine transformations in coupling and autoregressive layers, without the need to sacrifice analytic invertibility. Monotonic spline transforms enable models based on coupling layers to achieve density-estimation performance on par with the best autoregressive flows, while retaining exact one-pass sampling. These models strike a novel middle ground between flexibility and practicality, providing a useful off-the-shelf tool for the enhancement of architectures like the variational autoencoder, while also improving parameter efficiency in generative modeling.

Rational-quadratic transforms are also a useful differentiable and invertible module in their own right, which could be included in many models that can be trained end-to-end. For instance, monotonic warping functions with a tractable Jacobian determinant are useful for supervised learning [Snelson et al., 2004]. More generally, invertibility can be useful for training very large networks, since activations can be recomputed on-the-fly for backpropagation, meaning gradient computation requires memory which is constant instead of linear in the depth of the network [Gomez et al., 2017, MacKay et al., 2018]. Monotonic splines are one way of constructing invertible elementwise transformations, but there may be others. The benefits of research in this direction are clear, and so we look forward to future work in this area.

Acknowledgements

This work was supported in part by the EPSRC Centre for Doctoral Training in Data Science, funded by the UK Engineering and Physical Sciences Research Council (grant EP/L016427/1) and the University of Edinburgh. George Papamakarios was also supported by Microsoft Research through its PhD Scholarship Programme.

References

  • Blinn [2007] J. F. Blinn. How to solve a cubic equation, part 5: Back to numerics. IEEE Computer Graphics and Applications, 27(3):78–89, 2007.
  • Burda et al. [2016] Y. Burda, R. B. Grosse, and R. Salakhutdinov. Importance weighted autoencoders. International Conference on Learning Representations, 2016.
  • Chen et al. [2018] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 2018.
  • Chen and Gopinath [2001] S. S. Chen and R. A. Gopinath. Gaussianization. Advances in Neural Information Processing Systems, 2001.
  • Cohen et al. [2017] G. Cohen, S. Afshar, J. Tapson, and A. van Schaik. EMNIST: an extension of MNIST to handwritten letters. arXiv:1702.05373, 2017.
  • De Cao et al. [2019] N. De Cao, I. Titov, and W. Aziz. Block neural autoregressive flow. Conference on Uncertainty in Artificial Intelligence, 2019.
  • Dheeru and Karra Taniskidou [2017] D. Dheeru and E. Karra Taniskidou. UCI machine learning repository, 2017. URL http://archive.ics.uci.edu/ml.
  • Dinh et al. [2015] L. Dinh, D. Krueger, and Y. Bengio. NICE: Non-linear independent components estimation. International Conference on Learning Representations, Workshop track, 2015.
  • Dinh et al. [2017] L. Dinh, J. Sohl-Dickstein, and S. Bengio. Density estimation using Real NVP. International Conference on Learning Representations, 2017.
  • Durkan et al. [2019] C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios. Cubic-spline flows. Workshop on Invertible Neural Networks and Normalizing Flows, International Conference on Machine Learning, 2019.
  • Elidan [2013] G. Elidan. Copulas in machine learning. Copulae in Mathematical and Quantitative Finance, 2013.
  • Germain et al. [2015] M. Germain, K. Gregor, I. Murray, and H. Larochelle. MADE: Masked autoencoder for distribution estimation. International Conference on Machine Learning, 2015.
  • Gomez et al. [2017] A. N. Gomez, M. Ren, R. Urtasun, and R. B. Grosse. The reversible residual network: Backpropagation without storing activations. Advances in Neural Information Processing Systems, 2017.
  • Goodfellow et al. [2014] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. Advances in Neural Information Processing Systems, 2014.
  • Grathwohl et al. [2018] W. Grathwohl, R. T. Q. Chen, J. Bettencourt, I. Sutskever, and D. K. Duvenaud. FFJORD: Free-form continuous dynamics for scalable reversible generative models. International Conference on Learning Representations, 2018.
  • Gregory and Delbourgo [1982] J. Gregory and R. Delbourgo. Piecewise rational quadratic interpolation to monotonic data. IMA Journal of Numerical Analysis, 2(2):123–130, 1982.
  • He et al. [2016a] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. IEEE Conference on Computer Vision and Pattern Recognition, 2016a.
  • He et al. [2016b] K. He, X. Zhang, S. Ren, and J. Sun. Identity mappings in deep residual networks. European Conference on Computer Vision, 2016b.
  • Ho et al. [2019] J. Ho, X. Chen, A. Srinivas, Y. Duan, and P. Abbeel. Flow++: Improving flow-based generative models with variational dequantization and architecture design. International Conference on Machine Learning, 2019.
  • Hoogeboom et al. [2019] E. Hoogeboom, R. van den Berg, and M. Welling. Emerging convolutions for generative normalizing flows. International Conference on Machine Learning, 2019.
  • Huang et al. [2018] C.-W. Huang, D. Krueger, A. Lacoste, and A. Courville. Neural autoregressive flows. International Conference on Machine Learning, 2018.
  • Hutchinson [1990] M. F. Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statistics - Simulation and Computation, 19(2):433–450, 1990.
  • Ioffe and Szegedy [2015] S. Ioffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International Conference on Machine Learning, 2015.
  • Jaini et al. [2019] P. Jaini, K. A. Selby, and Y. Yu. Sum-of-squares polynomial flow. International Conference on Machine Learning, 2019.
  • Kim et al. [2018] S. Kim, S.-g. Lee, J. Song, and S. Yoon. FloWaveNet: A generative flow for raw audio. arXiv:1811.02155, 2018.
  • Kingma and Ba [2014] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2014.
  • Kingma and Dhariwal [2018] D. P. Kingma and P. Dhariwal. Glow: Generative flow with invertible 1×1 convolutions. Advances in Neural Information Processing Systems, 2018.
  • Kingma and Welling [2014] D. P. Kingma and M. Welling. Auto-encoding variational Bayes. International Conference on Learning Representations, 2014.
  • Kingma et al. [2016] D. P. Kingma, T. Salimans, R. Jozefowicz, X. Chen, I. Sutskever, and M. Welling. Improved variational inference with inverse autoregressive flow. Advances in Neural Information Processing Systems, 2016.
  • Krizhevsky and Hinton [2009] A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Master’s thesis, Department of Computer Science, University of Toronto, 2009.
  • Kumar et al. [2019] M. Kumar, M. Babaeizadeh, D. Erhan, C. Finn, S. Levine, L. Dinh, and D. Kingma. VideoFlow: A flow-based generative model for video. arXiv:1903.01434, 2019.
  • LeCun et al. [1998] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.
  • Loaiza-Ganem et al. [2017] G. Loaiza-Ganem, Y. Gao, and J. P. Cunningham. Maximum entropy flow networks. International Conference on Learning Representations, 2017.
  • Loshchilov and Hutter [2017] I. Loshchilov and F. Hutter. SGDR: Stochastic gradient descent with warm restarts. International Conference on Learning Representations, 2017.
  • Louizos and Welling [2017] C. Louizos and M. Welling. Multiplicative normalizing flows for variational Bayesian neural networks. International Conference on Machine Learning, 2017.
  • MacKay et al. [2018] M. MacKay, P. Vicol, J. Ba, and R. B. Grosse. Reversible recurrent neural networks. Advances in Neural Information Processing Systems, 2018.
  • Martin et al. [2001] D. Martin, C. Fowlkes, D. Tal, and J. Malik. A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics. International Conference on Computer Vision, 2001.
  • Müller et al. [2018] T. Müller, B. McWilliams, F. Rousselle, M. Gross, and J. Novák. Neural importance sampling. arXiv:1808.03856, 2018.
  • Nash and Durkan [2019] C. Nash and C. Durkan. Autoregressive energy machines. International Conference on Machine Learning, 2019.
  • Papamakarios [2018] G. Papamakarios. Preprocessed datasets for MAF experiments, 2018. URL https://doi.org/10.5281/zenodo.1161203.
  • Papamakarios et al. [2017] G. Papamakarios, T. Pavlakou, and I. Murray. Masked autoregressive flow for density estimation. Advances in Neural Information Processing Systems, 2017.
  • Papamakarios et al. [2019] G. Papamakarios, D. C. Sterratt, and I. Murray. Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows. International Conference on Artificial Intelligence and Statistics, 2019.
  • Prenger et al. [2018] R. Prenger, R. Valle, and B. Catanzaro. WaveGlow: A flow-based generative network for speech synthesis. arXiv:1811.00002, 2018.
  • Rezende and Mohamed [2015] D. J. Rezende and S. Mohamed. Variational inference with normalizing flows. International Conference on Machine Learning, 2015.
  • Rezende and Viola [2018] D. J. Rezende and F. Viola. Taming VAEs. arXiv:1810.00597, 2018.
  • Rezende et al. [2014] D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. International Conference on Machine Learning, 2014.
  • Russakovsky et al. [2015] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. C. Berg, and L. Fei-Fei. ImageNet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):211–252, 2015.
  • Salimans et al. [2017] T. Salimans, A. Karpathy, X. Chen, and D. P. Kingma. PixelCNN++: Improving the PixelCNN with discretized logistic mixture likelihood and other modifications. International Conference on Learning Representations, 2017.
  • Salman et al. [2018] H. Salman, P. Yadollahpour, T. Fletcher, and K. Batmanghelich. Deep diffeomorphic normalizing flows. arXiv:1810.03256, 2018.
  • Snelson et al. [2004] E. Snelson, C. E. Rasmussen, and Z. Ghahramani. Warped Gaussian processes. Advances in Neural Information Processing Systems, 2004.
  • Srivastava et al. [2014] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1):1929–1958, 2014.
  • Steffen [1990] M. Steffen. A simple method for monotonic interpolation in one dimension. Astronomy and Astrophysics, 239:443, 1990.
  • Tomczak and Welling [2016] J. M. Tomczak and M. Welling. Improving variational auto-encoders using Householder flow. arXiv:1611.09630, 2016.
  • Uria et al. [2013] B. Uria, I. Murray, and H. Larochelle. RNADE: The real-valued neural autoregressive density-estimator. Advances in Neural Information Processing Systems, 2013.
  • van den Berg et al. [2018] R. van den Berg, L. Hasenclever, J. M. Tomczak, and M. Welling. Sylvester normalizing flows for variational inference. Conference on Uncertainty in Artificial Intelligence, 2018.
  • van den Oord et al. [2016a] A. van den Oord, S. Dieleman, H. Zen, K. Simonyan, O. Vinyals, A. Graves, N. Kalchbrenner, A. Senior, and K. Kavukcuoglu. WaveNet: A generative model for raw audio. ISCA Speech Synthesis Workshop, 2016a.
  • van den Oord et al. [2016b] A. van den Oord, N. Kalchbrenner, L. Espeholt, K. Kavukcuoglu, O. Vinyals, and A. Graves. Conditional image generation with PixelCNN decoders. Advances in Neural Information Processing Systems, 2016b.
  • van den Oord et al. [2016c] A. van den Oord, N. Kalchbrenner, and K. Kavukcuoglu. Pixel recurrent neural networks. International Conference on Machine Learning, 2016c.
  • Ziegler and Rush [2019] Z. M. Ziegler and A. M. Rush. Latent normalizing flows for discrete sequences. arXiv:1901.10548, 2019.

Appendix A Monotonic rational-quadratic transforms

A.1 Parameterization of the spline

We here include an outline of the method of Gregory and Delbourgo [1982] which closely matches the original paper. Let {(x(k),y(k))}k=0K be a given set of knot points in the plane which satisfy

(x(0),y(0))=(-B,-B), (7)
(x(K),y(K))=(B,B), (8)
x(k)<x(k+1) and y(k)<y(k+1)for all k=0,,K-1. (9)

Let {d(k)}k=0K be the non-negative derivative values at these knot points (we take d(0)=d(K)=1 so that the spline matches the derivative of the linear tails). Given these quantities, the algorithm of Gregory and Delbourgo [1982] defines a monotonic rational-quadratic spline which passes through each knot and has the given derivative value at each knot as follows:

  1. 1.

    Let w(k)=x(k+1)-x(k) be the bin widths, and s(k)=(y(k+1)-y(k))/w(k) be the slopes of the lines joining the co-ordinates.

  2. 2.

    For x[x(k),x(k+1)], let ξ=(x-x(k))/w(k), so that ξ[0,1].

  3. 3.

    Then, for x[x(k),x(k+1)], k=0,,K-1, define

    g(x)=α(k)(ξ)β(k)(ξ), (10)

    where

    α(k)(ξ) =s(k)y(k+1)ξ2+[y(k)d(k+1)+y(k+1)d(k)]ξ(1-ξ)+s(k)y(k)(1-ξ)2, (11)
    β(k)(ξ) =s(k)ξ2+[d(k+1)+d(k)]ξ(1-ξ)+s(k)(1-ξ)2. (12)

Gregory and Delbourgo [1982] note that β(k)(ξ) can be rewritten as

β(k)(ξ)=s(k)+[d(k+1)+d(k)-2s(k)]ξ(1-ξ), (13)

so that the quotient can be written as

α(k)(ξ)β(k)(ξ)=y(k)+(y(k+1)-y(k))[s(k)ξ2+d(k)ξ(1-ξ)]s(k)+[d(k+1)+d(k)-2s(k)]ξ(1-ξ), (14)

which is less prone to numerical issues, especially for small values of s(k). Gregory and Delbourgo [1982] show that the spline defined by creftype 10 interpolates between the given knots, satisfies the derivative constraints at the knot points, and is monotonic on each bin. Rational-quadratic functions also provide flexibility over previous approaches: it is not possible to match arbitrary values and derivatives of a function at two boundary knots with a quadratic polynomial, or a monotonic segment of a cubic polynomial.

A.2 Computing the derivative

The derivative of creftype 10 is given by the quotient rule:

ddx[α(k)(ξ)β(k)(ξ)]=ddξ[α(k)(ξ)β(k)(ξ)]dξdx=1w(k)β(k)(ξ)ddξ[α(k)(ξ)]-α(k)(ξ)ddξ[β(k)(ξ)][β(k)(ξ)]2. (15)

It can be shown that

β(k)(ξ)dα(k)(ξ)dξ-α(k)(ξ)dβ(k)(ξ)dξ=w(k)(s(k))2[d(k+1)ξ2+2s(k)ξ(1-ξ)+d(k)(1-ξ)2], (16)

so that

ddx[α(k)(ξ)β(k)(ξ)]=(s(k))2[d(k+1)ξ2+2s(k)ξ(1-ξ)+d(k)(1-ξ)2][s(k)+[d(k+1)+d(k)-2s(k)]ξ(1-ξ)]2. (17)

Since the rational-quadratic transform is monotonic and acts elementwise, the logarithm of the absolute value of the determinant of its Jacobian is given by a sum of the logarithm of creftype 17 for each transformed x.

A.3 Computing the inverse

Computing the inverse of a monotonic rational-quadratic transformation when the value to invert lies in the tails is trivial. The problem of inversion is thus reduced to computing the inverse of the monotonic rational-quadratic spline. Consider a rational-quadratic function

y=α(ξ(x))β(ξ(x))=α0+α1x+α2x2β0+β1x+β2x2, (18)

which arises as the result of the algorithm outlined in creftype A.1. The coefficients are such that the function is monotonically-increasing in its associated bin. Inverting the function involves solving a quadratic equation:

q(x) =α(ξ(x))-yβ(ξ(x)) (19)
=ax2+bx+c= 0, (20)

where the coefficients depend on the target output y:

a=α2-β2y,b=α1-β1y,c=α0-β0y. (21)

Only one of the two solutions lies in the function’s associated bin. To identify the solution in general, we identify that along the line of corresponding (x,y) values, q(x)=0 and so dqdx=0, where

dqdx =qx+qyyx (22)
=qx-β(ξ(x))+veyx+ve=0. (23)

We substituted a partial derivative of creftype 19, noted creftype 12 is positive, and noted that the spline y(x) is monotonic and increasing. To satisfy creftype 23, qx>0, which corresponds to this solution to creftype 20:

x=-b+b2-4ac2a=2c-b-b2-4ac, (24)

where the first form is more commonly quoted, but the second form is numerically more precise when 4ac is small.

We can rearrange creftype 14 to show

a =(y(k+1)-y(k))[s(k)-d(k)]+(y-y(k))[d(k+1)+d(k)-2s(k)], (25)
b =(y(k+1)-y(k))d(k)-(y-y(k))[d(k+1)+d(k)-2s(k)], (26)
c =-s(k)(y-y(k)). (27)

Appendix B Experimental details

B.1 Tabular density estimation

Model selection is performed using the standard validation splits for these datasets. We clip the norm of gradients to the range [-5,5], and find this helps stabilize training. We modify MAF by replacing permutations with invertible linear layers. Hyperparameter settings are shown for coupling flows in creftype 4 and autoregressive flows in creftype 5. We include the dimensionality and number of training data points in each table for reference. For higher dimensional datasets such as Hepmass and BSDS300, we found increasing the number of coupling layers beneficial. This was not necessary for Miniboone, where overfitting was an issue due to the low number of data points.

Table 4: Hyperparameters for density-estimation results using coupling layers in creftype 5.1.
Power Gas Hepmass Miniboone BSDS300
Dimension 6 8 21 43 63
Train data points 1,615,917 852,174 315,123 29,556 1,000,000
Batch size 512 512 256 128 512
Training steps 400,000 400,000 400,000 200,000 400,000
Learning rate 0.0005 0.0005 0.0005 0.0003 0.0005
Flow steps 10 10 20 10 20
Residual blocks 2 2 1 1 1
Hidden features 256 256 128 32 128
Bins 8 8 8 4 8
Dropout 0.0 0.1 0.2 0.2 0.2
Table 5: Hyperparameters for density-estimation results using autoregressive layers in creftype 5.1.
Power Gas Hepmass Miniboone BSDS300
Dimension 6 8 21 43 63
Train data points 1,615,917 852,174 315,123 29,556 1,000,000
Batch size 512 512 512 64 512
Training steps 400,000 400,000 400,000 250,000 400,000
Learning rate 0.0005 0.0005 0.0005 0.0003 0.0005
Flow steps 10 10 10 10 10
Residual blocks 2 2 2 1 2
Hidden features 256 256 256 64 512
Bins 8 8 8 4 8
Dropout 0.0 0.1 0.2 0.2 0.2

B.2 Improving the variational autoencoder

We use the Adam optimizer [Kingma and Ba, 2014] with default hyperparameters, annealing an initial learning rate of 0.0005 to 0 using a cosine schedule [Loshchilov and Hutter, 2017] over 150,000 training steps with batch size 256. We use a ‘warm-up’ phase for the KL divergence term of the loss, where the multiplier for this term is initialized to 0.5 and linearly increased to 1 over the first 10% of training. This modification initially reduces the penalty incurred by the approximate posterior in deviating from the prior, and similar schemes have been shown to improve VAE training dynamics [Rezende and Viola, 2018]. Model selection is performed using a held-out validation set of 10,000 samples for MNIST, and 20,000 samples for EMNIST.

We use 32 latent features, and residual nets use 2 blocks, with 64 latent features for coupling layers, and 128 latent features for autoregressive layers. Both coupling and autoregressive flows use 10 steps. As with the tabular density-estimation experiments, we modify IAF [Kingma et al., 2016] and MAF [Papamakarios et al., 2017] by replacing permutations with invertible linear layers using an LU-decomposition. All NSF models use 8 bins. The encoder and decoder architectures are set up exactly as described by Nash and Durkan [2019], and are similar to those used in IAF [Kingma et al., 2016] and NAF [Huang et al., 2018].

Conditioning the approximate posterior distribution q(𝐳|𝐱) follows a multi-stage procedure. First, the encoder computes a context vector 𝐡 of dimension 64 as a function of the input 𝐱. This vector is then mapped to the mean and diagonal covariance of a Gaussian distribution in the latent space. Then, 𝐡 is also given as input to the residual nets in each of the flow’s coupling or autoregressive layers, where it is concatenated with the input 𝐳, mapped to the required number of hidden features, and also used to modulate the additive update of each residual block with a sigmoid gate. We found this scheme to work well across experiments.

B.3 Generative modeling of images

For image-modeling experiments we use a Glow-like model architecture introduced by Kingma and Dhariwal [2018, Section 3]. This involves stacking multiple steps for each level in the multi-scale architecture of Dinh et al. [2017], where each step consists of an actnorm layer, an invertible 1×1 convolution and a coupling transform. For our RQ-NSF (C) model, we make the following modifications to the original Glow model: we replace affine coupling transforms with rational-quadratic coupling transforms, we go back to residual convolutional networks as used in RealNVP Dinh et al. [2017], and we use an additional 1×1 convolution at the end of each level of transforms. The basline model is the same as RQ-NSF (C), except that it uses affine coupling transforms instead of rational-quadratic ones. For CIFAR-10 experiments we do not factor out dimensions at the end of each level, but still use the squeezing operation to trade spatial resolution for depth.

For all experiments we use 3 residual blocks and batch normalization [Ioffe and Szegedy, 2015] in the residual networks which parameterize the coupling transforms. We use 7 steps per level for all experiments, resulting in a total of 21 coupling transforms for CIFAR-10, and 28 coupling transform for ImageNet64 (Glow models used by Kingma and Dhariwal [2018] use 96 and 192 affine coupling transforms for CIFAR-10 and ImageNet64 respectively).

We use the Adam [Kingma and Ba, 2014] optimizer with default β1 and β2 values. An initial learning rate of 0.0005 is annealed to 0 following a cosine schedule [Loshchilov and Hutter, 2017]. We train for 100,000 steps for 5-bit experiments, and for 200,000 steps for 8-bit experiments. To track the performance of our models, we split off 1% of the training data to use as a development set. Due to the resource requirements of the experiments, we perform a limited manual hyper-parameter exploration. Final values are reported in creftype 6.

We use a single NVIDIA Tesla P100 GPU card per CIFAR-10 experiment, and two such cards per ImageNet64 experiment. Training for 200,000 steps takes about 5 days with this setup.

Table 6: Hyperparameters for generative image-modeling results in creftype 5.3.
Dataset Batch Size Levels Hidden Channels Bins Dropout
CIFAR-10 5-bit 512 3 64 2 0.2
8-bit 512 3 96 4 0.2
ImageNet64 5-bit 256 4 96 8 0.1
8-bit 256 4 96 8 0.0

Appendix C Additional experimental results

C.1 Affine coupling transforms for 2D datasets

Training data Flow density Flow samples
Figure 4: Qualitative results for two-dimensional synthetic datasets using two affine coupling layers.

C.2 Samples

Image samples for VAE experiments are shown in creftype 5. Additional samples for generative image-modeling experiments are shown in creftype 6.

(a) MNIST (b) EMNIST-letters
Figure 5: VAE samples. Top to bottom: training data, RQ-NSF (C), RQ-NSF (AR).
(a) CIFAR-10 5-bit
(b) ImageNet64 5-bit
(c) CIFAR-10 8-bit
(d) ImageNet64 8-bit
Figure 6: Additional image samples for generative image-modeling experiments.