Augmented Normalizing Flows: Bridging the Gap Between Generative Flows and Latent Variable Models

Introduction

Deep invertible models have recently gained increasing interest among machine learning researchers as they constitute a powerful probabilistic toolkit. They allow for the tracking of changes in probability density and have been widely applied in many tasks, including

generative models (Dinh et al., 2017; Kingma & Dhariwal, 2018; Chen et al., 2019),

variational inference (Rezende & Mohamed, 2015; Kingma et al., 2016; Berg et al., 2018),

density estimation (Papamakarios et al., 2017; Huang et al., 2018),

reinforcement learning (Mazoure et al., 2019; Ward et al., 2019), etc.

The main challenges in designing an invertible model for the above use cases are to ensure (1) the mapping $f$ is invertible, (2) the log-determinant of the Jacobian of $f$ is cheap to compute, and (3) $f$ is expressive. For use case (i), ideally we would also like to (4) invert $f$ efficiently.

In general, it is hard to design a family of functions that satisfy all of the above. Most work within this line of research is dedicated to improving the expressivity of the bijective mapping while maintaining the computational tractability of the log-determinant of Jacobian (Dinh et al., 2017; Kingma et al., 2016; Huang et al., 2018; Chen et al., 2019).

Aside from the unfortunate trade-off between the computational budget of inversion/Jacobian log-determinant and the expressivity of the invertible mapping, generative flows suffer from the limitation of local dependency. Unlike latent variable models such as Variational Autoencoders (VAEs; Kingma & Welling 2014; Rezende et al. 2014) and Generative Adversarial Networks (GANs; Goodfellow et al. 2014) which model the high dimensional data as coordinates in another space, most generative flows model the dependency among features only locally. Dependencies of features far away from each other can only be propagated through composition of mappings, which progressively enlarges the receptive field. Special design of parameterization like the attention mechanism can be made to address this issue (Ho et al., 2019).

In this paper, we propose to construct an invertible model on an augmented input space, which when combined with the block-wise coupling of Dinh et al. (2017) satisfies all criteria (1-4). The motivation is that to transform some distribution (such as the marginal distribution of $x$ pictured in Figure 1) into another (e.g. standard normal) in the original input space, $f$ needs to be capable of transporting the probability mass “non-uniformly” across its domain, whereas in an augmented input space it is possible to find a smoother transformation. For instance, if we couple the data $x$ with an independent noise $e$, we can first transform $e$ conditioned on $x$ into $z$, so that conditioned on different values of $z$, $x$ can be more easily centered and Gaussianized. Our proposed method also generalizes multiple variants of VAEs and possesses the advantage of transforming the data in a more globally coherent manner via first embedding the data in the augmented state space. Finally, operating on an augmented state space allows us to sidestep the topology preserving property of a diffeomorphism, which means input space can potentially be more freely deformed (Dupont et al., 2019).

we introduce Augmented Normalizing Flows (ANFs), an invertible generative model on the real-valued data $x$ coupled with an independent noise $e$. We propose a parameter estimation principle called Augmented Maximum Likelihood Estimation (AMLE), which we show amounts to maximizing a lower bound on the marginal likelihood of the original data $x$. Theoretically, we show that the family of ANFs with additive coupling can universally transform arbitrary data distribution into a standard Gaussian prior, augmented with a degenerate deterministic variable. To the best of our knowledge, this is the first attempt in understanding how expressivity can be improved via composing flow layers rather than widening the flow (Huang et al., 2018). Experimentally, we apply the proposed method to a suite of standard generative modelling tasks and demonstrate state-of-the-art performance in density estimation and image synthesis.

Background

Given a training set $(x_{i})_{i=1}^{n}\sim q(x)^{n}$, where $x_{i}\in{\mathcal{X}}$, and a family of density models $\{p_{\pi}(x):\pi\in\mathfrak{P}({\mathcal{X}})\}$, where $\mathfrak{P}({\mathcal{X}})$ is a collection of sets of parameters that can sufficiently describe the density function, the Maximum Likelihood Principle estimates the parameters by maximizing the chance of the data being generated by the assumed model:

where the latter expectation is over the empirical distribution $\hat{q}(x)$ ($x_{i}$ with uniformly distributed random index $i\in\{1,\cdots,n\}$). $\hat{\pi}$ is known as the maximum likelihood estimate (MLE) for the parameter $\pi$. Below, we review two families of likelihood-based density models.

Assume $y\sim{\mathcal{N}}(0,I)$. Assume the data is generated via a bijective mapping $x=f_{\theta}(y)$. Then the probability density function of $f_{\theta}(y)$ evaluated at $x$ can be written as

Equivalently, one can parameterize the inverse transformation $x\mapsto g_{\theta}(x)$ with invertible mapping $g_{\theta}$, and define the generative transformation as $f_{\theta}=g_{\theta}^{-1}$.

Much of the design effort has been dedicated to ensuring (1) the invertibility of the transformation $g$, and (2) efficiency in computing the log-determinant of the Jacobian in Equation 2. For example, Dinh et al. (2017) propose the affine coupling:

where $s_{\theta}$ and $m_{\theta}$ are parameterized by neural networks and $x_{a}$ and $x_{b}$ are two partitioning of the data vector, and compose multiple layers of transformations intertwined with permutation of elements of $x$.

Invertible models allow for exact computation of the likelihood, and can be composed to increase modelling capacity. Nevertheless, the expressivity of the transformation is limited due to the need to satisfy invertibility and to reduce the cost of computing the Jacobian determinant. For a comprehensive review of this topic, see Kobyzev et al. (2019) and Papamakarios et al. (2019).

Variational Autoencoders

Assume the data follows the generating process: $x\sim p_{\theta}(x|z)$ where $z\sim p_{\theta}(z)$. For simplicity, we assume $p_{\theta}(z)$ is the standard Gaussian distribution and drop the dependency on $\theta$ henceforward. Our goal is to find the MLE for $\theta$, but the log marginal density $\log p_{\theta}(x)=\log\int_{z}p_{\theta}(x|z)p(z)dz$ is generally not tractable since it involves integration. Instead, one can maximize a surrogate objective known as the evidence lower bound (ELBO):

where $q_{\phi}(z|x)$ is an inference network that amortizes the cost of parameterizing the variational distribution per input instance $x$ via conditioning. Learning and inference can be jointly achieved by drawing a stochastic estimate of the gradient of the ELBO via reparameterization (i.e. change of variable):

if $g_{\phi}(x,e)$ with $e\sim q(e)$ follows the same density as $q_{\phi}(z|x)$. Conventionally, $q_{\phi}(z|x)$ is a multivariate Gaussian distribution with diagonal covariance. We write it as ${\mathcal{N}}(z;\mu_{\phi}(x),\sigma^{2}_{\phi}(x))$. One choice of reparameterization is $g_{\phi}(x,e)=\mu_{\phi}(x)+\sigma_{\phi}(x)\odot e$ with $q(e)={\mathcal{N}}(0,I)$.

VAEs allow one to embed the data in another space (usually of lower dimensionality), and to generate via an arbitrarily parameterized mapping. However, the log likelihood of the data is no longer tractable, so we can only maximize an approximate log likelihood. The performance of the model highly depends on the choice of the encoding distribution and the decoding distribution, as they are closely related to the tightness of the lower bound (Cremer et al., 2018).

Augmented Maximum Likelihood

For augmented maximum likelihood, we couple each data point with an independent random variable $e\in{\mathcal{E}}$ drawn from $q(e)$ (in all our experiments we set $q(e)={\mathcal{N}}(0,I)$), and consider a family of joint density models $\{p_{\pi}(x,e):\pi\in\mathfrak{P}({\mathcal{X}}\times{\mathcal{E}})\}$. Instead of maximizing the marginal likelihood of $x_{i}$’s, we maximize the joint likelihood:

where the expectation is over $(x,e)\sim\hat{q}(x)q(e)$. We refer to this extremum estimator as the Augmented Maximum Likelihood Estimator (AMLE). The benefit of maximizing the joint likelihood is that it allows us to make use of the augmented state space to induce structure on the marginal distribution of $x$ in the original input space.

Since KL is non-negative, maximizing the joint likelihood according to Equation 5 is equivalent to maximizing a lower bound on the log marginal likelihood of $x$. We refer to this as the Augmentation Gap, as it reflects the incapability of the joint density to model the marginal of $e$ independently of $x$.

Estimating the log marginal likelihood

The log marginal likelihood $\log p_{\pi}(x)$ of the data can be estimated in a way similar to Burda et al. (2015), by drawing $K$ i.i.d. samples of $e_{j}\sim q(e)$ per $x$ to estimate the following stochastic lower bound:

which can be shown to be a consistent estimator for $\log p_{\pi}(x)$ and is monotonically tighter in expectation as we increase $K$.

Augmented Normalizing Flows (ANF)

We now demonstrate how to leverage the augmented input space to model the complex marginal distribution of the data. We consider maximizing the joint likelihood of $x$ coupled with a random noise $e\sim q(e)$. Let $(y,z)\sim p(y,z)$ be drawn from some simple distribution, such as independent Gaussian. Assume the data $x,e$ is deterministically generated via an invertible mapping $x,e=F_{\pi}(y,z)$, with inverse $G_{\pi}=F_{\pi}^{-1}$. Then analogous to Equation 2, $x,e$ has a joint density

For simplicity, we can choose $q(e)$ to be the standard normal distribution. What we are left with is the choice of an invertible $G_{\pi}$ that can harness the augmented state space ${\mathcal{E}}$ to induce a complex marginal on ${\mathcal{X}}$. Inspired by the affine coupling proposed by Dinh et al. (2017), we conditionally transform $x$ and $e$, hoping the structure in the marginal of $x$ can “leak” into ${\mathcal{E}}$ and make the joint more easily Gaussianized. Concretely, we define two types of affine coupling

We refer to the pair of encoding transform and decoding transform as the autoencoding transform. We stack them up in alternating order, i.e. $G_{\pi}=g_{\pi_{N}}^{\text{dec}}\circ g_{\pi_{N}}^{\text{enc}}\circ...\circ g_{\pi_{1}}^{\text{dec}}\circ g_{\pi_{1}}^{\text{enc}}$ for $N\geq 1$ steps, where $\pi=\{\pi_{1},...,\pi_{N}\}$ is the set of all parameters. See Figure 2-(a,b) for an illustration.

Variational Autoencoders are a special case of augmented normalizing flows with only “one step” of encoding and decoding transform (Dinh et al., 2014). To see this, assume the decoding distribution $p_{\theta}(x|z)$ is a factorized Gaussian with mean $\mu_{\theta}(z)$ and standard deviation $\sigma_{\theta}(z)$. By letting $z=\mu_{\phi}(x)+\sigma_{\phi}(x)\cdot e$ and $y=(x-\mu_{\theta}(z))/\sigma_{\theta}(z)$ and applying the change of variable formula to both $q_{\phi}(z|x)$ and $p_{\theta}(x|z)$, we get from Equation (4)

Averaging over the data distribution $\hat{q}(x)$, we obtain the expected joint likelihood (up to the constant $H(e)$)

The variational gap between the log marginal likelihood and the evidence lower bound is equal to the augmentation gap since the KL divergence is invariant under the transformation between $e\longleftrightarrow z$:

This gives us an alternative interpretation of inference suboptimality (Cremer et al., 2018): the inaccuracy of inferring the true posterior $p(z|x)$ can be attributed to the incapability of the joint density to model the augmented data $q(e)$.

To illustrate this phenomenon, we model the density of a one dimensional mixture of Gaussian (1D MoG). In Figure 3 (left), we plotted the density histograms of the MoG distribution (blue) and a one-step ANF, i.e. VAE with Gaussian encoder and decoder (orange), trained on the MoG samples. Not surprisingly, the latter fails to represent two well separated modes of probability mass. In Figure 3 (right), we visualize the joint density of the augmented data $x,e\sim q(x)q(e)$ throughout the transformation. We see that the transformed data $y,z=g^{dec}_{\pi_{1}}(g^{enc}_{\pi_{1}}(x,e))$ is not perfectly Gaussianized. In fact, if we project it horizontally we can see that the “aggregated posterior” (marginal of $z$) does not match the prior distribution $p(z)$. As a result, the pushforward $x,e=g^{enc,-1}_{\pi_{1}}(g^{dec,-1}_{\pi_{1}}(y,z))$ of $y,z\sim p(y,z)$ does not follow the augmented data distribution $q(x)q(e)$ well. When we fix different values of $x$, we have different slices of density functions for $e$, indicating that $e$ and $x$ are dependent and that $p_{\pi}(e|x)$ deviates from $q(e)$.

We carry out the same experiment on 1D MoG with multiple flow layers, which generalizes a VAE with Gaussian encoder and decoder. We set the number of flow layers (i.e. steps) to be $5$. To furthermore demonstrate the benefit of transformation composition, we also tie the parameters of each encoder and decoder step, separately. That is, the same set of parameters are used at different steps of encoding and decoding to make sure capacity stays constant. Since the conditional independence assumption in VAE is relaxed, the augmented data is more successfully Gaussianized, as can be seen in Figure 4. The generated samples also follow the target joint density more closely.

2 Hierarchical Augmented Normalizing Flows

The information flow of the encoding-decoding transform just described is limited to the size of the random vector $e$, which makes it hard to optimize for more realistic settings such as natural images. We thus propose a second architecture by following the hierarchical variational autoencoder, which is defined by two pairs of joint distributions This particular factorization of the variational distribution is known as the bottom-up inference. We leave the top-down inference (Kingma et al., 2016) and the bidirectional inference (Maaløe et al., 2019), which benefit more from parameter sharing, for future work.

When all the conditionals are Gaussian distributions, the corresponding ELBO can be similarly rearranged to be the loss function of an ANF (see Figure 2-(c)). The encoding transform for each $e_{l}$ is conditioned on the “transformed” preceding variables $s_{\pi,l}^{e}(x,z_{<l})\odot e_{l}+m_{\pi,l}^{e}(x,z_{<l})$ due to the conditioning in $q(z_{l}|z_{<l},x)$. The decoding transform on the other hand is conditioned on the “original” preceding variables $s_{\pi,l}^{d}(e_{>l})\odot e_{l}+m_{\pi,l}^{d}(e_{>l})$, which is block-wise inverse autoregressive (Kingma et al., 2016). When the conditioning mappings are convolutional, the lower level transformation preserves information of the input locally, which is then combined with the deterministic path of the decoding that “sees” more of the input. More details on the architecture are described in Appendix B.1.

ANFs as Approximate Hamiltonian ODE and Universality

The affine-coupling autoencoding transform with augmented variable is reminiscent of the leap-frog integration of the Hamiltonian system (Neal et al., 2011). More recently, it has been shown by Taghvaei & Mehta (2019) that solving a family of Hamiltonian ordinary differential equations (ODE) with an infinite time horizon gives us a transport map from the initial (data) distribution to an arbitrary target distribution with a log-concave density function. This suggests we can develop an approximation theorem by using ANFs to approximately, numerically solve the ODE.

where $\dot{x}_{t}$ and $\dot{e}_{t}$ are the time derivatives of $x$ and $e$ at time $t$, and $q_{t}$ is the marginal density of $x_{t}$.

Second, we construct a sequence of encoding and decoding functions $m^{\text{enc}}_{n}$ and $m^{\text{dec}}_{n}$ parameterized by neural networks, and define the following (additive) invertible mappings

with $e^{\pi}_{0}=0$ and $x^{\pi}_{0}\sim q_{0}$. The step size parameter $\epsilon$ will be chosen to depend on the depth coefficient $N$, i.e. the number of steps of the joint transformation.

Assume our target distribution lies within a family of distributions ${\mathcal{Q}}$ satisfying Assumption 1 in the Appendix F (some smoothness condition on the time derivatives and $\Phi$). We can then set the encoding and decoding functions to be arbitrarily close to the time derivatives by the universal approximation of neural networks (Cybenko, 1989), and by taking the depth $N$ to be arbitrarily large, we can approximate the transport map induced by the Hamiltonian ODE arbitrarily well, which gives rise to the following universal approximation theorem (the proof is relegated to the Appendix F): {thm}[] For any $q\in{\mathcal{Q}}$, we can find a sequence $(x^{\pi}_{N},e^{\pi}_{N})$ of ANFs of the additive form (7,8), such that if $x^{\pi}_{0},e^{\pi}_{0}\sim q(x)\delta_{0}(e)$ and $x_{\infty},e_{\infty}\sim p(x)\delta_{0}(e)$, then $(x^{\pi}_{N},e^{\pi}_{N})\rightarrow(x_{\infty},e_{\infty})$ in distribution.

Related Work

In the literature of normalizing flows, much work has been done to improve expressivity while maintaining computational tractability. For example, Dinh et al. (2014, 2017) introduce an affine coupling that partitions the features into two groups so that the Jacobian is a block-triangular matrix. The resulting mapping is relatively restricted since it only models partial dependency. Kingma et al. (2016) further exploits the ordered dependency by constructing an inverse autoregressive mapping but its inversion requires a computation time linear in dimensionality (Papamakarios et al., 2017), and does not even have a closed-form formula in the more general non-affine setting (Huang et al., 2018). Behrmann et al. (2018) propose a residual form of $f$ whose Jacobian log-determinant can be stochastically estimated (Chen et al., 2019) but inversion is achieved iteratively, not in one pass.

Normalizing flows have also been used as (1) an inference machine in the context of variational inference for continuous latent variable models (Kingma et al., 2016; Tomczak & Welling, 2016; Berg et al., 2018), and (2) a trainable component of the latent variable model (Chen et al., 2017; Agrawal & Dukkipati, 2016; Huang et al., 2017). ANFs lie at the intersection of normalizing flows and latent variable models when a specific type of block-conditioning transformation is applied, and allow us to unifyingly view flow-based priors as marginal transformation in the space of $e$, and amortized flows for improving posterior inference as different variants of the encoding transform. Another way of improving the inference machine’s expressivity is to consider a hierarchical model; in fact, ANFs can be viewed as a generalization of the auxiliary variable method for hierarchical variational inference (Agakov & Barber, 2004; Ranganath et al., 2016); see Appendix D for the connection and C for more discussion on future direction.

Finally, Dupont et al. (2019) also employs augmentation to improve the expressivity and stability of a neural ODE (Chen et al., 2018a), and they believe such a method can be used to reduce the cost of training a continuous normalizing flow (Grathwohl et al., 2019).

Large-Scale Experiments

In the more realistic settings, we augment the data with a hierarchy of noise, as described in the last part of Section 4. See Appendix B for more experimental details.

We conduct an ablation study on the effect of composing multiple encoding-decoding transformations ($N$ steps) versus increasing the number of stochastic layers ($L$ layers). We monitor the bits per dim (BPD) of the test set of CIFAR 10 (Krizhevsky et al., 2009) throughout training. Figure 5 shows that increasing the number of flow layers can more effectively improve the likelihood of the model than increasing the number of stochastic layers.

Density estimation

We perform density modelling on the MNIST handwritten digit dataset (LeCun et al., 1998), CIFAR 10 (Krizhevsky et al., 2009), downscaled versions of ImageNet ($32\times 32$ and $64\times 64$) (Oord et al., 2016) and the celebrity face dataset CelebA (Liu et al., 2015), and compare with other state-of-the-art density models. In Table 1, we see that ANFs set a few new records in terms of BPD on the standard benchmarks in the non-autoregressive category. We use the importance sampling described in Section 3 to estimate the log likelihood. The augmentation gap is around 0.01 BPD for all benchmarks, indicating the augmented flow is capable of achieving good likelihood estimate and high inference precision at the same time.

2 Qualitative results

For quantitative evaluation of sample quality, we report the Inception Score (IS) (Salimans et al., 2016) and the Fréchet Inception Distance (FID) (Heusel et al., 2017), expanding the table of Ostrovski et al. (2018). We found the FID score of WGAN-GP (Gulrajani et al., 2017a) reported in Ostrovski et al. (2018) is worse than the one reported in the literature, so we include the original values of IS and FID of GANs from the original works of Gulrajani et al. (2017a) and Heusel et al. (2017) for more realistic comparison. In Table 2, we see that ANF obtains better scores than all the other explicit density models, and is close to matching the FID of the orignal WGAN-GP by Gulrajani et al. (2017a). The generated samples are presented in Figure 6 and Appendix E. Since the encoding-decoding transformation has a receptive field that is wide enough to cover the entire raw data, the generated samples also look more globally coherent.

Lossy reconstruction

As a hierarchical model, ANF can be used to perform inference for the higher level representation, and sample the lower level details for reconstruction. We do this by sampling $e_{1},...,e_{L}$, obtaining the corresponding $y,z_{1},...,z_{L}\leftarrow G_{\pi}(x,e_{1},...,e_{L})$, randomizing all but the last representations $y^{\prime},z_{1}^{\prime},...,z_{L-1}^{\prime}\sim{\mathcal{N}}(0,I)$, and reconstructing from the new joint representation $x^{\prime},e_{1}^{\prime},...,e_{L}^{\prime}\leftarrow G_{\pi}^{-1}(y^{\prime},z_{1}^{\prime},...,z_{L-1}^{\prime},z_{L})$. Similar to other hierarchical models (Gulrajani et al., 2017b; Belghazi et al., 2018), ANF is also capable of retaining global, semantic information of the raw data stored in its higher level code; this is shown in Figure 7.

Interpolation

We also perform interpolation in the latent space between real images. Previous works such as Kingma & Dhariwal (2018) perform linear interpolation of the form $h(u,v,t)=tu+(1-t)v$ for $t\in$, which we observe has a non-smooth transition (e.g. sudden color change). We hypothesize this is due to the fact that convex combination of two vectors would result in an increase and then a decrease in the density of the standard Gaussian prior. This is undesirable since the interpolated points are atypical because Gaussian samples are known to concentrate around the shell (of radius proportional to square root dimensionality). Hence, we propose the rescaled interpolation

where $||\cdot||$ denotes the L2 norm, to make sure the scale of the resulting point is a linear interpolation of the scales of the two input vectors. The result in Figure 8 shows that the transition is extremely smooth (see Appendix A for a side-by-side comparison with linear interpolation) and the intermediate images are realistic looking.

Conclusion

In this work, we propose the Augmented Normalizing Flows and a corresponding variational lower bound on the marginal likelihood. We show that the proposed method can be used to approximate a Hamiltonian dynamical system as a universal transport map and achieves competitive or better results than state-of-the-art flow-based methods.

Acknowledgements

CW would like to thank Matt Hoffman for a discussion on deterministic Langevin transitions and Amirhossein Taghvaei for referencing the work of Wang & Li. Special thanks to people who have provided their feedback and advice during discussion or while reviewing the manuscript, including Valentin Thomas, Joey Bose, and Eeshan Dhekane; to Taoli Cheng and Bogdan Mazoure for volunteering for the internal review at Mila; and to Ahmed Touati, Christos Tsirigotis and Jose Gallego for proofreading parts of the proof.

References

Appendix A Interpolation

We compare linear interpolation with rescaled interpolation (rescaling is done separately for each stochastic layer). We see that the middle points of linear interpolation tend to be more yellowish, and rescaled interpolation results in a smoother and direct transition between two input vectors.

Appendix B Experiment Details

To model natural images, we employ a more intricate architecture with a higher modeling capacity described in B.1. Section B.3 describes the parameter initialization scheme and parameterization constraints that are imposed to stablize training. In B.4, we propose an objective-annealing technique that biases the autoencoding transform to focus on Gaussianizing the raw data more at the early stage of training. We found this to be helpful for optimization. All the hyperparameters used in the experiments are summarized in B.5.

We parameterize all of the feedforward layers with weight normalization (Salimans & Kingma, 2016). For the encoding transforms, we first map $x$ using a convolutional layer composed with an activation function, followed by a pooling layer to obtain the first set of conditional features $h$. We now use these features to conditionally transform each of the stochastic units in the order $e_{1},...,e_{L}$, using an encoding block. The encoding block outputs a set of modified conditional features, which is then used to modify the next $e_{l}$ (except when the feature map size is halved, in which case an additional pair of convolutional layer and pooling layer is first applied).

Each encoding block has a nested residual structure, taking in the conditional features as input to transform the corresponding stochastic unit $e$, illustrated in Figure 10. The conditional features are first convolved and then fed into a nested Real NVP block. The output of the Real NVP block (applied twice, see below) is then convolved and added to the convolved conditional features. We apply non-linearity before convolving again and adding to the the original conditional features via skip connection.

The decode transform is similar, except we convolve the incoming stochastic unit to modify the conditional features, thus having a shorter computational path to reconstruct the data.

Real NVP block

The stochastic unit is split into two halves ($e_{1}$ and $e_{2}$) using the checkerboard mask. The convolved conditional feature is concatenated with the masked stochastic unit (the part that is not masked is denoted as $e_{1}$) to transform the part of the stochastic unit being masked out ($e_{2}$). The same Real NVP block is reused (using the same set of parameters), alternating the pattern of the mask to transform the other half of the stochastic unit (with $e_{1}$ and $e_{2}$ swapped). We found sharing the parameters of these two consecutive Real NVP transformations to improve the convergence of the likelihood.

B.2 Variational dequantization

We use the variational dequantization proposed by Ho et al. (2019) in all our density estimation experiments. We first map the input image $x$ into a deterministic feature $x^{\prime}$ space using a convolutional network of the following form:

where act denotes activation function. Note that the input to this convolution network is rescaled to $$via$x/(2^{\texttt{n\_bits}}-1)$, where n_bits is the number of bits. We then use the Real NVP block to transform a standard Gaussian noise, where$x^{\prime}$acts as the conditional feature. We apply two Real NVP blocks to obtain$u\_logit$, and transform it into$u$using the logistic sigmoid activation function so that each element of$u$lies within$(0,1)$. We then perturb the data by$(x+u)/2^{\texttt{n\_bits}}$, which is then passed through a clip operator for numerical stability. We define clip to be

where $\delta=0.1$. Finally, we pass the clipped value into the logit function (inverse sigmoid) to obtain the dequantized data. We have taken into account the probability density of $u\_logit$ and all the changes of variable (i.e. sigmoid, rescaling, clipping, and logit) when computing the lower bound.

B.3 Initialization and parameterization constraint

Unless it is otherwise stated, we initialize all the convolutional kernels using truncated normal distribution with standard deviation $0.1$, and for weight normalization, the rescaling parameter is set to be $1.0$ and the shifting parameter $0.0$. Only for the last layer of the Real NVP block we initialize $g$ to be $0.0$. We apply a split operator to this last layer to obtain a “shift” coefficient and “log scale” coefficient for affine transformation. The last layer has double the dimensionality of the stochastic unit to be transformed. The split operator simply splits it in two parts. For the log scale coefficient, we apply the log-sigmoid function to make sure after exponentiation, it is bounded between and $1$ (similar to Kingma & Dhariwal (2018)). Since the pre-log-sigmoid is initialized to be , we add in a constant that depends on the total number of transformations that will be apply to the stochastic unit such that the overall transformation (after composition) will rescale the raw input by a factor of $0.95$ (without considering activation normalization). This is to ensure the entire transformation is more robust to variation of depths at initialization.

We also apply activation normalization (Kingma & Dhariwal, 2018) with data-dependent initialization that standardizes the transformed feature, after each encoding transform and each decoding transform. We clip the log scale coefficient at $\pm 2.5$.

B.4 Deterministic warm up

Due to our choice of flow, our instantiation of ANF resembles a VAE. It has been previously shown that starting off with less regularization and noise injection is beneficial to training, a technique known as deterministic warm up (Raiko et al., 2007; Sønderby et al., 2016). Similarly, if we expand the objective of ANF with autoencoding transform (affine coupling), we get

where $d$ is the dimensionality of the augmented data $e$, and $(y_{t},z_{t})$ are defined recursively as

with the initial values $z_{0}=e$, $y_{0}=x$. We modify the objective by lowering the weighting of $s_{\pi_{t}}^{\text{enc}}$ and $\log{\mathcal{N}}(z_{T};0,I)$ such that the network can focus more on Gaussianizing the raw input $x$. We defined the modified objective as

where $\pi$ is all the trainable parameters. We linearly anneal the weighting coefficient $\beta$ from to $1$ for the first $\alpha$ iterations of the training. Note that in practice we apply the same $\beta$ to all augmented data in the hierarchical setup.

B.5 Hyperparameters

$L$: number of stochastic units ($e_{1},...,e_{L}$).

$N$: number of steps (encoding-decoding pairs).

$K$: number of samples for importance sampling.

$\lambda$: decoupled weight decay coefficient for Adam.

$c$: number of channels (all deterministic features).

$c^{\prime}$: number of channels for the $l$’th stochastic unit (stochastic features). Power denotes repetition.

$k$: kernel size (except for the data space layer).

$a$: annealing schedule (number of parameter updates).

Appendix C Extended related work and future direction

The term Normalizing Flow was originally coined by Tabak et al. (2010); Tabak & Turner (2013) where it was used for density estimation. Differentiable bijective models were first introduced to the deep learning community as likelihood-based generative models by Rippel & Adams (2013); Dinh et al. (2014), and as an inference machine by Rezende & Mohamed (2015). Most development within this line of research is dedicated to improving the expressivity of the bijective mapping while maintaining computational tractability of the log-determinant of the Jacobian. Each family of flows can be characterized by the “trick” used to achieve this, e.g.

Partial ordered dependency. If the mapping has a partial and ordered dependency, its Jacobian matrix will be a triangular matrix, the determinant of which can be computed in linear time. This includes the following:

Dinh et al. (2014, 2017); Kingma & Dhariwal (2018); Ho et al. (2019) use a block-wise conditioning in the mapping, and

Kingma et al. (2016); Chen et al. (2017); Papamakarios et al. (2017); Huang et al. (2018) generalize block-wise dependency to temporal dependency wherein all the variables prior to the current variable of a given ordering are inputs of the conditioning to transform the current variable.

Low rank transform. If the mapping is of a particular residual form, the Jacobian determinant can be computed readily using the matrix determinant lemma (Rezende & Mohamed, 2015) or its higher rank generalization (Berg et al., 2018).

Lipschitz residual flow. If the nonlinear block of a residual mapping is no more than 1-Lipschitz, the overall mapping is invertible, and its Jacobian can be estimated using power series expansion, the Hutchinsons trace estimator (Behrmann et al., 2018) and the Russian roulette estimator (Chen et al., 2019).

Special convolutional forms. Certain structure of convolutional kernels can also be designed to to ensure tractability, such as via using $1\times 1$ convolution (Kingma & Dhariwal, 2018), masking (Oord et al., 2016; van den Oord et al., 2016; Hoogeboom et al., 2019; Song et al., 2019; Ma et al., 2019) or imposing certain repeated structure (Karami et al., 2019).

In this work, we introduce the augmentation trick, which generalizes flow-based methods in an orthogonal yet complementary manner. In particular, we employ the coupling proposed by Dinh et al. (2017) to transform the augmented data; one potential alternative is to replace it with any of the tricks mentioned above.

Architectures and parameter sharing.

As the autoencoding transform we use generalizes VAEs and hierarchical VAEs, one potential direction is to consider parameterizations that have shared components which are shown to be conducive to training, such as the ResNet with top-downn inference (Kingma et al., 2016) and the bidirectional inference machine (Maaløe et al., 2019). As a generalization of VAEs, ANFs can also be applied to latent variable models of different graphical representations, such as variational recurrent neural networks (Chung et al., 2015) and models of sets (Edwards & Storkey, 2017); for example, the set flow proposed by Rasul et al. (2019) is an instance of permutation-invariant ANF applied to sets. Another avenue for improving parameter efficiency is to consider tying the weights of different steps of transformations. As our theory suggests, consecutive transformations of the discretized Hamiltonian ODE would differ only slightly if the time derivatives are smooth enough. This means it would be sufficient to consider a single network which also takes in time embedding as input for all transformations.

Approximate Hamiltonian flows.

Our approximation theory builds on the result of Wang & Li (2019), which follows the optimal control framework of Wibisono et al. (2016). The augmented variable is treated as the costate, which is deterministically dependent on the state, i.e. the input data. Therefore we set the initial augmented distribution to be a Dirac point mass for the approximation theory to hold. Our theorem can be improved if one can show some time trajectories with the augmented variable drawn independently from a non-degenerate initial distribution are convergent to the prior distribution. We leave this for future work. Meanwhile, the same proof technique can be used to study the approximation capability of different families of flows. In particular, the residual flows (Behrmann et al., 2018; Chen et al., 2019) and their continuous counterpart (Chen et al., 2018a; Grathwohl et al., 2019) can be used to approximate the deterministic Langevin diffusion, since (1) one can replace the Brownian motion term with the gradient of the log marginal density without modifying its Fokker-Planck equation (see Hoffman & Ma (2019) or the appendix of Wang & Li (2019)) and (2) the first-order Langevin dynamic is known to be convergent to its stationary distribution (Roberts et al., 1996).

Gradient-based flows.

As the theory suggests, gradient of the potential can be used to guide the evolution of the particle. This has been previously explored by Duvenaud et al. (2016). Salimans et al. (2015) on the other hand propose a hierarchical model inspired by the Hamiltonian dynamic, and Song et al. (2017); Levy et al. (2018) generalize Hamiltonian Monte Carlo (HMC) with trainable neural components. Similarly, one can parameterize a Gradient-based augmented generative flow to model the data distribution.

Normalizing flows for variational inference.

The most well-known application of normalizing flows is to improve the variational distribution to approximate posterior distribution of (1) the latent representations (Rezende & Mohamed, 2015; Kingma et al., 2016; Tomczak & Welling, 2016; Berg et al., 2018) and (2) the parameters of neural networks (Louizos & Welling, 2017; Krueger et al., 2017; Huang et al., 2019). ANF can also be applied to inference problems, with slight modification of the target potential. We show in Appendix D that one can augment the target distribution with an independent distribution and infer the joint target altogether. This boils down to the hierarchical variational method (Agakov & Barber, 2004; Ranganath et al., 2016) as a special case when one step of autoencoding transform is applied.

Variational gap.

The joint likelihood that we maximize is a variational objective lower-bounding the marginal likelihood of the data. One potential avenue for improvement is to reduce this bias (the augmentation gap) throughout training, by closing up the gap via importance sampling (Burda et al., 2015) or using an unbiased estimate of the marginal likelihood (Luo et al., 2020).

Representation learning.

Considering invertible transformations in an augmented data space allows us to sidestep the topology-preserving property of a homeomorphism. The issue of this property is discussed and addressed by Cornish et al. (2019) by converting the flow into a latent-variable model. Dupont et al. (2019) adopt the same technique by augmenting the data space and apply the augmented continuous time flow to discriminative tasks. We hypothesize this can potentially improve the representation learned by an invertible model, for example in a semi-supervised setting (Nalisnick et al., 2019; Atanov et al., 2019) or as a component of a reversible model for memory-efficient backpropagation (Gomez et al., 2017).

Appendix D Augmented Normalizing Flows for Variational Inference

where $|_{1}$ and $|_{2}$ denote the first and the second coordinates, respectively. This lower-bounds the ELBO since

Auxiliary variable for hierarchical variational inference. The above derivation for applying ANF to variational inference is reminiscent of the auxiliary variable method (Agakov & Barber, 2004; Ranganath et al., 2016). To see this, assume we parameterize $G(e,u)$ as the composition $g^{\text{enc}}\circ g^{\text{dec}}$, where

with $s^{\text{enc}},s^{\text{dec}}>0$. Then Equation (10) becomes

where $z:=s^{\text{dec}}(u)\odot e+m^{\text{dec}}(u)$, which is equivalent to

where $q(z|u)={\mathcal{N}}(z;m^{\text{dec}}(u),s^{\text{dec}}(u)^{2})$ and $r(u|z)={\mathcal{N}}(u;-m^{\text{enc}}(e)/s^{\text{enc}}(e),s^{\text{enc}}(e)^{-2})$. This shows hierarchical variational methods are a special case of ANF, and the latter can potentially be used to improve the joint expressivity of the former through additional composition.

Appendix E More samples

E.2 Celeba 64

Appendix F Proofs

where $\dot{x}_{t}$ and $\dot{e}_{t}$ are the time derivatives of $x$ and $e$ at time $t$, and $q_{t}$ is the marginal density of $x_{t}$.

For some convex $\Phi$, the trajectories of $x_{t}$ and $e_{t}$ following (11,12) converge in distribution to $x_{\infty}$ and $e_{\infty}$, respectively, where $x_{\infty}\sim p(x)$ and $e_{\infty}\sim\delta_{0}$ (i.e. a point mass at ).

This implies for any bounded, continuous $f$,

which converges to as $t\rightarrow\infty$. By Portmanteau’s Lemma, $x_{t}\rightarrow x_{\infty}$ in distribution. ∎

We first construct a sequence of encoding functions $m^{\text{enc}}_{n}$ and decoding functions $m^{\text{dec}}_{n}$ parameterized by neural networks, and define the following (volume preserving) invertible mappings

with $e^{\pi}_{0}=0$ and $x^{\pi}_{0}\sim q_{0}$. The step size parameter $\epsilon$ will be chosen to depend on the depth coefficient $N$, i.e. the number of layers of the joint transformation.

Below we prove ANF of the above form can universally transform $q(x)\delta_{0}(e)$ into $p(x)\delta_{0}(e)$. We make the following assumption on the family of $q$:

We assume the gradient of the convex function in Proposition (1) $\nabla\Phi$ is continuous, and that $f(e,t):=e^{\alpha_{t}-\gamma_{t}}e$ and $g(x,t):=-e^{\alpha_{t}+\beta_{t}+\gamma_{t}}\log\frac{q_{t}(x)}{p(x)}$ have a bounded second time derivative (on the trajectories $x_{t}$ and $e_{t}$ which are also functions of time), and are uniformly Lipschitz; that is,

for some $K\geq 0$, where we define the single-argument vector functions $f(t)=f(e_{t},t)$ and $g(t)=g(x_{t},t)$ as the time derivatives of the trajectories $(x_{t},e_{t})$.

We denote by ${\mathcal{Q}}$ the family of probability measures that satisfies this assumption.

Before we move on to approximation, we start with a lemma for bounding approximation error by solving recursion using the technique of generating functions.

If for any $N>0$, $\{d_{n}:0\leq n\leq N\}$ is a sequence of real numbers satisfying

We would like to bound the error $d_{n}$ explicitly. To do so, we first note that the sequence $\{d_{n}\}$ is no larger than $\{D_{n}\}$, which is recursively defined as

for $n\geq 0$, where for simplicity we let $C=c/N^{2}$.

Now to express $D_{n+1}$ explicitly, we use the method of generating function, following the recipe of Wilf (2005) (see Chapter 1 for a brief introduction). Define function $f$ to be a power series whose coefficients are $D_{n}$’s; that is, $f(x)=\sum_{n\geq 0}D_{n}x^{n}$. Multiply both sides of (13) by $x^{n}$ and summing over the indices of non-negative integers $n\geq 0$ give us

which can be decomposed into the partial fractions

where $a_{1}$ and $a_{2}$ are the roots of the quadratic function $x^{2}-(2+C)x+1$, which satisfy $a_{1}+a_{2}=2+C$ and $a_{1}a_{2}=1$.

For sufficiently small $x$, we can break (14) into the geometric series

This means for $n>0$, since $a_{1}a_{2}=1$, the coefficient of $f(x)$ can be expressed as

Now let $a_{1}$ be the larger root. Solving $x^{2}-(2+C)x+1$ yields

where $r:=\frac{C}{2}+\sqrt{\frac{C^{2}}{4}+C}$.

We show that the parenthesis in (15) can be controlled asymptotically (i.e. does not exceed certain constant for sufficiently large $N$), and that since $C$ diminishes, $D_{n}$ converges. First, since $r>0$, $a_{1}>1$ and

Second, since $(1+r)^{n}\leq e^{nr}$ for $n\geq 0$ and $r\geq-1$,

which converges to $\exp(\sqrt{c})$ as $N\rightarrow\infty$.

Finally, since $C\rightarrow 0$ as $N\rightarrow\infty$ and $d_{n}\leq D_{n}$, $d_{n}\rightarrow 0$ for all $n\leq N$ as $N\rightarrow\infty$. ∎

We are now ready to show the result of the pointwise approximation of the Hamiltonian ODE using ANFs with affine (more specifically, additive) coupling.

Fix $q_{0}\in{\mathcal{Q}}$ and $T>0$ and some compact subset ${\mathcal{X}}_{0}\subset{\mathcal{X}}$. We first consider all points $x_{0}$ in ${\mathcal{X}}_{0}$, and show that $(x^{\pi}_{n},e^{\pi}_{n})$ can be used to approximate $(x_{T},e_{T})$ uniformly well.

We consider a $N$-step joint transformation, and set $\epsilon=\frac{T}{2N}>0$. We start with approximating $e_{\epsilon}$ by $e^{\pi}_{1}$. Since $e^{\pi}_{0}$ is , by the universal approximation theorem (UAT) of neural networks (Cybenko, 1989), we can choose some $m^{\text{enc}}_{1}$ such that $||e_{\epsilon}-e^{\pi}_{1}||=||e_{\epsilon}-m^{\text{enc}}_{1}||\leq\epsilon^{2}$ for all $x_{0}\in{\mathcal{X}}_{0}$.

We proceed with an approximate leap-frog integration of the dynamic, using the neural encoders and decoders to approximate the time derivatives. Let ${\mathcal{E}}_{1}:=e^{\pi}_{1}({\mathcal{X}}_{0})$ where $e^{\pi}_{1}:=m^{\text{enc}}_{1}$, which is compact, since ${\mathcal{X}}_{0}$ is compact and $e^{\pi}_{1}$ is continuous wrt ${\mathcal{X}}_{0}$. Again, by the UAT, we can choose some $m^{\text{dec}}_{1}$ such that $||f(e,\epsilon)-m^{\text{dec}}_{1}(e)||<\epsilon^{2}$ for all $e\in{\mathcal{E}}_{1}$. Likewise, we let ${\mathcal{X}}_{1}:=x^{\pi}_{1}({\mathcal{X}}_{0})$ where $x^{\pi}_{1}:=(2\epsilon m^{\text{dec}}_{1}\circ e^{\pi}_{1}+Id)({\mathcal{X}}_{0})$ with $Id$ being the identity map, such that ${\mathcal{X}}_{1}$ is also compact since $x^{\pi}_{1}$ is continuous wrt ${\mathcal{X}}_{0}$, and choose $m^{\text{enc}}_{2}$ such that $||g(x,2\epsilon)-m^{\text{enc}}_{2}(x)||<\epsilon^{2}$ for all $x\in{\mathcal{X}}_{1}$.

Repeating the same construction for $m^{\text{dec}}_{n}$ and $m^{\text{enc}}_{n}$ for $n\leq N$, we have

with $m^{\text{dec}}_{n}$ and $m^{\text{enc}}_{n}$ chosen such that

$||f(e,2n\epsilon+\epsilon)-m^{\text{dec}}_{n+1}(e)||<\epsilon^{2}$ for all $e\in{\mathcal{E}}_{n+1}:=e^{\pi}_{n+1}({\mathcal{X}}_{0})$ where $e^{\pi}_{n+1}:=2\epsilon m^{\text{enc}}_{n+1}\circ x^{\pi}_{n}+e^{\pi}_{n}$ is a continuous map of ${\mathcal{X}}_{0}$; and

$||g(x,2n\epsilon)-m^{\text{enc}}_{n+1}(x)||<\epsilon^{2}$ for all $x\in{\mathcal{X}}_{n}:=x^{\pi}_{n}({\mathcal{X}}_{0})$ where $x^{\pi}_{n}:=2\epsilon m^{\text{dec}}_{n}\circ e^{\pi}_{n}+x^{\pi}_{n-1}$ is a continuous map of ${\mathcal{X}}_{0}$.

Such choices of $m^{\text{enc}}_{n}$ and $m^{\text{dec}}_{n}$ are possible since by construction ${\mathcal{X}}_{n-1}$ and ${\mathcal{E}}_{n}$ are compact.

Equations (16,17) are approximate midpoint methods as they use functions to approximate the time derivatives evaluated at midpoints of their counterparts. The exact midpoint method has a cubic error rate of $\frac{h^{3}}{24}f^{\prime\prime}(\xi)$, for some $\xi$ between the midpoint and the approximating point, where $h$ is the interval width of each iteration; see Section 5.4 of Epperson (2013). That is,

for some $\xi^{e}_{n+1}$ between the two steps.

The error on the RHS consists of two parts: (1) the first two terms constitute the propagated error from the previous steps and (2) the third term is a newly introduced truncation error due to the Taylor expansion.

Again the RHS can be decomposed into two error parts: (1) a midpoint deviation resulting from performing midpoint numerical integration which would not vanish even if the neural network is replaced with the true time derivative, and (2) an approximation error due to the inaccuracy of approximating the time derivative.

Letting $d^{x}_{n}=||x_{2n\epsilon}-x^{\pi}_{n}||$ and $d^{e}_{n}=||e_{2n\epsilon-\epsilon}-e^{\pi}_{n}||$, and applying the properties of the Assumption 1, we have

owing to the uniform error bound of the neural decoder $||f(e,2n\epsilon+\epsilon)-m^{\text{dec}}_{n+1}(e)||<\epsilon^{2}$ for all $e\in{\mathcal{E}}_{n+1}$ and the fact that $e^{\pi}_{n+1}(x_{0})\in{\mathcal{E}}_{n+1}$ since $x_{0}\in{\mathcal{X}}_{0}$.

The same can be done to obtain a bound on $d^{e}_{n+1}$ by subtracting (17) from (19), which yields

where $K^{\prime}=\max\{K,\frac{K}{3}+2\}$.

Summing $d^{x}_{1},...,d^{x}_{n}$ and subtracting $d^{x}_{1}+...+d^{x}_{n-1}$ from both sides yield

Note that $d^{x}_{0}=0$. Similarly, summing $d^{e}_{2},...,d^{e}_{n}$ and subtracting $d^{e}_{2}+...+d^{e}_{n-1}$ from both sides yield

To recursively express $d^{x}_{n}$ in terms of itself (except for $d^{e}_{1}$), we sum over the sequence $d^{e}_{1},...,d^{e}_{n}$ again

Since $n\leq N$, $\sum_{t=1}^{n}t\leq n^{2}$, $d^{e}_{1}\leq\epsilon^{2}$ and $\epsilon=\frac{T}{2N}$, the above can be rearranged and further bounded by

The same can be done for (24) to analyze $d^{e}_{n}$.

pointwise as $B\rightarrow\infty$. The same holds for the augmented variable $e$. ∎

The lemma below shows if one can approximate the solution of an ODE ($||y_{n}-x_{n}||\rightarrow 0$, i.e. $x_{n}$ and $y_{n}$ are asymptotically indistinguishable) and if the limit of the solution is a transport map ($x_{n}\overset{d}{\rightarrow}x_{\infty}$), then the approximation also forms a transport map ($y_{n}\overset{d}{\rightarrow}x_{\infty}$).

Let $x_{\infty}$, $(x_{n}:n\geq 0)$ and $(y_{n}:n\geq 0)$ be random variables. If $x_{n}\rightarrow x_{\infty}$ in distribution and if $||x_{n}-y_{n}||\rightarrow 0$ almost surely as $n\rightarrow\infty$, then $y_{n}\rightarrow x_{\infty}$ in distribution.

First, since $x_{n}\rightarrow x_{\infty}$ in distribution and since $\Lambda$ is bounded and continuous, by the Portmanteau Lemma the first term of the RHS converges to as $n\rightarrow\infty$. Second, since $y_{n}$ is almost surely asymptotically indistinguishable from $x_{n}$ (let $\Omega$ be the almost sure set), and since the Lipschitzness of $\Lambda$ implies uniform continuity, the following are true

For all $\epsilon>0$, there exists a $\delta>0$ such that $||x-y||\leq\delta$ implies $|\Lambda(x)-\Lambda(y)|\leq\epsilon$.

For any $\delta>0$, there exists a integer $N>0$ such that for all $n\geq N$, $||x_{n}-y_{n}||\leq\delta$ for all $\omega\in\Omega$.

These imply $||\Lambda(x_{n})-\Lambda(y_{n})||\rightarrow 0$ on $\Omega$. Then

We now are ready to prove Theorem 5, which we restate below. The main idea is to notice that ANFs can be made pointwise inseparable from the Hamiltonian ODE, which implies weak convergence since the Hamiltonian ODE converges in distribution. See 5

First, by Proposition 1, $x_{B}\rightarrow x_{\infty}$ in distribution as $B\rightarrow\infty$. Second, $x_{B}$ and $x^{\pi}_{N(B,1/B,[-B,B]^{d})}$ chosen from Proposition 2 are almost surely asymptotically indistinguishable. Thus, by Lemma 2, $x^{\pi}_{N(B,1/B,[-B,B]^{d})}$ converges in distribution to $x_{\infty}$. The same holds for the augmented variable $e$. Let $(x^{\pi}_{N})$ and $(e^{\pi}_{N})$ denote such sequences. By Theorem 2.7 of Van der Vaart (2000), $(x^{\pi}_{N},e^{\pi}_{N})\rightarrow(x_{\infty},e_{\infty})$ in distribution (as $e_{\infty}=0$ is a constant). ∎