Code for Neural Spline Flows paper
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.READ FULL TEXT VIEW PDF
Code for Neural Spline Flows paper
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 :
The probability density of under the flow is obtained by a change of 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 are often implemented by composing a series of invertible neural-network modules. Given a dataset , the flow is trained by maximizing the total log likelihood 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 , usually with a neural network. Depending on the flow’s intended use cases, there are practical constraints in addition to formal invertibility:
[leftmargin=*, topsep=0pt, itemsep=0pt]
To train a density estimator, we need to be able to evaluate the Jacobian determinant and the inverse function quickly. We don’t evaluate , so the flow is usually defined by specifying .
If we wish to draw samples using eq. 1, we would like to be available analytically, rather than having to invert with iterative or approximate methods.
Ideally, we would like both and 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 times slower to invert than to evaluate, where 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 fig. 1.
A coupling transform [Dinh et al., 2015] maps an input to an output in the following way:
Split the input into two parts, .
Compute parameters , where NN is an arbitrary neural network.
Compute for in parallel, where is an invertible function parameterized by .
Set , and return .
The Jacobian matrix of a coupling transform is lower triangular, since is given by transforming elementwise as a function of , and is equal to . Thus, the Jacobian determinant of the coupling transform is given by , 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 to compute 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 , as noted by Papamakarios et al. . In this way, advances in flows based on coupling transforms can be applied to autoregressive flows, and vice versa.
and is usually constrained to be positive. The additive transformation corresponds to the special case . 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 section C.1). We aim to choose a more flexible , that is still differentiable and easy to invert.
Recently, Müller et al.  proposed a powerful generalization of the above affine transformations, based on monotonic piecewise polynomials. The idea is to restrict the input domain of to the interval , partition the input domain into bins, and define to be a simple polynomial segment within each bin. Müller et al.  restrict themselves to monotonically-increasing linear and quadratic polynomial segments, whose coefficients are parameterized by . Moreover, the polynomial segments are restricted to match at the bin boundaries so that
is continuous. Functions of this form, which interpolate between data using piecewise polynomials, are known aspolynomial 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. . We proposed to implement as a monotonic cubic spline Steffen , where
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. , our spline transform and its inverse were defined only on the interval . To ensure that the input is always between and
, 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, 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 section 3.1 we propose a modified method based on rational-quadratic splines which overcomes these difficulties.
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 . Kingma and Dhariwal  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 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 time where is the batch size. Alternatively, we can pay a one-time cost of to explicitly compute , which can then be cached for re-use.
We propose to implement the function 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  to parameterize a monotonic rational-quadratic spline. The spline itself maps an interval to . 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 different rational-quadratic functions, with boundaries set by coordinates known as knots. The knots monotonically increase between and . We give the spline arbitrary positive values for the derivatives at the internal points, and set the boundary derivatives to 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 fig. 1, and section A.1 gives full details.
The practical implementation of the monotonic rational-quadratic coupling transform is as follows:
A neural network NN takes
as input and outputs an unconstrained parameter vectorof length for each .
Vector is partitioned as , where and have length , and has length .
Vectors and are each passed through a softmax and multiplied by ; the outputs are interpreted as the widths and heights of the bins, which must be positive and span the interval. Cumulative sums of the bin widths and heights, each starting at , yield the knots .
The vector 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 requires finding the bin in which 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 section A.2 and section 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 , 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 can be computed as a function of 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
(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|
all elements of at once as follows:
Figure 2 demonstrates the flexibility of our rational-quadratic coupling transform on synthetic two-dimensional datasets. Using just two coupling layers, each with 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 section C.1).
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
, 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.
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 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 convolution presented in Glow, Hoogeboom et al.  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. Tomczak and Welling, 2016].
Outside of those discussed in section 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.
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.
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.
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 . Preliminary results indicated only minor differences in setting the tail bound within the range , and so we fix a value across experiments, and find this to work robustly. We also fix the number of bins 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 appendix B. Code is available online at https://github.com/bayesiains/nsf.
We first evaluate our proposed flows using a selection of datasets from the UCI machine-learning repository Dheeru and Karra Taniskidou  and BSDS300 collection of natural images [Martin et al., 2001]. We follow the experimental setup and pre-processing of Papamakarios et al. , 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.  to match the rational-quadratic transforms, by defining them on the range instead of , 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 table 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.
, 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 table 2, and image samples in appendix 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.
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 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 appendix B. Quantitative results are shown in table 3, and samples are shown in fig. 3 and appendix 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. Figure 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. .
|CIFAR-10 5-bit||CIFAR-10 8-bit||ImageNet64 5-bit||ImageNet64 8-bit|
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.
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.
Conference on Uncertainty in Artificial Intelligence, 2019.
Reversible recurrent neural networks.Advances in Neural Information Processing Systems, 2018.
We here include an outline of the method of Gregory and Delbourgo  which closely matches the original paper. Let be a given set of knot points in the plane which satisfy
Let be the non-negative derivative values at these knot points (we take so that the spline matches the derivative of the linear tails). Given these quantities, the algorithm of Gregory and Delbourgo  defines a monotonic rational-quadratic spline which passes through each knot and has the given derivative value at each knot as follows:
Let be the bin widths, and be the slopes of the lines joining the co-ordinates.
For , let , so that .
Then, for , , define
Gregory and Delbourgo  note that can be rewritten as
so that the quotient can be written as
which is less prone to numerical issues, especially for small values of . Gregory and Delbourgo  show that the spline defined by eq. 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.
The derivative of eq. 10 is given by the quotient rule:
It can be shown that
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 eq. 17 for each transformed .
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
which arises as the result of the algorithm outlined in section A.1. The coefficients are such that the function is monotonically-increasing in its associated bin. Inverting the function involves solving a quadratic equation:
where the coefficients depend on the target output:
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 values, and so , where
where the first form is more commonly quoted, but the second form is numerically more precise when is small.
We can rearrange eq. 14 to show
Model selection is performed using the standard validation splits for these datasets. We clip the norm of gradients to the range
, and find this helps stabilize training. We modify MAF by replacing permutations with invertible linear layers. Hyperparameter settings are shown for coupling flows intable 4 and autoregressive flows in table 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.
|Train data points||1,615,917||852,174||315,123||29,556||1,000,000|
|Train data points||1,615,917||852,174||315,123||29,556||1,000,000|
We use the Adam optimizer [Kingma and Ba, 2014] with default hyperparameters, annealing an initial learning rate of to using a cosine schedule [Loshchilov and Hutter, 2017] over training steps with batch size . We use a ‘warm-up’ phase for the KL divergence term of the loss, where the multiplier for this term is initialized to and linearly increased to over the first 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 samples for MNIST, and 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 , and are similar to those used in IAF [Kingma et al., 2016] and NAF [Huang et al., 2018].
Conditioning the approximate posterior distribution follows a multi-stage procedure. First, the encoder computes a context vector of dimension 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.
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. , where each step consists of an actnorm layer, an invertible 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. , and we use an additional 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  use 96 and 192 affine coupling transforms for CIFAR-10 and ImageNet64 respectively).
We use the Adam [Kingma and Ba, 2014] optimizer with default and values. An initial learning rate of 0.0005 is annealed to 0 following a cosine schedule [Loshchilov and Hutter, 2017]. We train for steps for 5-bit experiments, and for 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 table 6.
We use a single NVIDIA Tesla P100 GPU card per CIFAR-10 experiment, and two such cards per ImageNet64 experiment. Training for steps takes about days with this setup.
|Dataset||Batch Size||Levels||Hidden Channels||Bins||Dropout|
|Training data||Flow density||Flow samples|
|(a) MNIST||(b) EMNIST-letters|