Masked Autoregressive Flow for Density Estimation

Introduction

The joint density $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ of a set of variables $\mathbf{x}$ is a central object of interest in machine learning. Being able to access and manipulate $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ enables a wide range of tasks to be performed, such as inference, prediction, data completion and data generation. As such, the problem of estimating $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ from a set of examples $\left\{\mathbf{x}_{n}\right\}$ is at the core of probabilistic unsupervised learning and generative modelling.

In recent years, using neural networks for density estimation has been particularly successful. Combining the flexibility and learning capacity of neural networks with prior knowledge about the structure of data to be modelled has led to impressive results in modelling natural images and audio data . State-of-the-art neural density estimators have also been used for likelihood-free inference from simulated data , variational inference , and as surrogates for maximum entropy models .

Neural density estimators differ from other approaches to generative modelling—such as variational autoencoders and generative adversarial networks —in that they readily provide exact density evaluations. As such, they are more suitable in applications where the focus is on explicitly evaluating densities, rather than generating synthetic data. For instance, density estimators can learn suitable priors for data from large unlabelled datasets, for use in standard Bayesian inference . In simulation-based likelihood-free inference, conditional density estimators can learn models for the likelihood or the posterior from simulated data. Density estimators can learn effective proposals for importance sampling or sequential Monte Carlo ; such proposals can be used in probabilistic programming environments to speed up inference . Finally, conditional density estimators can be used as flexible inference networks for amortized variational inference and as part of variational autoencoders .

A challenge in neural density estimation is to construct models that are flexible enough to represent complex densities, but have tractable density functions and learning algorithms. There are mainly two families of neural density estimators that are both flexible and tractable: autoregressive models and normalizing flows . Autoregressive models decompose the joint density as a product of conditionals, and model each conditional in turn. Normalizing flows transform a base density (e.g. a standard Gaussian) into the target density by an invertible transformation with tractable Jacobian.

Our starting point is the realization (as pointed out by Kingma et al. ) that autoregressive models, when used to generate data, correspond to a differentiable transformation of an external source of randomness (typically obtained by random number generators). This transformation has a tractable Jacobian by design, and for certain autoregressive models it is also invertible, hence it precisely corresponds to a normalizing flow. Viewing an autoregressive model as a normalizing flow opens the possibility of increasing its flexibility by stacking multiple models of the same type, by having each model provide the source of randomness for the next model in the stack. The resulting stack of models is a normalizing flow that is more flexible than the original model, and that remains tractable.

In this paper we present Masked Autoregressive Flow (MAF), which is a particular implementation of the above normalizing flow that uses the Masked Autoencoder for Distribution Estimation (MADE) as a building block. The use of MADE enables density evaluations without the sequential loop that is typical of autoregressive models, and thus makes MAF fast to evaluate and train on parallel computing architectures such as Graphics Processing Units (GPUs). We show a close theoretical connection between MAF and Inverse Autoregressive Flow (IAF) , which has been designed for variational inference instead of density estimation, and show that both correspond to generalizations of the successful Real NVP . We experimentally evaluate MAF on a wide range of datasets, and we demonstrate that MAF outperforms RealNVP and achieves state-of-the-art performance on a variety of general-purpose density estimation tasks.

Background

Using the chain rule of probability, any joint density $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ can be decomposed into a product of one-dimensional conditionals as $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}=\prod_{i}{p\mathopen{}\left(x_{i}\,|\,\mathbf{x}_{1:i-1}\right)\mathclose{}}$. Autoregressive density estimators model each conditional $p\mathopen{}\left(x_{i}\,|\,\mathbf{x}_{1:i-1}\right)\mathclose{}$ as a parametric density, whose parameters are a function of a hidden state $\mathbf{h}_{i}$. In recurrent architectures, $\mathbf{h}_{i}$ is a function of the previous hidden state $\mathbf{h}_{i-1}$ and the $i^{\text{th}}$ input variable $x_{i}$. The Real-valued Neural Autoregressive Density Estimator (RNADE) uses mixtures of Gaussian or Laplace densities for modelling the conditionals, and a simple linear rule for updating the hidden state. More flexible approaches for updating the hidden state are based on Long Short-Term Memory recurrent neural networks .

A drawback of autoregressive models is that they are sensitive to the order of the variables. For example, the order of the variables matters when learning the density of Figure 1a if we assume a model with Gaussian conditionals. As Figure 1b shows, a model with order $\mathopen{}\left(x_{1},x_{2}\right)\mathclose{}$ cannot learn this density, even though the same model with order $\mathopen{}\left(x_{2},x_{1}\right)\mathclose{}$ can represent it perfectly. In practice is it hard to know which of the factorially many orders is the most suitable for the task at hand. Autoregressive models that are trained to work with an order chosen at random have been developed, and the predictions from different orders can then be combined in an ensemble . Our approach (Section 3) can use a different order in each layer, and using random orders would also be possible.

Straightforward recurrent autoregressive models would update a hidden state sequentially for every variable, requiring $D$ sequential computations to compute the probability $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ of a $D$-dimensional vector, which is not well-suited for computation on parallel architectures such as GPUs. One way to enable parallel computation is to start with a fully-connected model with $D$ inputs and $D$ outputs, and drop out connections in order to ensure that output $i$ will only be connected to inputs $1,2,\ldots,i\!-\!1$. Output $i$ can then be interpreted as computing the parameters of the $i^{\text{th}}$ conditional $p\mathopen{}\left(x_{i}\,|\,\mathbf{x}_{1:i-1}\right)\mathclose{}$. By construction, the resulting model will satisfy the autoregressive property, and at the same time it will be able to calculate $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ efficiently on a GPU. An example of this approach is the Masked Autoencoder for Distribution Estimation (MADE) , which drops out connections by multiplying the weight matrices of a fully-connected autoencoder with binary masks. Other mechanisms for dropping out connections include masked convolutions and causal convolutions .

2 Normalizing flows

A normalizing flow represents $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ as an invertible differentiable transformation $f$ of a base density $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$. That is, $\mathbf{x}=f\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ where $\mathbf{u}\sim\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$. The base density $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ is chosen such that it can be easily evaluated for any input $\mathbf{u}$ (a common choice for $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ is a standard Gaussian). Under the invertibility assumption for $f$, the density $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ can be calculated as

In order for Equation (1) to be tractable, the transformation $f$ must be constructed such that (a) it is easy to invert, and (b) the determinant of its Jacobian is easy to compute. An important point is that if transformations $f_{1}$ and $f_{2}$ have the above properties, then their composition $f_{1}\circ f_{2}$ also has these properties. In other words, the transformation $f$ can be made deeper by composing multiple instances of it, and the result will still be a valid normalizing flow.

There have been various approaches in developing normalizing flows. An early example is Gaussianization , which is based on successive application of independent component analysis. Enforcing invertibility with nonsingular weight matrices has been proposed , however in such approaches calculating the determinant of the Jacobian scales cubicly with data dimensionality in general. Planar/radial flows and Inverse Autoregressive Flow (IAF) are models whose Jacobian is tractable by design. However, they were developed primarily for variational inference and are not well-suited for density estimation, as they can only efficiently calculate the density of their own samples and not of externally provided datapoints. The Non-linear Independent Components Estimator (NICE) and its successor Real NVP have a tractable Jacobian and are also suitable for density estimation. IAF, NICE and Real NVP are discussed in more detail in Section 3.

Masked Autoregressive Flow

Consider an autoregressive model whose conditionals are parameterized as single Gaussians. That is, the $i^{\text{th}}$ conditional is given by

In the above, $f_{\mu_{i}}$ and $f_{\alpha_{i}}$ are unconstrained scalar functions that compute the mean and log standard deviation of the $i^{\text{th}}$ conditional given all previous variables. We can generate data from the above model using the following recursion:

In the above, $\mathbf{u}=\mathopen{}\left(u_{1},u_{2},\ldots,u_{I}\right)\mathclose{}$ is the vector of random numbers the model uses internally to generate data, typically by making calls to a random number generator often called randn().

Equation (3) provides an alternative characterization of the autoregressive model as a transformation $f$ from the space of random numbers $\mathbf{u}$ to the space of data $\mathbf{x}$. That is, we can express the model as $\mathbf{x}=f\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ where $\mathbf{u}\sim\mathcal{N}\mathopen{}\left(\mathbf{0},\mathbf{I}\right)\mathclose{}$. By construction, $f$ is easily invertible. Given a datapoint $\mathbf{x}$, the random numbers $\mathbf{u}$ that were used to generate it are obtained by the following recursion:

Due to the autoregressive structure, the Jacobian of $f^{-1}$ is triangular by design, hence its absolute determinant can be easily obtained as follows:

It follows that the autoregressive model can be equivalently interpreted as a normalizing flow, whose density $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ can be obtained by substituting Equations (4) and (5) into Equation (1). This observation was first pointed out by Kingma et al. .

A useful diagnostic for assessing whether an autoregressive model of the above type fits the target density well is to transform the train data $\left\{\mathbf{x}_{n}\right\}$ into corresponding random numbers $\left\{\mathbf{u}_{n}\right\}$ using Equation (4), and assess whether the $u_{i}$’s come from independent standard normals. If the $u_{i}$’s do not seem to come from independent standard normals, this is evidence that the model is a bad fit. For instance, Figure 1b shows that the scatter plot of the random numbers associated with the train data can look significantly non-Gaussian if the model fits the target density poorly.

Here we interpret autoregressive models as a flow, and improve the model fit by stacking multiple instances of the model into a deeper flow. Given autoregressive models $M_{1},M_{2},\ldots,M_{K}$, we model the density of the random numbers $\mathbf{u}_{1}$ of $M_{1}$ with $M_{2}$, model the random numbers $\mathbf{u}_{2}$ of $M_{2}$ with $M_{3}$ and so on, finally modelling the random numbers $\mathbf{u}_{K}$ of $M_{K}$ with a standard Gaussian. This stacking adds flexibility: for example, Figure 1c demonstrates that a flow of $5$ autoregressive models is able to learn multimodal conditionals, even though each model has unimodal conditionals. Stacking has previously been used in a similar way to improve model fit of deep belief nets and deep mixtures of factor analyzers .

We choose to implement the set of functions $\left\{f_{\mu_{i}},f_{\alpha_{i}}\right\}$ with masking, following the approach used by MADE . MADE is a feedforward network that takes $\mathbf{x}$ as input and outputs $\mu_{i}$ and $\alpha_{i}$ for all $i$ with a single forward pass. The autoregressive property is enforced by multiplying the weight matrices of MADE with suitably constructed binary masks. In other words, we use MADE with Gaussian conditionals as the building layer of our flow. The benefit of using masking is that it enables transforming from data $\mathbf{x}$ to random numbers $\mathbf{u}$ and thus calculating $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ in one forward pass through the flow, thus eliminating the need for sequential recursion as in Equation (4). We call this implementation of stacking MADEs into a flow Masked Autoregressive Flow (MAF).

2 Relationship with Inverse Autoregressive Flow

Like MAF, Inverse Autoregressive Flow (IAF) is a normalizing flow which uses MADE as its component layer. Each layer of IAF is defined by the following recursion:

Similarly to MAF, functions $\left\{f_{\mu_{i}},f_{\alpha_{i}}\right\}$ are computed using a MADE with Gaussian conditionals. The difference is architectural: in MAF $\mu_{i}$ and $\alpha_{i}$ are directly computed from previous data variables $\mathbf{x}_{1:i-1}$, whereas in IAF $\mu_{i}$ and $\alpha_{i}$ are directly computed from previous random numbers $\mathbf{u}_{1:i-1}$.

The consequence of the above is that MAF and IAF are different models with different computational trade-offs. MAF is capable of calculating the density $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ of any datapoint $\mathbf{x}$ in one pass through the model, however sampling from it requires performing $D$ sequential passes (where $D$ is the dimensionality of $\mathbf{x}$). In contrast, IAF can generate samples and calculate their density with one pass, however calculating the density $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ of an externally provided datapoint $\mathbf{x}$ requires $D$ passes to find the random numbers $\mathbf{u}$ associated with $\mathbf{x}$. Hence, the design choice of whether to connect $\mu_{i}$ and $\alpha_{i}$ directly to $\mathbf{x}_{1:i-1}$ (obtaining MAF) or to $\mathbf{u}_{1:i-1}$ (obtaining IAF) depends on the intended usage. IAF is suitable as a recognition model for stochastic variational inference , where it only ever needs to calculate the density of its own samples. In contrast, MAF is more suitable for density estimation, because each example requires only one pass through the model whereas IAF requires $D$.

The inverse transformation $f^{-1}$ from $\mathbf{x}$ to $\mathbf{u}$ can be seen as describing an implicit IAF with base density $\pi_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$, which defines the following implicit density over the $\mathbf{u}$ space:

3 Relationship with Real NVP

Real NVP (NVP stands for Non Volume Preserving) is a normalizing flow obtained by stacking coupling layers. A coupling layer is an invertible transformation $f$ from random numbers $\mathbf{u}$ to data $\mathbf{x}$ with a tractable Jacobian, defined by

In the above, $\odot$ denotes elementwise multiplication, and the $\exp$ is applied to each element of $\bm{\alpha}$. The transformation copies the first $d$ elements, and scales and shifts the remaining $D\!-\!d$ elements, with the amount of scaling and shifting being a function of the first $d$ elements. When stacking coupling layers into a flow, the elements are permuted across layers so that a different set of elements is copied each time. A special case of the coupling layer where $\bm{\alpha}\!=\!\mathbf{0}$ is used by NICE .

We can see that the coupling layer is a special case of both the autoregressive transformation used by MAF in Equation (3), and the autoregressive transformation used by IAF in Equation (6). Indeed, we can recover the coupling layer from the autoregressive transformation of MAF by setting $\mu_{i}=\alpha_{i}=0$ for $i\leq d$ and making $\mu_{i}$ and $\alpha_{i}$ functions of only $\mathbf{x}_{1:d}$ for $i>d$ (for IAF we need to make $\mu_{i}$ and $\alpha_{i}$ functions of $\mathbf{u}_{1:d}$ instead for $i>d$). In other words, both MAF and IAF can be seen as more flexible (but different) generalizations of Real NVP, where each element is individually scaled and shifted as a function of all previous elements. The advantage of Real NVP compared to MAF and IAF is that it can both generate data and estimate densities with one forward pass only, whereas MAF would need $D$ passes to generate data and IAF would need $D$ passes to estimate densities.

4 Conditional MAF

Given a set of example pairs $\left\{\mathopen{}\left(\mathbf{x}_{n},\mathbf{y}_{n}\right)\mathclose{}\right\}$, conditional density estimation is the task of estimating the conditional density $p\mathopen{}\left(\mathbf{x}\,|\,\mathbf{y}\right)\mathclose{}$. Autoregressive modelling extends naturally to conditional density estimation. Each term in the chain rule of probability can be conditioned on side-information $\mathbf{y}$, decomposing any conditional density as $p\mathopen{}\left(\mathbf{x}\,|\,\mathbf{y}\right)\mathclose{}=\prod_{i}{p\mathopen{}\left(x_{i}\,|\,\mathbf{x}_{1:i-1},\mathbf{y}\right)\mathclose{}}$. Therefore, we can turn any unconditional autoregressive model into a conditional one by augmenting its set of input variables with $\mathbf{y}$ and only modelling the conditionals that correspond to $\mathbf{x}$. Any order of the variables can be chosen, as long as $\mathbf{y}$ comes before $\mathbf{x}$. In masked autoregressive models, no connections need to be dropped from the $\mathbf{y}$ inputs to the rest of the network.

We can implement a conditional version of MAF by stacking MADEs that were made conditional using the above strategy. That is, in a conditional MAF, the vector $\mathbf{y}$ becomes an additional input for every layer. As a special case of MAF, Real NVP can be made conditional in the same way.

Experiments

We systematically evaluate three types of density estimator (MADE, Real NVP and MAF) in terms of density estimation performance on a variety of datasets. Code for reproducing our experiments (which uses Theano ) can be found at https://github.com/gpapamak/maf.

MADE. We consider two versions: (a) a MADE with Gaussian conditionals, denoted simply by MADE, and (b) a MADE whose conditionals are each parameterized as a mixture of $C$ Gaussians, denoted by MADE MoG. We used $C\!=\!10$ in all our experiments. MADE can be seen either as a MADE MoG with $C\!=\!1$, or as a MAF with only one autoregressive layer. Adding more Gaussian components per conditional or stacking MADEs to form a MAF are two alternative ways of increasing the flexibility of MADE, which we are interested in comparing.

Real NVP. We consider a general-purpose implementation of the coupling layer, which uses two feedforward neural networks, implementing the scaling function $f_{\alpha}$ and the shifting function $f_{\mu}$ respectively. Both networks have the same architecture, except that $f_{\alpha}$ has hyperbolic tangent hidden units, whereas $f_{\mu}$ has rectified linear hidden units (we found this combination to perform best). Both networks have a linear output. We consider Real NVPs with either $5$ or $10$ coupling layers, denoted by Real NVP ($5$) and Real NVP ($10$) respectively, and in both cases the base density is a standard Gaussian. Successive coupling layers alternate between (a) copying the odd-indexed variables and transforming the even-indexed variables, and (b) copying the even-indexed variables and transforming the odd-indexed variables.

It is important to clarify that this is a general-purpose implementation of Real NVP which is different and thus not comparable to its original version , which was designed specifically for image data. Here we are interested in comparing coupling layers with autoregressive layers as building blocks of normalizing flows for general-purpose density estimation tasks, and our design of Real NVP is such that a fair comparison between the two can be made.

MAF. We consider three versions: (a) a MAF with $5$ autoregressive layers and a standard Gaussian as a base density $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$, denoted by MAF ($5$), (b) a MAF with $10$ autoregressive layers and a standard Gaussian as a base density, denoted by MAF ($10$), and (c) a MAF with $5$ autoregressive layers and a MADE MoG with $C\!=\!10$ Gaussian components in each conditional as a base density, denoted by MAF MoG ($5$). MAF MoG ($5$) can be thought of as a MAF ($5$) stacked on top of a MADE MoG and trained jointly with it.

In all experiments, MADE and MADE MoG order the inputs using the order that comes with the dataset by default; no alternative orders were considered. MAF uses the default order for the first autoregressive layer (i.e. the layer that directly models the data) and reverses the order for each successive layer (the same was done for IAF by Kingma et al. ).

MADE, MADE MoG and each layer in MAF is a feedforward neural network with masked weight matrices, such that the autoregressive property holds. The procedure for designing the masks (due to Germain et al. ) is as follows. Each input or hidden unit is assigned a degree, which is an integer ranging from $1$ to $D$, where $D$ is the data dimensionality. The degree of an input is taken to be its index in the order. The $D$ outputs have degrees that sequentially range from to $D\!-\!1$. A unit is allowed to receive input only from units with lower or equal degree, which enforces the autoregressive property. In order for output $i$ to be connected to all inputs with degree less than $i$, and thus make sure that no conditional independences are introduced, it is both necessary and sufficient that every hidden layer contains every degree. In all experiments except for CIFAR-10, we sequentially assign degrees within each hidden layer and use enough hidden units to make sure that all degrees appear. Because CIFAR-10 is high-dimensional, we used fewer hidden units than inputs and assigned degrees to hidden units uniformly at random (as was done by Germain et al. ).

We added batch normalization after each coupling layer in Real NVP and after each autoregressive layer in MAF. Batch normalization is an elementwise scaling and shifting operation, which is easily invertible and has a tractable Jacobian, and thus it is suitable for use in a normalizing flow. We found that batch normalization in Real NVP and MAF reduces training time, increases stability during training and improves performance (a fact that was also observed by Dinh et al. for Real NVP). Section B of the appendix discusses our implementation of batch normalization and its use in normalizing flows.

2 Unconditional density estimation

Following Uria et al. , we perform unconditional density estimation on four UCI datasets (POWER, GAS, HEPMASS, MINIBOONE) and on a dataset of natural image patches (BSDS300).

UCI datasets. These datasets were taken from the UCI machine learning repository . We selected different datasets than Uria et al. , because the ones they used were much smaller, resulting in an expensive cross-validation procedure involving a separate hyperparameter search for each fold. However, our data preprocessing follows Uria et al. . The sample mean was subtracted from the data and each feature was divided by its sample standard deviation. Discrete-valued attributes were eliminated, as well as every attribute with a Pearson correlation coefficient greater than $0.98$. These procedures are meant to avoid trivial high densities, which would make the comparison between approaches hard to interpret. Section D of the appendix gives more details about the UCI datasets and the individual preprocessing done on each of them.

Image patches. This dataset was obtained by extracting random $8\!\times\!8$ monochrome patches from the BSDS300 dataset of natural images . We used the same preprocessing as by Uria et al. . Uniform noise was added to dequantize pixel values, which was then rescaled to be in the range $\left[0,1\right]$. The mean pixel value was subtracted from each patch, and the bottom-right pixel was discarded.

Table 1 shows the performance of each model on each dataset. A Gaussian fitted to the train data is reported as a baseline. We can see that on $3$ out of $5$ datasets MAF is the best performing model, with MADE MoG being the best performing model on the other $2$. On all datasets, MAF outperforms Real NVP. For the MINIBOONE dataset, due to overlapping error bars, a pairwise comparison was done to determine which model performs the best, the results of which are reported in Section E of the appendix. MAF MoG ($5$) achieves the best reported result on BSDS300 for a single model with $156.36$ nats, followed by Deep RNADE with $155.2$. An ensemble of $32$ Deep RNADEs was reported to achieve $157.0$ nats . The UCI datasets were used for the first time in the literature for density estimation, so no comparison with existing work can be made yet.

3 Conditional density estimation

For conditional density estimation, we used the MNIST dataset of handwritten digits and the CIFAR-10 dataset of natural images . In both datasets, each datapoint comes from one of $10$ distinct classes. We represent the class label as a $10$-dimensional, one-hot encoded vector $\mathbf{y}$, and we model the density $p\mathopen{}\left(\mathbf{x}\,|\,\mathbf{y}\right)\mathclose{}$, where $\mathbf{x}$ represents an image. At test time, we evaluate the probability of a test image $\mathbf{x}$ by $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}\!=\!\sum_{\mathbf{y}}{p\mathopen{}\left(\mathbf{x}\,|\,\mathbf{y}\right)\mathclose{}p\mathopen{}\left(\mathbf{y}\right)\mathclose{}}$, where $p\mathopen{}\left(\mathbf{y}\right)\mathclose{}\!=\!\frac{1}{10}$ is a uniform prior over the labels. For comparison, we also train every model as an unconditional density estimator and report both results.

For both MNIST and CIFAR-10, we use the same preprocessing as by Dinh et al. . We dequantize pixel values by adding uniform noise, and then rescale them to $\left[0,1\right]$. We transform the rescaled pixel values into logit space by $\mathbf{x}\mapsto\operatorname{logit}\mathopen{}\left(\lambda+\mathopen{}\left(1-2\lambda\right)\mathclose{}\mathbf{x}\right)\mathclose{}$, where $\lambda\!=\!10^{-6}$ for MNIST and $\lambda\!=\!0.05$ for CIFAR-10, and perform density estimation in that space. In the case of CIFAR-10, we also augment the train set with horizontal flips of all train examples (as also done by Dinh et al. ).

Table 2 shows the results on MNIST and CIFAR-10. The performance of a class-conditional Gaussian is reported as a baseline for the conditional case. Log likelihoods are calculated in logit space. MADE MoG is the best performing model on MNIST, whereas MAF is the best performing model on CIFAR-10. On CIFAR-10, both MADE and MADE MoG performed significantly worse than the Gaussian baseline. MAF outperforms Real NVP in all cases. To facilitate comparison with the literature, Section E of the appendix reports results in bits/pixel.In earlier versions of the paper, results marked with * were reported incorrectly, due to an error in the calculation of the test log likelihood of conditional MAF. Thus error has been corrected in the current version.

Discussion

We showed that we can improve MADE by modelling the density of its internal random numbers. Alternatively, MADE can be improved by increasing the flexibility of its conditionals. The comparison between MAF and MADE MoG showed that the best approach is dataset specific; in our experiments MAF outperformed MADE MoG in $5$ out of $9$ cases, which is strong evidence of its competitiveness. MADE MoG is a universal density approximator; with sufficiently many hidden units and Gaussian components, it can approximate any continuous density arbitrarily well. It is an open question whether MAF with a Gaussian base density has a similar property (MAF MoG clearly does).

We also showed that the coupling layer used in Real NVP is a special case of the autoregressive layer used in MAF. In fact, MAF outperformed Real NVP in all our experiments. Real NVP has achieved impressive performance in image modelling by incorporating knowledge about image structure. Our results suggest that replacing coupling layers with autoregressive layers in the original version of Real NVP is a promising direction for further improving its performance. Real NVP maintains however the advantage over MAF (and autoregressive models in general) that samples from the model can be generated efficiently in parallel.

Density estimation is one of several types of generative modelling, with the focus on obtaining accurate densities. However, we know that accurate densities do not necessarily imply good performance in other tasks, such as in data generation . Alternative approaches to generative modelling include variational autoencoders , which are capable of efficient inference of their (potentially interpretable) latent space, and generative adversarial networks , which are capable of high quality data generation. Choice of method should be informed by whether the application at hand calls for accurate densities, latent space inference or high quality samples. Masked Autoregressive Flow is a contribution towards the first of these goals.

We thank Maria Gorinova for useful comments, and Johann Brehmer for discovering the error in the calculation of the test log likelihood for conditional MAF. George Papamakarios and Theo Pavlakou were supported by the Centre for Doctoral Training in Data Science, funded by EPSRC (grant EP/L016427/1) and the University of Edinburgh. George Papamakarios was also supported by Microsoft Research through its PhD Scholarship Programme.

Appendix A Equivalence between MAF and IAF

In this section, we present the equivalence between MAF and IAF in full mathematical detail. Let $\pi_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ be the true density the train data $\left\{\mathbf{x}_{n}\right\}$ is sampled from. Suppose we have a MAF whose base density is $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$, and whose transformation from $\mathbf{u}$ to $\mathbf{x}$ is $f$. The MAF defines the following density over the $\mathbf{x}$ space:

Using the definition of $p_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ in Equation (11), we can write the Kullback–Leibler divergence from $\pi_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ to $p_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ as follows:

The inverse transformation $f^{-1}$ from $\mathbf{x}$ to $\mathbf{u}$ can be seen as describing an implicit IAF with base density $\pi_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$, which would define the following density over the $\mathbf{u}$ space:

By making the change of variables $\mathbf{x}\mapsto\mathbf{u}$ in Equation (13) and using the definition of $p_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ in Equation (14) we obtain

Equation (16) is the definition of the KL divergence from $p_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ to $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$, hence

Suppose now that we wish to fit the implicit density $p_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ to the base density $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ by minimizing the above KL. This corresponds exactly to the objective minimized when employing IAF as a recognition network in stochastic variational inference , where $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ would be the (typically intractable) posterior. The first step in stochastic variational inference would be to rewrite the expectation in Equation (16) with respect to the base distribution $\pi_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ used by IAF, which corresponds exactly to Equation (13). This is often referred to as the reparameterization trick . The second step would be to approximate Equation (13) with Monte Carlo, using samples $\left\{\mathbf{x}_{n}\right\}$ drawn from $\pi_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$, as follows:

Using the definition of $p_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ in Equation (11), we can rewrite Equation (19) as

Since samples $\left\{\mathbf{x}_{n}\right\}$ drawn from $\pi_{x}\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ correspond precisely to the train data for MAF, we can recognize in Equation (20) the training objective for MAF. In conclusion, training a MAF by maximizing its total log likelihood $\sum_{n}\log{p_{x}\mathopen{}\left(\mathbf{x}_{n}\right)\mathclose{}}$ on train data $\left\{\mathbf{x}_{n}\right\}$ is equivalent to variationally training an implicit IAF with MAF’s base distribution $\pi_{u}\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ as its target.

Appendix B Batch normalization

In our implementation of MAF, we inserted a batch normalization layer between every two autoregressive layers, and between the last autoregressive layer and the base distribution. We did the same for Real NVP (the original implementation of Real NVP also uses batch normalization layers between coupling layers ). The purpose of a batch normalization layer is to normalize its inputs $\mathbf{x}$ to have approximately zero mean and unit variance. In this section, we describe in full detail our implementation of batch normalization and its use as a layer in normalizing flows.

A batch normalization layer can be thought of as a transformation between two vectors of the same dimensionality. For consistency with our notation for autoregressive and coupling layers, let $\mathbf{x}$ be the vector closer to the data, and $\mathbf{u}$ be the vector closer to the base distribution. Batch normalization implements the transformation $\mathbf{x}=f\mathopen{}\left(\mathbf{u}\right)\mathclose{}$ defined by

In the above, $\odot$ denotes elementwise multiplication. All other operations are to be understood elementwise. The inverse transformation $f^{-1}$ is given by

and the absolute determinant of its Jacobian is

Vectors $\bm{\beta}$ and $\bm{\gamma}$ are parameters of the transformation that are learnt during training. In typical implementations of batch normalization, parameter $\bm{\gamma}$ is not exponentiated. In our implementation, we chose to exponentiate $\bm{\gamma}$ in order to ensure its positivity and simplify the expression of the log absolute determinant. Parameters $\mathbf{m}$ and $\mathbf{v}$ correspond to the mean and variance of $\mathbf{x}$ respectively. During training, we set $\mathbf{m}$ and $\mathbf{v}$ equal to the sample mean and variance of the current minibatch (we used minibatches of 100 examples). At validation and test time, we set them equal to the sample mean and variance of the entire train set. Other implementations use averages over minibatches or maintain running averages during training . Finally, $\epsilon$ is a hyperparameter that ensures numerical stability if any of the elements of $\mathbf{v}$ is near zero. In our experiments, we used $\epsilon=10^{-5}$.

Appendix C Number of parameters

To get a better idea of the computational trade-offs between different model choices versus the performance gains they achieve, we compare the number of parameters for each model. We only count connection weights, as they contribute the most, and ignore biases and batch normalization parameters. We assume that masking reduces the number of connections by approximately half.

For all models, let $D$ be the number of inputs, $H$ be the number of units in a hidden layer and $L$ be the number of hidden layers. We assume that all hidden layers have the same number of units (as we did in our experiments). For MAF MoG, let $C$ be the number of components per conditional. For Real NVP and MAF, let $K$ be the number of coupling layers/autoregressive layers respectively. Table 3 lists the number of parameters for each model.

For each extra component we add to MADE MoG, we increase the number of parameters by $DH$. For each extra autoregressive layer we add to MAF, we increase the number of parameters by $\frac{3}{2}DH+\frac{1}{2}\mathopen{}\left(L-1\right)\mathclose{}H^{2}$. If we have one or two hidden layers $L$ (as we did in our experiments) and assume that $D$ is comparable to $H$, the number of extra parameters in both cases is about the same. In other words, increasing flexibility by stacking has a parameter cost that is similar to adding more components to the conditionals, as long as the number of hidden layers is small.

Comparing Real NVP with MAF, we can see that Real NVP has about $1.3$ to $2$ times more parameters than a MAF of comparable size. Given that our experiments show that Real NVP is less flexible than a MAF of comparable size, we can conclude that MAF makes better use of its available capacity. The number of parameters of Real NVP could be reduced by tying weights between the scaling and shifting networks.

Appendix D Additional experimental details

MADE, MADE MoG and each autoregressive layer in MAF is a feedforward neural network (with masked weight matrices), with $L$ hidden layers of $H$ hidden units each. Similarly, each coupling layer in Real NVP contains two feedforward neural networks, one for scaling and one for shifting, each of which also has $L$ hidden layers of $H$ hidden units each. For each dataset, we gave a number of options for $L$ and $H$ (the same options where given to all models) and for each model we selected the option that performed best on the validation set. Table 4 lists the combinations of $L$ and $H$ that were given as options for each dataset.

In terms of nonlinearity for the hidden units, MADE, MADE MoG and MAF used rectified linear units, except for the GAS datasets where we used hyperbolic tangent units. In the coupling layer of Real NVP, we used hyberbolic tangent hidden units for the scaling network and rectified linear hidden units for the shifting network.

D.2 Datasets

In the following paragraphs, we give a brief description of the four UCI datasets (POWER, GAS, HEPMASS, MINIBOONE) and of the way they were preprocessed.

POWER. The POWER dataset contains measurements of electric power consumption in a household over a period of $47$ months. It is actually a time series but was treated as if each example were an i.i.d. sample from the marginal distribution. The time feature was turned into an integer for the number of minutes in the day and then uniform random noise was added to it. The date was discarded, along with the global reactive power parameter, which seemed to have many values at exactly zero, which could have caused arbitrarily large spikes in the learnt distribution. Uniform random noise was added to each feature in the interval $\left[0,\epsilon_{i}\right]$, where $\epsilon_{i}$ is large enough to ensure that with high probability there are no identical values for the $i^{\text{th}}$ feature but small enough to not change the data values significantly.

GAS. Created by Fonollosa et al. , this dataset represents the readings of an array of $16$ chemical sensors exposed to gas mixtures over a $12$ hour period. Similarly to POWER, it is a time series but was treated as if each example were an i.i.d. sample from the marginal distribution. Only the data from the file ethylene_CO.txt was used, which corresponds to a mixture of ethylene and carbon monoxide. After removing strongly correlated attributes, the dimensionality was reduced to $8$.

HEPMASS. Used by Baldi et al. , this dataset describes particle collisions in high energy physics. Half of the data are examples of particle-producing collisions (positive), whereas the rest come from a background source (negative). Here we used the positive examples from the “$1000$” dataset, where the particle mass is $1000$. Five features were removed because they had too many reoccurring values; values that repeat too often can result in spikes in the density and misleading results.

MINIBOONE. Used by Roe et al. , this dataset comes from the MiniBooNE experiment at Fermilab. Similarly to HEPMASS, it contains a number of positive examples (electron neutrinos) and a number of negative examples (muon neutrinos). Here we use the positive examples. These had some obvious outliers ($11$) which had values at exactly $-1000$ for every column and were removed. Also, seven of the features had far too high a count for a particular value, e.g. $0.0$, so these were removed as well.

Appendix E Additional results

On the MINIBOONE dataset, the model with highest average test log likelihood is MAF MoG ($5$). However, due to the relatively small size of this dataset, the average test log likelihoods of some other models have overlapping error bars with that of MAF MoG ($5$). To assess whether the differences are statistically significant, we performed a pairwise comparison, which is a more powerful statistical test. In particular, we calculated the difference in test log probability between every other model and MAF MoG ($5$) on each test example, and assessed whether this difference is significantly positive, which would indicate that MAF MoG ($5$) performs significantly better. The results of this comparison are shown in Table 6. We can see that MAF MoG ($5$) is significantly better than all other models except for MAF ($5$).

E.2 Bits per pixel

In the main text, the results for MNIST and CIFAR-10 were reported in log likelihoods in logit space, since this is the objective that the models were trained to optimize. For comparison with other results in the literature, in Table 7 we report the same results in bits per pixel. For CIFAR-10, different colour components count as different pixels (i.e. an image is thought of as having $32\!\times\!32\!\times\!3$ pixels).

In order to calculate bits per pixel, we need to transform the densities returned by a model (which refer to logit space) back to image space in the range $\left[0,256\right]$. Let $\mathbf{x}$ be an image of $D$ pixels in logit space and $\mathbf{z}$ be the corresponding image in $\left[0,256\right]$ image space. The transformation from $\mathbf{z}$ to $\mathbf{x}$ is

where $\lambda\!=\!10^{-6}$ for MNIST and $\lambda\!=\!0.05$ for CIFAR-10. If $p\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ is the density in logit space as returned by the model, using the above transformation the density of $\mathbf{z}$ can be calculated as

where $\sigma\mathopen{}\left(\cdot\right)\mathclose{}$ is the logistic sigmoid function. From that, we can calculate the bits per pixel $b\mathopen{}\left(\mathbf{x}\right)\mathclose{}$ of image $\mathbf{x}$ as follows:

The above equation was used to convert between the average log likelihoods reported in the main text and the results of Table 7.

E.3 Generated images

Figures 2, 3 and 4 show generated images and real examples for BSDS300, MNIST and CIFAR-10 respectively. Images were generated by MAF MoG ($5$) for BSDS300, conditional MAF ($5$) for MNIST, and conditional MAF ($10$) for CIFAR-10.

The BSDS300 generated images are visually indistinguishable from the real ones. For MNIST and CIFAR-10, generated images lack the fidelity produced by modern image-based generative approaches, such as RealNVP or PixelCNN++ . This is because our version of MAF has no knowledge about image structure, as it was designed for general-purpose density estimation and not for realistic-looking image synthesis. However, if the latter is desired, it would be possible to incorporate image modelling techniques in the design of MAF (such as convolutions or a multi-scale architecture as used by Real NVP ) in order to improve quality of generated images.

References