Normalizing Flows for Probabilistic Modeling and Inference

Introduction

The search for well-specified probabilistic models—models that correctly describe the processes that produce data—is one of the enduring ideals of the statistical sciences. Yet, in only the simplest of settings are we able to achieve this goal. A central need in all of statistics and machine learning is then to develop the tools and theories that allow ever-richer probabilistic descriptions to be made, and consequently, that make it possible to develop better-specified models.

This paper reviews one tool we have to address this need: building probability distributions as normalizing flows. Normalizing flows operate by pushing a simple density through a series of transformations to produce a richer, potentially more multi-modal distribution—like a fluid flowing through a set of tubes. As we will see, repeated application of even simple transformations to a unimodal initial density leads to models of exquisite complexity. This flexibility means that flows are ripe for use in the key statistical tasks of modeling, inference, and simulation.

Normalizing flows are an increasingly active area of machine learning research. Yet there is an absence of a unifying lens with which to understand the latest advancements and their relationships to previous work. The thesis of Papamakarios (2019) and the survey by Kobyzev et al. (2020) have made steps in establishing this broader understanding. Our review complements these existing papers. In particular, our treatment of flows is more comprehensive than Papamakarios (2019)’s but shares some organizing principles. Kobyzev et al. (2020)’s article is commendable in its coverage and synthesis of the literature, discussing both finite and infinitesimal flows (as we do) and curating the latest results in density estimation. Our review is more tutorial in nature and provides in-depth discussion of several areas that Kobyzev et al. (2020) label as open problems (such as extensions to discrete variables and Riemannian manifolds).

Our exploration of normalizing flows attempts to illuminate enduring principles that will guide their construction and application for the foreseeable future. Specifically, our review begins by establishing the formal and conceptual structure of normalizing flows in Section 2. Flow construction is then discussed in detail, both for finite (Section 3) and infinitesimal (Section 4) variants. A more general perspective is then presented in Section 5, which in turn allows for extensions to structured domains and geometries. Lastly, we discuss commonly encountered applications in Section 6.

Normalizing Flows

We begin by outlining basic definitions and properties of normalizing flows. We establish the expressive power of flow-based models, explain how to use flows in practice, and provide some historical background. This section doesn’t assume prior familiarity with normalizing flows, and can serve as an introduction to the field.

The defining property of flow-based models is that the transformation $T$ must be invertible and both $T$ and $T^{-1}$ must be differentiable. Such transformations are known as diffeomorphisms and require that ${\mathbf{u}}$ be $D$-dimensional as well (Milnor and Weaver, 1997). Under these conditions, the density of ${\mathbf{x}}$ is well-defined and can be obtained by a change of variables (Rudin, 2006; Bogachev, 2007):

The Jacobian $J_{T}({\mathbf{u}})$ is the $D\times D$ matrix of all partial derivatives of $T$ given by:

An important property of invertible and differentiable transformations is that they are composable. Given two such transformations $T_{1}$ and $T_{2}$, their composition $T_{2}\circ T_{1}$ is also invertible and differentiable. Its inverse and Jacobian determinant are given by:

2 Expressive Power of Flow-Based Models

The Jacobian of $F$ is lower triangular since $\frac{\partial F_{i}}{\partial{\textnormal{x}}_{j}}=0$ for $i<j$. Hence, the Jacobian determinant of $F$ is equal to the product of its diagonal elements:

which implies ${\mathbf{z}}$ is distributed uniformly in the open unit cube $(0,1)^{D}$.

3 Using Flows for Modeling and Inference

Minimizing the above Monte Carlo approximation of the KL divergence is equivalent to fitting the flow-based model to the samples $\{{\mathbf{x}}_{n}\}_{n=1}^{N}$ by maximum likelihood estimation.

In practice, we often optimize the parameters ${\bm{\theta}}$ iteratively with stochastic gradient-based methods. We can obtain an unbiased estimate of the gradient of the KL divergence with respect to the parameters as follows:

3.2 Reverse KL Divergence

Alternatively, we may fit the flow-based model by minimizing the reverse KL divergence, which can be written as follows:

Similarly, we can estimate the gradient with respect to ${\bm{\psi}}$ by reparameterizing ${\mathbf{u}}$ as:

3.3 Relationship Between Forward and Reverse KL Divergence

3.4 Alternative Divergences

Learning the parameters of flow models is not restricted to the use of the KL divergence. Many alternative measures of difference between distributions are available. These alternatives are often grouped into two general families, the $f$-divergences that use density ratios to compare distributions, and the integral probability metrics (IPMs) that use differences for comparison:

For the $f$-divergences, the function $f$ is convex; when this function is $f(r)=r\log r$ we recover the KL divergence. For IPMs, the function $s$ can be chosen from a set of test statistics, or can be a witness function chosen adversarially.

4 Brief Historical Overview

Whitening transformations (Johnson, 1966; Friedman, 1987)—so-called as they transform data into white noise—are the clearest intellectual predecessor to the use of normalizing flows within machine learning. Chen and Gopinath (2000) were perhaps the first to use whitening as a density estimation technique rather than for feature pre-processing, calling the method Gaussianization. Tabak and Vanden-Eijnden (2010) approached Gaussianization from the perspective of diffusion processes, making connections to statistical mechanics—specifically, using Liouville’s equation to characterize the flow. In a follow-up paper, Tabak and Turner (2013) introduce what can be thought of as the modern conception of normalizing flows: introducing the very term normalizing flow and defining the flow generally as a composition of $K$ simple maps. As we will discuss in Section 3, this definition via composition is essential for enabling flows to be expressive while preserving computational and analytical tractability.

The idea of composition saw its recent emergence in machine learning starting with Rippel and Adams (2013), who were perhaps the first to recognize that parameterizing flows with deep neural networks could result in quite general and expressive distribution classes. Dinh et al. (2015) then introduced a scalable and computationally efficient architecture, demonstrating further improvements to image modeling and inference. Rezende and Mohamed (2015) used the idea and language from Tabak and Turner (2013) to apply normalizing flows in the setting of variational inference. Following this, as the papers that are reviewed here will show, normalizing flows now constitute a broad literature in their own right, with much work expanding the applicability and scalability of these initial efforts.

Because of the connection to change of variables, we can also connect the current use of normalizing flows in machine learning with its pre-history in many other settings. The change of measure has been studied intensely in the development of statistical mechanics—a famous example being the aforementioned Liouville’s theorem (1838). Copulas (Sklar, 1959; Elidan, 2013) can be viewed as rudimentary flows, where each dimension is transformed independently using the empirically-estimated marginal cumulative distribution function. Optimal transport and the Wasserstein metric (Villani, 2008) can also be formulated in terms of transformations of measures (‘transport of measures’)—also known as the Monge problem. In particular, triangular maps (a concept deeply related to autoregressive flows) can be shown to be a limiting solution to a class of Monge–Kantorovich problems (Carlier et al., 2010). This class of triangular maps itself has a long history, with Rosenblatt (1952) studying their properties for transforming multivariate distributions uniformly over the hypercube. Optimal transport could be a tutorial unto itself and therefore we mostly sidestep this framework, instead choosing to think in terms of the change of variables.

Constructing Flows Part I: Finite Compositions

Having described some high-level properties and uses of flows, we transition into describing, categorizing, and unifying the various ways to construct a flow. As discussed in Section 2.1, normalizing flows are composable; that is, we can construct a flow with transformation $T$ by composing a finite number of simple transformations $T_{k}$ as follows:

The idea is to use simple transformations as building blocks—each having a tractable inverse and Jacobian determinant—to define a complex transformation with more expressive power than any of its constituent components. Importantly, the flow’s forward and inverse evaluation and Jacobian-determinant computation can be localized to the sub-flows. As illustrated in Figure 2, assuming ${\mathbf{z}}_{0}={\mathbf{u}}$ and ${\mathbf{z}}_{K}={\mathbf{x}}$, the forward evaluation is:

and the Jacobian-determinant computation (in the log domain) is:

Increasing the ‘depth’ (i.e. number of composed sub-flows) of the transformation crucially results in only $\mathcal{O}\mathopen{}\left(K\right)\mathclose{}$ growth in the computational complexity—a pleasantly practical cost to pay for the increased expressivity.

In practice we implement either $T_{k}$ or $T_{k}^{-1}$ using a model (such as a neural network) with parameters ${\bm{\phi}}_{k}$, which we will denote as $f_{{\bm{\phi}}_{k}}$. That is, we can take the model $f_{{\bm{\phi}}_{k}}$ to implement either $T_{k}$, in which case it will take in ${\mathbf{z}}_{k-1}$ and output ${\mathbf{z}}_{k}$, or $T_{k}^{-1}$, in which case it will take in ${\mathbf{z}}_{k}$ and output ${\mathbf{z}}_{k-1}$. In either case, we must ensure that the model is invertible and has a tractable Jacobian determinant. In the rest of this section, we will describe several approaches for constructing $f_{{\bm{\phi}}_{k}}$ so that these requirements are satisfied. An overview of all the methods discussed in this section is shown in Table 1.

Ensuring that $f_{{\bm{\phi}}_{k}}$ is invertible and explicitly calculating its inverse are not synonymous. In many implementations, even though the inverse of $f_{{\bm{\phi}}_{k}}$ is guaranteed to exist, it can be expensive or even intractable to compute exactly. As discussed in Section 2, the forward transformation $T$ is used when sampling, and the inverse transformation $T^{-1}$ is used when evaluating densities. If the inverse of $f_{{\bm{\phi}}_{k}}$ is not efficient, either density evaluation or sampling will be inefficient or even intractable, depending on whether $f_{{\bm{\phi}}_{k}}$ implements $T_{k}$ or $T_{k}^{-1}$. Whether $f_{{\bm{\phi}}_{k}}$ should be designed to have an efficient inverse and whether it should be taken to implement $T_{k}$ or $T_{k}^{-1}$ are decisions that ought to be based on intended usage.

We should also clarify what we mean by ‘tractable Jacobian determinant’. We can always compute the Jacobian matrix of a differentiable function with $D$ inputs and $D$ outputs, using $D$ passes of either forward-mode or reverse-mode automatic differentiation. Then, we can explicitly calculate the determinant of that Jacobian. However, this computation has a time cost of $\mathcal{O}\mathopen{}\left(D^{3}\right)\mathclose{}$, which can be intractable for large $D$. For most applications of flow-based models, the Jacobian-determinant computation should be at most $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$. Hence, in the following sections, we will describe functional forms that allow the Jacobian determinant to be computed in linear time with respect to the input dimensionality.

To simplify notation from here on, we will drop the dependence of the model parameters on $k$ and denote the model by $f_{{\bm{\phi}}}$. Also, we will denote the model’s input by ${\mathbf{z}}$ and its output by ${\mathbf{z}}^{\prime}$, regardless of whether the model implements $T_{k}$ or $T_{k}^{-1}$.

where $\tau$ is termed the transformer and $c_{i}$ the $i$-th conditioner. This is illustrated in Figure 3(a). The transformer is a strictly monotonic function of ${\textnormal{z}}_{i}$ (and therefore invertible), is parameterized by ${\bm{h}}_{i}$, and specifies how the flow acts on ${\textnormal{z}}_{i}$ in order to output ${\textnormal{z}}_{i}^{\prime}$. The conditioner determines the parameters of the transformer, and in turn, can modify the transformer’s behavior. The conditioner does not need to be a bijection. Its one constraint is that the $i$-th conditioner can take as input only the variables with dimension indices less that $i$. The parameters ${\bm{\phi}}$ of $f_{\bm{\phi}}$ are typically the parameters of the conditioner (not shown above for notational simplicity), but sometimes the transformer has its own parameters too (in addition to ${\bm{h}}_{i}$).

It is easy to check that the above construction is invertible for any choice of $\tau$ and $c_{i}$ as long as the transformer is invertible. Given ${\mathbf{z}}^{\prime}$, we can compute ${\mathbf{z}}$ iteratively as follows:

This is illustrated in Figure 3(b). In the forward computation, each ${\bm{h}}_{i}$ and therefore each ${\textnormal{z}}_{i}^{\prime}$ can be computed independently in any order or in parallel. In the inverse computation however, all ${\mathbf{z}}_{<i}$ need to have been computed before ${\textnormal{z}}_{i}$, so that ${\mathbf{z}}_{<i}$ is available to the conditioner for computing ${\bm{h}}_{i}$.

It is also easy to show that the Jacobian of the above transformation is triangular, and thus the Jacobian determinant is tractable. Since each ${\textnormal{z}}_{i}^{\prime}$ doesn’t depend on ${\mathbf{z}}_{>i}$, the partial derivative of ${\textnormal{z}}_{i}^{\prime}$ with respect to ${\textnormal{z}}_{j}$ is zero whenever $j>i$. Hence, the Jacobian of $f_{{\bm{\phi}}}$ can be written in the following form:

The Jacobian is a lower-triangular matrix whose diagonal elements are the derivatives of the transformer for each of the $D$ elements of ${\mathbf{z}}$. Since the determinant of any triangular matrix is equal to the product of its diagonal elements, the log-absolute-determinant of $J_{f_{{\bm{\phi}}}}({\mathbf{z}})$ can be calculated in $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ time as follows:

The lower-triangular part of the Jacobian— denoted here by ${\mathbf{L}}({\mathbf{z}})$—is irrelevant. The derivatives of the transformer can be computed either analytically or via automatic differentiation, depending on the implementation.

Autoregressive flows are universal approximators (under the conditions discussed in Section 2.2) provided the transformer and the conditioner are flexible enough to represent any function arbitrarily well. This follows directly from the fact that the universal transformation from Section 2.2, which is based on the cumulative distribution functions of the conditionals, is indeed an autoregressive flow. Yet, this is just a statement of representational power and makes no guarantees about the flow’s behavior in practice.

An alternative, but mathematically equivalent, formulation of autoregressive flows is to have the conditioner $c_{i}$ take in ${\mathbf{z}}_{<i}^{\prime}$ instead of ${\mathbf{z}}_{<i}$. This is equivalent to swapping $\tau$ with $\tau^{-1}$ and ${\mathbf{z}}$ with ${\mathbf{z}}^{\prime}$ in the formulation presented above. Both formulations are common in the literature; here we use the convention that $c_{i}$ takes in ${\mathbf{z}}_{<i}$ without loss of generality. The computational differences between the two alternatives are discussed in more detail by e.g. Kingma et al. (2016); Papamakarios et al. (2017).

Implementing an autoregressive flow boils down to (a) implementing the transformer $\tau$ and (b) implementing the conditioner $c_{i}$. These are independent choices: any type of transformer can be paired up with any type of conditioner, yielding the various combinations that have appeared in the literature. In the following paragraphs, we will list a number of transformer implementations (Section 3.1.1) and an number of conditioner implementations (Section 3.1.2). We will discuss their pros and cons, and point out the choices that have been used by the specific models in the literature.

One of the simplest possible choices for the transformer—and one of the first used—is the class of affine functions:

Non-affine transformers can be constructed from simple components based on the observation that conic combinations as well as compositions of monotonic functions are also monotonic. Given monotonic functions $\tau_{1},\ldots,\tau_{K}$ of a real variable z, the following functions are also monotonic:

Conic combination: $\tau({\textnormal{z}})=\sum_{k=1}^{K}w_{k}\tau_{k}({\textnormal{z}})$, where $w_{k}>0$ for all $k$.

Composition: $\tau({\textnormal{z}})=\tau_{K}\circ\cdots\circ\tau_{1}({\textnormal{z}})$.

For example, a non-affine transformer can be constructed using a conic combination of monotonically increasing activation functions $\sigma(\cdot)$ (such as logistic sigmoid, tanh, leaky ReLU (Maas et al., 2013), and others):

provided $\alpha_{ik}>0$ and $w_{ik}>0$ for all $k\geq 1$. Clearly, the above construction corresponds to a monotonic single-layer perceptron. By repeatedly combining and composing monotonic activation functions, we can construct a multi-layer perceptron that is monotonic, provided that all its weights are strictly positive.

Transformers such as the above can represent any monotonic function arbitrarily well, which follows directly from the universal-approximation capabilities of multi-layer perceptrons (see e.g. Huang et al., 2018, for details). The derivatives of the transformer needed for the computation of the Jacobian determinant are in principle analytically obtainable, but more commonly they are computed via backpropagation. A drawback of combination-based transformers is that in general they cannot be inverted analytically, and can be inverted only iteratively e.g. using bisection search (Burden and Faires, 1989). Variants of combination-based transformers have been used in models such as NAF (Huang et al., 2018), block-NAF (De Cao et al., 2019), and Flow++ (Ho et al., 2019).

Another way to define a non-affine transformer is by recognizing that the integral of some positive function is a monotonically increasing function. For example, Wehenkel and Louppe (2019) define the transformer as:

where $g(\cdot;{\bm{\alpha}}_{i})$ can be any positive-valued neural network parameterized by ${\bm{\alpha}}_{i}$. Typically $g(\cdot;{\bm{\alpha}}_{i})$ will have its own parameters in addition to ${\bm{\alpha}}_{i}$. The derivative of the transformer required for the computation of the Jacobian determinant is simply equal to $g({\textnormal{z}}_{i};{\bm{\alpha}}_{i})$. This approach results in arbitrarily flexible transformers, but the integral lacks analytical tractability. One possibility is to resort to a numerical approximation.

An analytically tractable integration-based transformer can be obtained by taking the integrand $g(\cdot;{\bm{\alpha}}_{i})$ to be a positive polynomial of degree $2L$. The integral will be a polynomial of degree $2L+1$ in ${\textnormal{z}}_{i}$, and thus can be computed analytically. Since every positive polynomial of degree $2L$ can be written as a sum of $2$ (or more) squares of polynomials of degree $L$ (Marshall, 2008, Proposition 1.1.2), this fact can be exploited to define a sum-of-squares polynomial transformer (Jaini et al., 2019):

where $\alpha_{i}=\sum_{k=1}^{K}\alpha_{ik0}^{2}$. It can be shown that, for large enough $L$, the sum-of-squares polynomial transformer can approximate arbitrarily well any monotonically increasing function (Jaini et al., 2019, Theorem 3). Nonetheless, since only polynomials of degree up to $4$ can be solved analytically, the sum-of-squares polynomial transformer is not analytically invertible for $L\geq 2$, and can only be inverted iteratively using e.g. bisection search (Burden and Faires, 1989).

So far, we have discussed non-affine transformers that can be made arbitrarily flexible, but don’t have an analytic inverse. One way to overcome this limitation is by implementing the transformer as a monotonic spline, i.e. a piecewise function consisting of $K$ segments, where each segment is a simple function that is easy to invert. Specifically, given a set of $K+1$ input locations ${\textnormal{z}}_{i0},\ldots,{\textnormal{z}}_{iK}$, the transformer $\tau({\textnormal{z}}_{i};{\bm{h}}_{i})$ is taken to be a simple monotonic function (e.g. a low-degree polynomial) in each interval $[{\textnormal{z}}_{i(k-1)},{\textnormal{z}}_{ik}]$, under the constraint that the $K$ segments meet at the endpoints ${\textnormal{z}}_{i1},\ldots,{\textnormal{z}}_{i(K-1)}$. Outside the interval $[{\textnormal{z}}_{i0},{\textnormal{z}}_{iK}]$, the transformer can default to a simple function such as the identity. Typically, the parameters ${\bm{h}}_{i}$ of the transformer are the input locations ${\textnormal{z}}_{i0},\ldots,{\textnormal{z}}_{iK}$, the corresponding output locations ${\textnormal{z}}_{i0}^{\prime},\ldots,{\textnormal{z}}_{iK}^{\prime}$, and (depending on the type of spline) the derivatives (i.e. slopes) at ${\textnormal{z}}_{i0},\ldots,{\textnormal{z}}_{iK}$. See Figure 4 for an illustration.

Spline-based transformers are distinguished by the type of spline they use, i.e. by the functional form of the segments. The following options have been explored thus far, in order of increasing flexibility: linear and quadratic splines (Müller et al., 2019), cubic splines (Durkan et al., 2019a), linear-rational splines (Dolatabadi et al., 2020), and rational-quadratic splines (Durkan et al., 2019b). Spline-based transformers are as fast to invert as to evaluate, while maintaining exact analytical invertibility. Evaluating or inverting a spline-based transformer is done by first locating the right segment—which can be done in $\mathcal{O}\mathopen{}\left(\log K\right)\mathclose{}$ time using binary search—and then evaluating or inverting that segment, which is assumed to be analytically tractable. By increasing the number of segments $K$, a spline-based transformer can be made arbitrarily flexible.

1.2 Implementing the Conditioner

The conditioner $c_{i}({\mathbf{z}}_{<i})$ can be any function of ${\mathbf{z}}_{<i}$, meaning that each conditioner can, in principle, be implemented as an arbitrary model with input ${\mathbf{z}}_{<i}$ and output ${\bm{h}}_{i}$. However, a naïve implementation in which each $c_{i}({\mathbf{z}}_{<i})$ is a separate model would scale poorly with the dimensionality $D$, requiring $D$ model evaluations, each with a vector of average size $D/2$. This is in addition to the cost of storing and estimating the parameters of $D$ independent models. In fact, early work on flow precursors (Chen and Gopinath, 2000) dismissed the autoregressive approach as prohibitively expensive.

Nonetheless, this problem can be effectively addressed in practice by sharing parameters across the conditioners $c_{i}({\mathbf{z}}_{<i})$, or even by combining the conditioners into a single model. In the following paragraphs, we will discuss some practical implementations of the conditioner that allow it to scale to high dimensions.

One way to share parameters across conditioners is by implementing them jointly using a recurrent neural network (RNN). The $i$-th conditioner is implemented as:

The RNN processes ${\mathbf{z}}_{<D}=({\textnormal{z}}_{1},\ldots,{\textnormal{z}}_{D-1})$ one element at a time, and at each step it updates a fixed-size internal state ${\bm{s}}_{i}$ that summarizes the subsequence ${\mathbf{z}}_{<i}=({\textnormal{z}}_{1},\ldots,{\textnormal{z}}_{i-1})$. The network $c$ that computes ${\bm{h}}_{i}$ from ${\bm{s}}_{i}$ can be the same for each step. The initial state ${\bm{s}}_{1}$ can be fixed or it can be a learned parameter of the RNN. Any RNN architecture can be used, such as LSTM (Hochreiter and Schmidhuber, 1997) or GRU (Cho et al., 2014).

RNNs have been used extensively to share parameters across the conditional distributions of autoregressive models. Examples of RNN-based autoregressive models include distribution estimators (Larochelle and Murray, 2011; Uria et al., 2013, 2014), sequence models (Mikolov et al., 2010; Graves, 2013; Sutskever et al., 2014), and image/video models (Theis and Bethge, 2015; van den Oord et al., 2016b; Kalchbrenner et al., 2017). Section 3.1.3 discusses the relationship between autoregressive models and autoregressive flows in detail.

In the autoregressive-flows literature, RNN-based conditioners have been proposed by e.g. Oliva et al. (2018) and Kingma et al. (2016), but are relatively uncommon compared to alternatives. The main downside of RNN-based conditioners is that they turn an inherently parallel computation into a sequential one: the states ${\bm{s}}_{1},\ldots,{\bm{s}}_{D}$ must be computed sequentially, even though each ${\bm{h}}_{i}$ can in principle be computed independently and in parallel from ${\mathbf{z}}_{<i}$. Since this recurrent computation involves $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ sequential steps, it can be slow for high-dimensional data such as images or videos.

Another approach that shares parameters across conditioners but avoids the sequential computation of an RNN is based on masking. This approach uses a single, typically feedforward neural network that takes in ${\mathbf{z}}$ and outputs the entire sequence $({\bm{h}}_{1},\ldots,{\bm{h}}_{D})$ in one pass. The only requirement is that this network must obey the autoregressive structure of the conditioner: an output ${\bm{h}}_{i}$ cannot depend on inputs ${\mathbf{z}}_{\geq i}$.

To construct such a network, one takes an arbitrary neural network and removes connections until there is no path from input ${\textnormal{z}}_{i}$ to outputs $({\bm{h}}_{1},\ldots,{\bm{h}}_{i})$. A simple way to remove connections is by multiplying each weight matrix elementwise with a binary matrix of the same size. This has the effect of removing the connections corresponding to weights that are multiplied by zero, while leaving all other connections unmodified. These binary matrices can be thought of as ‘masking out’ connections, hence the term masking. The masked network will have the same architecture and size as the original network. In turn, it retains the computational properties of the original network, such as parallelism or ability to evaluate efficiently on a GPU.

A general procedure for constructing masks for multilayer perceptrons with arbitrarily many hidden layers or hidden units was proposed by Germain et al. (2015). The key idea is to assign a ‘degree’ between $1$ and $D$ to each input, hidden, and output unit, and mask-out the weights between subsequent layers such that no unit feeds into a unit with lower or equal degree. In convolutional networks, masking can be done by multiplying the filter with a binary matrix of the same size, which leads to a type of convolution often referred to as autoregressive or causal convolution (van den Oord et al., 2016c, a; Hoogeboom et al., 2019b). In architectures that use self-attention, masking can be done by zeroing out the softmax probabilities (Vaswani et al., 2017).

Masked autoregressive flows have two main advantages. First, they are efficient to evaluate. Given ${\mathbf{z}}$, the parameters $({\bm{h}}_{1},\ldots,{\bm{h}}_{D})$ can be obtained in one neural-network pass, and then each dimension of ${\mathbf{z}}^{\prime}$ can be computed in parallel via ${\textnormal{z}}_{i}^{\prime}=\tau({\textnormal{z}}_{i};{\bm{h}}_{i})$. Second, masked autoregressive flows are universal approximators. Given a large enough conditioner and a flexible enough transformer, they can represent any autoregressive transformation with monotonic transformers and thus transform between any two distributions (as discussed in Section 2.2).

On the other hand, the main disadvantage of masked autoregressive flows is that they are not as efficient to invert as to evaluate. This is because the parameters ${\bm{h}}_{i}$ that are needed to obtain ${\textnormal{z}}_{i}=\tau^{-1}({\textnormal{z}}_{i}^{\prime};{\bm{h}}_{i})$ cannot be computed until all $({\textnormal{z}}_{i},\ldots,{\textnormal{z}}_{i-1})$ have been obtained. Following this logic, we must first compute ${\bm{h}}_{1}$ by which we obtain ${\textnormal{z}}_{1}$, then compute ${\bm{h}}_{2}$ by which we obtain ${\textnormal{z}}_{2}$, and so on until ${\textnormal{z}}_{D}$ has been obtained. Using a masked conditioner $c$, the above procedure can be implemented in pseudocode as follows:

To see why this procedure is correct, observe that if ${\mathbf{z}}_{\leq i-1}$ is correct before the $i$-th iteration, then ${\bm{h}}_{i}$ will be computed correctly (due to the autoregressive structure of $c$) and thus ${\mathbf{z}}_{\leq i}$ will be correct before the $(i+1)$-th iteration. Since ${\mathbf{z}}_{\leq 0}=\varnothing$ is correct before the first iteration (in a degenerate sense, but still), by induction it follows that ${\mathbf{z}}_{\leq D}={\mathbf{z}}$ will be correct at the end of the loop. Even though the above procedure can invert the flow exactly (provided the transformer is easy to invert), it requires calling the conditioner $D$ times. This means that inverting a masked autoregressive flow using the above method is about $D$ times more expensive than evaluating the forward transformation. For high-dimensional data such as images or video, this can be prohibitively expensive.

An alternative way to invert the flow, proposed by Song et al. (2019), is to solve the equation ${\mathbf{z}}^{\prime}=f_{{\bm{\phi}}}({\mathbf{z}})$ approximately, by iterating the following Newton-style fixed-point update:

Despite the computational difficulties associated with inversion, masking remains one of the most popular techniques for implementing autoregressive flows. It is well suited to situations for which inverting the flow is not needed or the data dimensionality is not too large. Examples of flow-based models that use masking include IAF (Kingma et al., 2016), MAF (Papamakarios et al., 2017), NAF (Huang et al., 2018), block-NAF (De Cao et al., 2019), MintNet (Song et al., 2019) and MaCow (Ma et al., 2019). Masking can also be used and has been popular in implementing non-flow-based autoregressive models such as MADE (Germain et al., 2015), PixelCNN (van den Oord et al., 2016c; Salimans et al., 2017) and WaveNet (van den Oord et al., 2016a, 2018).

As we have seen, masked autoregressive flows have computational asymmetry that impacts their application and usability. Either sampling or density evaluation will be $D$ times slower than the other. If both of these operations are required to be fast, a different implementation of the conditioner is needed. One such implementation that is computationally symmetric, i.e. equally fast to evaluate or invert, is the coupling layer (Dinh et al., 2015, 2017). The idea is to choose an index $d$ (a common choice is $D/2$ rounded to an integer) and design the conditioner such that:

Parameters $({\bm{h}}_{1},\ldots,{\bm{h}}_{d})$ are constants, i.e. not a function of ${\mathbf{z}}$.

Parameters $({\bm{h}}_{d+1},\ldots,{\bm{h}}_{D})$ are functions of ${\mathbf{z}}_{\leq d}$ only, i.e. they don’t depend on ${\mathbf{z}}_{>d}$.

This can be easily implemented using an arbitrary function approximator $F$ (such as a neural network) as follows:

In other words, the coupling layer splits ${\mathbf{z}}$ into two parts such that ${\mathbf{z}}=[{\mathbf{z}}_{\leq d},{\mathbf{z}}_{>d}]$. The first part is transformed elementwise independently of other dimensions. The second part is transformed elementwise in a way that depends on the first part. We can also think of the coupling layer as implementing an aggressive masking strategy that allows only $({\bm{h}}_{d+1},\ldots,{\bm{h}}_{D})$ to depend only on ${\mathbf{z}}_{\leq d}$.

Common implementations of coupling layers fix the transformers $\tau(\cdot;{\bm{h}}_{1}),\ldots,\tau(\cdot;{\bm{h}}_{D})$ to the identity function. In this case, the transformation can be written as follows:

In turn, the inverse transformation is straightforward, given by:

These are illustrated in Figure 5. Like all autoregressive flows, the Jacobian of the transformation is lower triangular, but in addition it has the following special structure:

where ${\mathbf{I}}$ is the $d\times d$ identity matrix, $\mathbf{0}$ is the $d\times(D-d)$ zero matrix, ${\mathbf{A}}$ is a $(D-d)\times d$ full matrix, and ${\mathbf{D}}$ is a $(D-d)\times(D-d)$ diagonal matrix. The Jacobian determinant is simply the product of the diagonal elements of ${\mathbf{D}}$, which are equal to the derivatives of the transformers $\tau(\cdot;{\bm{h}}_{d+1}),\ldots,\tau(\cdot;{\bm{h}}_{D})$.

Coupling layers and fully autoregressive flows are two extremes on a spectrum of possible implementations. A coupling layer splits ${\mathbf{z}}$ into two parts and transforms the second part elementwise as a function of the first part, whereas a fully autoregressive flow splits the input into $D$ parts (each with one element in it) and transforms each part as a function of all previous parts. Clearly, there are intermediate choices: one can split the input into $K$ parts and transform the $k$-th part elementwise as a function of parts $1$ to $k-1$, with $K=2$ corresponding to a coupling layer and $K=D$ to a fully autoregressive flow. Using masking, inverting the transformation will be $\mathcal{O}\mathopen{}\left(K\right)\mathclose{}$ times more expensive than evaluating it, hence $K$ could be chosen based on the computational requirements of the task.

The efficiency of coupling layers comes at the cost of reduced expressive power. Unlike a recurrent or masked autoregressive flow, a single coupling layer can no longer represent any autoregressive transformation, regardless of how expressive the function $F$ is. As a result, an autoregressive flow with a single coupling layer is no longer a universal approximator. Nonetheless, the expressivity of the flow can be increased by composing multiple coupling layers. When composing coupling layers, the elements of ${\mathbf{z}}$ need to be permuted between layers so that all dimensions have a chance to be transformed as well as interact with one another. Previous work across various domains (see e.g. Kingma and Dhariwal, 2018; Prenger et al., 2019; Durkan et al., 2019b) has shown that composing coupling layers can indeed create flexible flows.

Theoretically, it is easy to show that a composition of $D$ coupling layers is indeed a universal approximator, as long as the index $d$ of the $i$-th coupling layer is equal to $i-1$. Observe that the $i$-th coupling layer can express any transformation of the form ${\textnormal{z}}^{\prime}_{i}=\tau({\textnormal{z}}_{i};c_{i}({\mathbf{z}}_{<i}))$, hence a composition of $D$ such layers will have transformed each dimension fully autoregressively. However, this construction involves $D$ sequential computations (one for each layer) in both the forward and inverse directions, so it doesn’t provide an improvement over recurrent or masked autoregressive flows. It is an open problem whether it’s possible to obtain a universal approximator by composing strictly fewer than $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ coupling layers.

Coupling layers are one of the most popular methods for implementing flow-based models because they allow both density evaluation and sampling to be fast. A flow based on coupling layers can be tractably fitted with maximum likelihood and then be sampled from efficiently. Thus coupling layers are often found in generative models of high-dimensional data such as images, audio and video. Examples of flow-based models with coupling layers include NICE (Dinh et al., 2015), Real NVP (Dinh et al., 2017), Glow (Kingma and Dhariwal, 2018), WaveGlow (Prenger et al., 2019), FloWaveNet (Kim et al., 2019) and Flow++ (Ho et al., 2019).

1.3 Relationship with Autoregressive Models

Alongside normalizing flows, another popular class of models for high-dimensional distributions is the class of autoregressive models. Autoregressive models have a long history, from the general framework of Bayesian networks (Pearl, 1988; Frey, 1998) to more recent neural-network-based implementations (Bengio and Bengio, 2000; Uria et al., 2016).

We then model each conditional by some parametric distribution with parameters ${\bm{h}}_{i}$:

is always distributed uniformly in $(0,1)^{D}$. The above expression has exactly the same form as the definition of an autoregressive flow in Equation 29, with ${\mathbf{z}}={\mathbf{x}}$ and ${\mathbf{z}}^{\prime}={\mathbf{u}}$. Therefore, an autoregressive model is in fact an autoregressive flow with a single autoregressive layer. Moreover, the layer’s transformers are the cumulative distribution functions of the conditionals of the autoregressive model, and the layer’s base distribution is a uniform in $(0,1)^{D}$. We can make the connection explicit by writing the density under the change of variables:

The term involving the uniform base density drops from the expression, leaving just the Jacobian determinant.Inouye and Ravikumar (2018) termed flows of this form—whereby the density is fully determined by the Jacobian determinant—density destructors. Following Equation 30, the inverse autoregressive flow that maps ${\mathbf{u}}$ to ${\mathbf{x}}$ is obtained by iterating the following for $i\in\left\{1,\ldots,D\right\}$:

The above corresponds exactly to sampling from the autoregressive model one element at a time, where at each step the corresponding conditional is sampled from using inverse transform sampling.

Yet the transformer is not necessarily limited to being the inverse CDF. We can make further connections between specific types of autoregressive models and the transformers discussed in Section 3.1.1. For example, consider an autoregressive model with Gaussian conditionals of the form:

The above conditional can be reparameterized as follows:

Hence, the entire autoregressive model can be reparameterized as an affine autoregressive flow as shown in Equation 33, where $\alpha_{i}=\sigma_{i}$, $\beta_{i}=\mu_{i}$, and the base distribution is a standard Gaussian (Kingma et al., 2016; Papamakarios et al., 2017). In the same way, we can relate other types of autoregressive models to non-affine transformers. For example, an autoregressive model whose conditionals are mixtures of Gaussians can be reparameterized as an autoregressive flow with combination-based transformers such as those in Equation 35. Similarly, an autoregressive model whose conditionals are histograms can be reparameterized as an autoregressive flow with transformers given by linear splines.

In consequence, we can think of autoregressive flows as subsuming and further extending autoregressive models for continuous variables. There are several benefits of viewing autoregressive models as flows. First, this view decouples the model architecture from the source of randomness, which gives us freedom in specifying the base distribution. Thus, we can enhance the flexibility of an autoregressive model by choosing a more flexible base distribution; for example, the base distribution can be another autoregressive model with its own learnable parameters. This provides a framework for composing autoregressive models, like layers in a flow (Papamakarios et al., 2017). Also, it allows us to compose autoregressive models with other types of flows, potentially non-autoregressive ones.

2 Linear Flows

As discussed in the previous section, autoregressive flows restrict an output variable ${\textnormal{z}}^{\prime}_{i}$ to depend only on inputs ${\mathbf{z}}_{\leq i}$, making the flow dependent on the order of the input variables. As we showed, in the limit of infinite capacity, this restriction doesn’t limit the flexibility of the flow-based model. However, in practice we don’t operate at infinite capacity. The order of the input variables will determine the set of distributions the model can represent. Moreover, the target transformation may be easy to learn for some input orderings and hard to learn for others. The problem is further exacerbated when using coupling layers since only part of the input variables is transformed.

To cope with these limitations in practice, it often helps to permute the input variables between successive autoregressive layers. For coupling layers it is in fact necessary: if we don’t permute the input variables between successive layers, part of the input will never be transformed. A permutation of the input variables is itself an easily invertible transformation, and its absolute Jacobian determinant is always $1$ (i.e. it is volume-preserving). Hence, permutations can seamlessly be composed with other invertible and differentiable transformations the usual way.

An approach that generalizes the idea of a permutation of input variables is that of a linear flow. A linear flow is essentially an invertible linear transformation of the form:

where ${\mathbf{W}}$ is a $D\times D$ invertible matrix that parameterizes the transformation. The Jacobian of the above transformation is simply ${\mathbf{W}}$, making the Jacobian determinant equal to $\det{\mathbf{W}}$. A permutation is a special case of a linear flow, where ${\mathbf{W}}$ is a permutation matrix (i.e. a binary matrix with exactly one entry of $1$ in each row and column and s everywhere else). Alternating invertible linear transformations with autoregressive/coupling layers is often used in practice (see e.g. Kingma and Dhariwal, 2018; Durkan et al., 2019b).

A straightforward implementation of a linear flow is to directly parameterize and learn the matrix ${\mathbf{W}}$. However, this approach can be problematic. First, ${\mathbf{W}}$ is not guaranteed to be invertible. Second, inverting the flow, which amounts to solving the linear system ${\mathbf{W}}{\mathbf{z}}={\mathbf{z}}^{\prime}$ for ${\mathbf{z}}$, takes time $\mathcal{O}\mathopen{}\left(D^{3}\right)\mathclose{}$ in general, which is prohibitively expensive for high-dimensional data. Third, computing $\det{\mathbf{W}}$ also takes time $\mathcal{O}\mathopen{}\left(D^{3}\right)\mathclose{}$ in general.

To address the above challenges, many approaches restrict ${\mathbf{W}}$ to a structured matrix, or a product of structured matrices. For example, if we restrict ${\mathbf{W}}$ to be triangular, we can guarantee its invertibility by e.g. making the diagonal elements positive. Moreover, inversion then costs $\mathcal{O}\mathopen{}\left(D^{2}\right)\mathclose{}$ (i.e. about the same as a matrix multiplication) and computing the determinant costs $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$. In Appendix B, we discuss in more detail this and a few more parameterizations that restrict the form of ${\mathbf{W}}$ in various ways.

3 Residual Flows

In this section, we consider a class of invertible transformations of the general form:

where $g_{{\bm{\phi}}}$ is a function that outputs a $D$-dimensional translation vector, parameterized by ${\bm{\phi}}$ (Figure 6). This structure bears a strong similarity to residual networks (He et al., 2016), and thus we use the term residual flow to refer to a normalizing flow composed of such transformations. Residual transformations are not always invertible, but can be made invertible if $g_{{\bm{\phi}}}$ is constrained appropriately. In what follows, we discuss two general approaches to designing invertible residual transformations: the first is based on contractive maps, and the second is based on the matrix determinant lemma.

In other words, a contractive map brings any two inputs closer together (as measured by $\delta$) by at least a factor $L$. It directly follows that $F$ is Lipschitz continuous with a Lipschitz constant equal to $L$. The Banach fixed-point theorem (Rudin, 1976, Theorem 9.23) states that any such contractive map has exactly one fixed point ${\mathbf{z}}_{*}=F({\mathbf{z}}_{*})$. Furthermore, this fixed point is the limit of any sequence $({\mathbf{z}}_{0},{\mathbf{z}}_{1},\ldots)$ that is formed by an arbitrary starting point ${\mathbf{z}}_{0}$ and repeated application of $F$, i.e. ${\mathbf{z}}_{k+1}=F({\mathbf{z}}_{k})$ for all $k\geq 0$.

Invertibility of the residual transformation ${\mathbf{z}}^{\prime}=f_{{\bm{\phi}}}({\mathbf{z}})={\mathbf{z}}+g_{\bm{\phi}}({\mathbf{z}})$ follows directly from $g_{\bm{\phi}}$ being contractive. Given ${\mathbf{z}}^{\prime}$, consider the map:

If $g_{\bm{\phi}}$ is contractive with Lipschitz constant $L$, then $F$ is also contractive with the same Lipschitz constant. Hence, from the Banach fixed-point theorem, there exists a unique ${\mathbf{z}}_{*}$ such that ${\mathbf{z}}_{*}={\mathbf{z}}^{\prime}-g_{\bm{\phi}}({\mathbf{z}}_{*})$. By rearranging, we see that ${\mathbf{z}}^{\prime}=f_{\bm{\phi}}({\mathbf{z}}_{*})$, and since ${\mathbf{z}}_{*}$ is unique, it follows that $f_{\bm{\phi}}$ is invertible.

In addition to a proof of the invertibility of $f_{\bm{\phi}}$, the above argument also gives us an algorithm for inversion. Starting from an arbitrary input ${\mathbf{z}}_{0}$ (a convenient choice is ${\mathbf{z}}_{0}={\mathbf{z}}^{\prime}$), we can iteratively apply $F$ as follows:

The Banach fixed-point theorem guarantees that the above procedure converges to ${\mathbf{z}}_{*}=f_{\bm{\phi}}^{-1}({\mathbf{z}}^{\prime})$ for any choice of starting point ${\mathbf{z}}_{0}$. Moreover, it can be shown that the rate of convergence (with respect to $\delta$) is exponential in the number of iterations $k$, and can be quantified as follows:

The smaller the Lipschitz constant is, the faster ${\mathbf{z}}_{k}$ converges to ${\mathbf{z}}_{*}$. We can think of $L$ as trading off flexibility for efficiency: as $L$ gets smaller, the fewer iterations it takes to approximately invert the flow, but the residual transformation becomes more constrained, i.e. less flexible. In the extreme case of $L=0$, the inversion procedure converges after one iteration, but the transformation reduces to adding a constant.

A challenge in building contractive residual flows is designing the function $g_{\bm{\phi}}$ to be contractive without impinging upon its flexibility. It is easy to see that the composition of $K$ Lipschitz-continuous functions $F_{1},\ldots,F_{K}$ is also Lipschitz continuous with a Lipschitz constant equal to $\prod_{k=1}^{K}L_{k}$, where $L_{k}$ is the Lipschitz constant of $F_{k}$. Hence, if $g_{\bm{\phi}}$ is a composition of neural-network layers (as is common in deep learning), it is sufficient to make each layer Lipschitz continuous with $L_{k}\leq 1$, with at least one layer having $L_{k}<1$, for the entire network to be contractive. Many elementwise nonlinearities used in deep learning—including the logistic sigmoid, hyperbolic tangent (tanh), and rectified linear (ReLU)—are in fact already Lipschitz continuous with a constant no greater than $1$. Furthermore, linear layers (including dense layers and convolutional layers) can be made contractive with respect to a norm by dividing them with a constant strictly greater than their induced operator norm. One such implementation was proposed by Behrmann et al. (2019): spectral normalization (Miyato et al., 2018) was used to make linear layers contractive with respect to the Euclidean norm.

One drawback of contractive residual flows is that there is no known general, efficient procedure for computing their Jacobian determinant. Rather, one would have to revert to automatic differentiation to obtain the Jacobian and an explicit determinant computation to obtain the Jacobian determinant, which costs $\mathcal{O}\mathopen{}\left(D^{3}\right)\mathclose{}$ as discussed earlier. Without an efficient way to compute the Jacobian determinant, exactly evaluating the density of the flow model is costly and potentially infeasible for high-dimensional data such as images.

Nonetheless, it is possible to obtain an unbiased estimate of the log absolute Jacobian determinant, and hence of the log density, which is enough to train the flow model e.g. with maximum likelihood using stochastic gradients. We begin by writing the log absolute Jacobian determinant as a power series:This power series is essentially the Maclaurin series $\log(1+x)=x-\frac{x^{2}}{2}+\frac{x^{3}}{3}-\ldots$ extended to matrices.

where $J^{k}_{g_{\bm{\phi}}}({\mathbf{z}})$ is the $k$-th power of the Jacobian of $g_{\bm{\phi}}$ evaluated at ${\mathbf{z}}$. The above series converges if $\left\|J_{g_{\bm{\phi}}}({\mathbf{z}})\right\|<1$ for some submultiplicative matrix norm $\left\|\cdot\right\|$, which in our case holds due to $g_{\bm{\phi}}$ being contractive. The trace of $J^{k}_{g_{\bm{\phi}}}({\mathbf{z}})$ can be efficiently estimated using the Hutchinson trace estimator (Hutchinson, 1990):

where ${\mathbf{v}}$ can be any $D$-dimensional random vector with zero mean and unit covariance. The Jacobian-vector product ${\mathbf{v}}^{\top}J^{k}_{g_{\bm{\phi}}}({\mathbf{z}})$ can then be computed with $k$ backpropagation passes. Finally, the infinite sum can be estimated by a finite sum of appropriately re-weighted terms using the Russian-roulette estimator (Chen et al., 2019).

Unlike autoregressive flows, which are based on constraining the Jacobian to be sparse, contractive residual flows have a dense Jacobian in general, which allows all input variables to affect all output variables. As a result, contractive residual flows can be very flexible and have demonstrated good results in practice. On the other hand, unlike the one-pass density evaluation and sampling offered by flows based on coupling layers, exact density evaluation is computationally expensive and sampling is done iteratively, which limits the applicability of contractive residual flows in certain tasks.

3.2 Residual Flows Based on the Matrix Determinant Lemma

Suppose ${\mathbf{A}}$ is an invertible matrix of size $D\times D$ and ${\mathbf{V}}$, ${\mathbf{W}}$ are matrices of size $D\times M$. The matrix determinant lemma states:

If the determinant and inverse of ${\mathbf{A}}$ are tractable and $M$ is less than $D$, the matrix determinant lemma can provide a computationally efficient way to compute the determinant of ${\mathbf{A}}+{\mathbf{V}}{\mathbf{W}}^{\top}$. For example, if ${\mathbf{A}}$ is diagonal, computing the left-hand side costs $\mathcal{O}\mathopen{}\left(D^{3}+D^{2}M\right)\mathclose{}$, whereas computing the right-hand side costs $\mathcal{O}\mathopen{}\left(M^{3}+DM^{2}\right)\mathclose{}$, which is preferable if $M<D$. In the context of flows, the matrix determinant lemma can be used to efficiently compute the Jacobian determinant. In this section, we will discuss examples of residual flows that are specifically designed such that application of the matrix determinant lemma leads to efficient Jacobian-determinant computation.

One early example is the planar flow (Rezende and Mohamed, 2015), where the function $g_{\bm{\phi}}$ is a one-layer neural network with a single hidden unit:

where $\sigma^{\prime}$ is the derivative of the activation function. The Jacobian has the form of a diagonal matrix plus a rank-$1$ update. Using the matrix determinant lemma, the Jacobian determinant can be computed in time $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ as follows:

In general, the planar flow is not invertible for all values of ${\mathbf{v}}$ and ${\mathbf{w}}$. However, assuming that $\sigma^{\prime}$ is positive everywhere and bounded from above (which is the case if $\sigma$ is the hyperbolic tangent, for example), a sufficient condition for invertibility is ${\mathbf{w}}^{\top}{\mathbf{v}}>-\frac{1}{\sup_{x}\sigma^{\prime}(x)}$.

Planar flows can be extended to $M$ hidden units, in which case they are known as Sylvester flows (van den Berg et al., 2018) and can be written as:

where ${\mathbf{S}}({\mathbf{z}})$ is an $M\times M$ diagonal matrix whose diagonal is equal to $\sigma^{\prime}({\mathbf{W}}^{\top}{\mathbf{z}}+{\mathbf{b}})$. Applying the matrix determinant lemma we get:

which can be computed in $\mathcal{O}\mathopen{}\left(M^{3}+DM^{2}\right)\mathclose{}$. To further reduce the computational cost, van den Berg et al. (2018) proposed the parameterization ${\mathbf{V}}={\mathbf{Q}}{\mathbf{U}}$ and ${\mathbf{W}}={\mathbf{Q}}{\mathbf{L}}$, where ${\mathbf{Q}}$ is a $D\times M$ matrix whose columns are an orthonormal set of vectors (this requires $M\leq D$), ${\mathbf{U}}$ is $M\times M$ upper triangular, and ${\mathbf{L}}$ is $M\times M$ lower triangular. Since ${\mathbf{Q}}^{\top}{\mathbf{Q}}={\mathbf{I}}$ and the product of upper-triangular matrices is also upper triangular, the Jacobian determinant becomes:

Similar to planar flows, Sylvester flows are not invertible for all values of their parameters. Assuming $\sigma^{\prime}$ is positive everywhere and bounded from above, a sufficient condition for invertibility is $L_{ii}U_{ii}>-\frac{1}{\sup_{x}\sigma^{\prime}(x)}$ for all $i\in\left\{1,\ldots,D\right\}$.

Radial flows (Tabak and Turner, 2013; Rezende and Mohamed, 2015) take the following form:

which is a diagonal matrix plus a rank-$1$ update. Applying the matrix determinant lemma and rearranging, we get the following expression for the Jacobian determinant, which can be computed in $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$:

The radial flow is not invertible for all values of $\beta$. A sufficient condition for invertibility is $\beta>-\alpha$.

In summary, planar, Sylvester and radial flows have Jacobian determinants that cost $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ to compute, and can be made invertible by suitably restricting their parameters. However, there is no analytical way of computing their inverse, which is why these flows have mostly been used to approximate posteriors for variational autoencoders. Moreover, each individual transformation is fairly simple, and it’s not clear how the flexibility of the flow can be increased other than by increasing the number of transformations.

4 Practical Considerations when Combining Transformations

Implementing a flow often amounts to composing as many transformations as computation and memory will allow. For instance, Kingma and Dhariwal (2018)’s Glow architecture employs as many as $320$ sub-transformations distributed across $40$ GPUs to achieve state-of-the-art image generation. Working with such deep flows introduces additional challenges of a practical nature. In this section, we summarize two techniques that, respectively, stabilize the optimization and ease the computational demands of deep flows.

Like with deep neural networks trained with gradient-based methods, normalizing the intermediate representations ${\mathbf{z}}_{k}$ is crucial for maintaining stable gradients throughout the flow. Batch normalization or batch norm (Ioffe and Szegedy, 2015) has been widely demonstrated to be effective in stabilizing and improving neural-network training, thus making it attractive for use in deep flows as well. Viewing the batch statistics as fixed, batch norm is essentially a composition of two affine transformations. The first has scale and translation parameters set by the batch statistics, and the second has free parameters ${\bm{\alpha}}$ (scale) and ${\bm{\beta}}$ (translation):

Moreover, batch norm has an easy-to-compute Jacobian determinant due to it acting elementwise (and thus having a diagonal Jacobian):

The above formulas assume that batch statistics are fixed, which is true for a trained model. During training, however, batch statistics are not fixed, but are functions of all examples in the batch. This makes batch norm not invertible, unless the batch statistics have been cached during a forward pass. Also, the Jacobian determinant as written above makes little sense mathematically, since batch norm is now a function of the whole batch. Yet, using this Jacobian-determinant formula as an approximation often suffices for training, at least if the batch is large enough (Dinh et al., 2017; Papamakarios et al., 2017).

Glow employs a variant termed activation normalization or act norm (Kingma and Dhariwal, 2018) that doesn’t use batch statistics $\hat{{\bm{\mu}}}$ and $\hat{{\bm{\sigma}}}$. Instead, before training begins, a batch is passed through the flow, and ${\bm{\alpha}}$ and ${\bm{\beta}}$ are set such that the transformed batch has zero mean and unit variance. After this data-dependent initialization, ${\bm{\alpha}}$ and ${\bm{\beta}}$ are optimized as model parameters. Act norm is preferable when training with small mini-batches since batch norm’s statistics become noisy and can destabilize training.

As mentioned in Section 2.1, ${\mathbf{x}}$ and ${\mathbf{u}}$ must have the same dimensionality and every sub-transformation $T_{k}$ must preserve dimensionality. This means that evaluating $T$ incurs an increasing computational cost as dimensionality grows. This constraint is at direct odds with our desire to use as many steps in the flow as possible. Dinh et al. (2017) proposed side-stepping this issue by way of a multi-scale architecture. At regular intervals in the steps of the flow when going from ${\mathbf{x}}$ to ${\mathbf{u}}$, some number of sub-dimensions of ${\mathbf{z}}_{k}$ are clamped and no additional transformation is applied. One can think of this as implementing a skip-connection, mapping those dimensions directly to the corresponding dimensions in the final representation ${\mathbf{u}}$: $({\textnormal{u}}_{j},\ldots,{\textnormal{u}}_{i})=({\textnormal{z}}_{k,j},\ldots,{\textnormal{z}}_{k,i})$ where $k$ is the step at which the clamping is applied. All $K$ steps are applied to only a small subset of dimensions, which is less costly than applying all steps to all dimensions. Dinh et al. (2017) also argue that these skip-connections help with optimization, distributing the objective throughout the full depth of the flow.

Besides having this practical benefit, multi-scale architectures are a natural modeling choice for granular data types such as pixels (Dinh et al., 2017; Kingma and Dhariwal, 2018) and waveforms (Prenger et al., 2019; Kim et al., 2019). The macro-structures that we often care about—such as shapes and textures, in the case of images—typically do not need all $D$ dimensions to be described. Dinh et al. (2017) showed that multi-scale architectures do indeed encode more global, semantic information in the dimensions that undergo all transformations. On the other hand, dimensions that are factored out earlier in the flow represent lower-level information; see Dinh et al. (2017)’s Appendix D for demonstrations.

Constructing Flows Part II: Continuous-Time Transformations

In the preceding section, we considered constructing flows by parameterizing a one-step transformation ${\mathbf{z}}_{k}=T_{k}({\mathbf{z}}_{k-1})$, several of which are then composed to create a flow of $K$ discrete steps. An alternative strategy is to construct flows in continuous time by parameterizing the flow’s infinitesimal dynamics, and then integrating to find the corresponding transformation. In other words, we construct the flow by defining an ordinary differential equation (ODE) that describes the flow’s evolution in time. We call these ‘continuous-time’ flows as they evolve according to a real-valued scalar variable analogous to the number of steps. We call this scalar ‘time’ as it determines how long the dynamics are run. In this section, we will describe this class of continuous-time flows and summarize numerical tools necessary for their implementation.

Let ${\mathbf{z}}_{t}$ denote the flow’s state at time $t$ (or ‘step’ $t$, thinking in the discrete setting). Time $t$ is assumed to run continuously from $t_{0}$ to $t_{1}$, such that ${\mathbf{z}}_{t_{0}}={\mathbf{u}}$ and ${\mathbf{z}}_{t_{1}}={\mathbf{x}}$. A continuous-time flow is constructed by parameterizing the time derivative of ${\mathbf{z}}_{t}$ with a function $g_{\bm{\phi}}$ with parameters ${\bm{\phi}}$, yielding the following ordinary differential equation (ODE):

The function $g_{{\bm{\phi}}}$ takes as inputs both the time $t$ and the flow’s state ${\mathbf{z}}_{t}$, and outputs the time derivative of ${\mathbf{z}}_{t}$ at time $t$. The only requirements for $g_{{\bm{\phi}}}$ are that it be uniformly Lipschitz continuous in ${\mathbf{z}}_{t}$ (meaning that there is a single Lipschitz constant that works for all $t$) and continuous in $t$ (Chen et al., 2018). From Picard’s existence theorem, it follows that satisfying these requirements ensures that the above ODE has a unique solution (Coddington and Levinson, 1955). Many neural-network layers meet these requirements (Gouk et al., 2018), and unlike the architectures described in Section 3 that require careful structural assumptions to ensure invertibility and tractability of their Jacobian determinant, $g_{{\bm{\phi}}}$ has no such requirements.

To compute the transformation ${\mathbf{x}}=T({\mathbf{u}})$, we need to run the dynamics forward in time by integrating:

where in the right-most expression we used the fact that switching the limits of integration is equivalent to negating the integral. We write the inverse in this last form to show that, unlike many flows comprised of discrete compositions (Section 3), continuous-time flows have the same computational complexity in each direction. In consequence, choosing which direction is the forward and which is the inverse is not a crucial implementation choice as it is for e.g. autoregressive flows.

The change in log density for continuous-time flows can be characterized directly as (Chen et al., 2018):

where ${\mathbf{v}}$ can be any $D$-dimensional random vector with zero mean and unit covariance. The Jacobian-vector product ${\mathbf{v}}^{\top}J_{g_{\bm{\phi}}(t,\cdot)}({\mathbf{z}}_{t})$ can be computed in a single backpropagation pass, which makes the Hutchinson trace estimator about $D$ times more efficient than calculating the trace exactly. Chen and Duvenaud (2019) propose an alternative solution in which the architecture of $g_{\bm{\phi}}$ is carefully constrained so that the exact Jacobian trace can be computed in a single backpropagation pass.

By integrating the derivative of $\log p({\mathbf{z}}_{t})$ over time, we obtain an expression for the log density of ${\mathbf{x}}$ under the continuous-time flow:

Evaluating the forward transform and the log density can be done simultaneously by computing the following combined integral:

Computation is not analytically feasible for a general $g_{{\bm{\phi}}}$. In practice, a numerical integrator is used, as we discuss next.

2 Solving and Optimizing Continuous-Time Flows

Due to continuous-time flows being defined by an ODE, the vast literature on numerical ODE solvers and the corresponding software can be leveraged to implement these flows. See Süli (2010) for an accessible introduction to numerical methods for ODEs. While there are numerous numerical methods that could be employed, below we briefly describe Euler’s method and the adjoint method.

Perhaps the simplest numerical technique one can apply is Euler’s method. The idea is to first discretize the ODE using a small step-size $\epsilon>0$ as follows:

with the approximation becoming exact as $\epsilon\rightarrow 0$. The ODE can then be (approximately) solved by iterating the above computation starting from ${\mathbf{z}}_{t_{0}}={\mathbf{u}}$ until obtaining ${\mathbf{z}}_{t_{1}}={\mathbf{x}}$. This way, the parameters ${\bm{\phi}}$ can be optimized with gradients computed via backpropagation through the ODE solver. It is relatively straightforward to use other discrete solvers such as any in the Runge–Kutta family. The discretized forward solution would be backpropagated through just as with Euler’s method.

Interestingly, the Euler approximation implements the continuous-time flow as a discrete-time residual flow of the class described in Equation 55. Having assumed that $g_{\bm{\phi}}(t,\cdot)$ is uniformly Lipschitz continuous with a Lipschitz constant $L$ independent of $t$, it immediately follows that $\epsilon\,g_{\bm{\phi}}(t,\cdot)$ is contractive for any $\epsilon<1/L$. Hence, for small enough $\epsilon$ we can think of the above Euler discretization as a particular instance of a contractive residual flow (Section 3.3.1).

Approximating continuous-time flows using discrete-time residual flows gives us insight on Equation 78’s description of the time evolution of $\log p({\mathbf{z}}_{t})$. Using the Taylor-series expansion of Equation 60, we can write the log absolute Jacobian determinant of $f_{\bm{\phi}}$ as follows:

Substituting the above into the change-of-variables formula and rearranging we get:

from which we directly obtain Equation 78 by letting $\epsilon\rightarrow 0$.

2.2 The Adjoint Method

Chen et al. (2018) proposed an elegant alternative to the discrete, fixed-step methods mentioned above. For a general optimization target $\mathcal{L}({\mathbf{x}};{\bm{\phi}})$ (such as log likelihood), they show that the gradient $\partial\mathcal{L}/\partial{\mathbf{z}}_{t}$ with respect to the flow’s intermediate state ${\mathbf{z}}_{t}$ can be characterized by the following ODE:

a result known widely as the adjoint sensitivity method (Pontryagin, 1962). In neural-network terminology, the adjoint method can be thought of as the continuous-time analog of backpropagation. The gradient with respect to ${\bm{\phi}}$ can be computed by:

Optimization of ${\bm{\phi}}$ can then be done with stochastic gradients.

Formulating gradient computation as a separate ODE means that both forward evaluation and gradient computation can be treated by black-box ODE solvers, without the need to backpropagate through the solver’s computational graph. This results in significant practical benefits since backpropagating through a solver is costly both in terms of computation and memory requirements. Another benefit is that a more sophisticated solver can allocate instance-dependent computation based on a user-specified tolerance, a common hyperparameter in most off-the-shelf solvers. At test time, this tolerance level can be tuned based on runtime or other constraints on computation.

Generalizations

So far, we have discussed normalizing flows as invertible transformations in Euclidean space. In this section, we go beyond the standard definition and explore more general classes of probability transformations. We show how more general types of flow can be derived from a unifying framework, point to specific implementations that extend the standard definition, and explore the frontiers of normalizing-flow research.

Since $T$ is a diffeomorphism, we can express the LHS integral via the change of variable ${\mathbf{u}}=T^{-1}({\mathbf{x}})$ as follows (Rudin, 2006):

Since the above must be true for any $\gamma\subseteq\mathcal{X}$, it follows that:

which is the familiar formula for the density of a flow-based model.

Using the conservation of probability measure as a general framework, we study different transformations $T$, sets $\mathcal{U}$ and $\mathcal{X}$, and integration measures $d\mu({\mathbf{u}})$ and $d\nu({\mathbf{x}})$. In the following sections, we explore further possibilities and potential implementations.

2 Piecewise-Invertible Transformations and Mixtures of Flows

We can extend the transformation $T$ to be many-to-one, whereby multiple ${\mathbf{u}}$’s are mapped to the same ${\mathbf{x}}$. One possibility is for $T$ to be piecewise invertible, in which case we can partition $\mathcal{U}$ into a countable collection of non-overlapping subsets $\left\{\mathcal{U}_{i}\right\}_{i\in\mathcal{I}}$ such that the restriction of $T$ to $\mathcal{U}_{i}$ is an invertible transformation $T_{i}:\mathcal{U}_{i}\rightarrow\mathcal{X}$ (Figure 8(a)). Then, the conservation of probability measure can be written as:

where each $\omega_{i}\subseteq\mathcal{U}_{i}$ is the image of $\gamma$ under $T_{i}^{-1}$. Using the change of variables ${\mathbf{x}}=T_{i}({\mathbf{u}})$ for each corresponding integral on the LHS, we obtain:

Since the above must be true for all $\gamma\subseteq\mathcal{X}$, it follows that:

where $\gamma_{i}\subseteq\mathcal{X}_{i}$ is the image of $\omega$ under $T_{i}$. Using the change of variables ${\mathbf{x}}=T_{i}({\mathbf{u}})$ for each corresponding integral on the LHS, we obtain:

The above must be true for all $\xi\subseteq\mathcal{I}$ and $\gamma_{i}\subseteq\mathcal{X}$, and by definition $T_{i}^{-1}({\mathbf{x}})=R({\mathbf{x}})$ for all ${\mathbf{x}}\in\mathcal{X}_{i}$, therefore:

3 Flows for Discrete Random Variables

Using the general probability-transformation formula, we can extend normalizing flows to discrete sets $\mathcal{U}$ and $\mathcal{X}$. Taking the integration measures $d\mu({\mathbf{u}})$ and $d\nu({\mathbf{x}})$ to be the counting measure, we can write the conservation of measure as follows:

Letting $\gamma=\left\{{\mathbf{x}}\right\}$ for an arbitrary ${\mathbf{x}}$, we obtain:

We will refer to the above type of normalizing flow as a discrete flow. Unlike standard flows, discrete flows don’t involve a Jacobian term in their density calculation.

A possible way to overcome this limitation, pointed out by van den Berg et al. (2020), is to embed $\mathcal{U}$ and $\mathcal{X}$ into extended spaces $\mathcal{U}^{\prime}$ and $\mathcal{X}^{\prime}$, such that the base distribution factorizes in the extended space. In the above example, we could take $\mathcal{U}^{\prime}=\mathcal{X}^{\prime}=\left\{0,1,2,3\right\}^{2}$, so that the probability table becomes:

Then, a discrete flow on the extended space can rearrange the probability table into:

which is rank $1$ and thus can be factorized.

4 Flows on Riemannian Manifolds

An alternative approach was proposed by Falorsi et al. (2019) for the special case where the manifold is also a group, in which case it is known as a Lie group. A Lie group can be reparameterized with respect to its Lie algebra (i.e. the tangent space at the identity element) using the exponential map, which maps the Lie algebra to the Lie group. Falorsi et al. (2019) show that this parameterization can be done analytically in certain cases; in general however, computing the exponential map has a cubic cost with respect to the dimensionality $D$.

5 Bypassing Topological Constraints

For another example, if the target density is comprised of disconnected modes, the base density must have the same number of disconnected modes. This is true for all flows, not just continuous-time ones. If the base density has fewer modes than the target (which will likely be the case in practice), the flow will be forced to assign a non-zero amount of probability mass to the ‘empty’ space between the disconnected modes. To alleviate this issue, Cornish et al. (2019) propose a latent-variable flow defined as follows:

6 Symmetric Densities and Equivariant Flows

In many real-world applications, we have domain knowledge of the symmetries of the target density. For example, if we want to model a physical system composed of many interacting particles, the physical interactions between particles (e.g. gravitational or electrostatic forces) are invariant to translations. Another example is modeling the configuration of a molecular structure, where there may be rotational symmetries in the molecule. Such symmetries result in a probability density (e.g. of molecular configurations) that is invariant to specific transformations (e.g. rotations).

In such situations, it is desirable to build the known symmetries directly into the model, not only so that the model has the right properties, but also so that estimating the model is more data efficient. However, in general it is not trivial to build expressive probability densities that are also invariant to a prescribed set of symmetries. Flows are a good candidate class of models to combine with such domain knowledge because we have a lot of control of how they deform an initial density. In this section, we explore ways to build flow-based models that respect a prescribed set of symmetries.

The set of all symmetries of a target density is closed under composition, is associative, has an identity element, and each element has an inverse—therefore forming a group $G$.

These observations lead to the result shown in Lemma 1, which provides a general mechanism for constructing flow-based models whose density is invariant with respect to a prescribed symmetry group $G$.

Proof The flow density evaluated at $\mathbf{R}_{g}{\mathbf{x}}$ for some $g\in G$ is equal to:

From the equivariance of $T$ and hence of $T^{-1}$, we have that $T^{-1}\mathopen{}\left(\mathbf{R}_{g}{\mathbf{x}}\right)\mathclose{}=\mathbf{R}_{g}T^{-1}({\mathbf{x}})$. Taking the Jacobian of both sides, we obtain:

Proof From the invariance of $f$ we have that $f(\mathbf{R}_{g}{\mathbf{u}})=f({\mathbf{u}})$. Taking the gradient on both sides, we obtain:

Applications

Normalizing flows have two primitive operations: density calculation and sampling. In turn, flows are effective in any application requiring a probabilistic model with either of those capabilities. In this section, we summarize applications to probabilistic modeling, inference, supervised learning, and reinforcement learning.

Often, the data $\left\{{\mathbf{x}}_{n}\right\}_{n=1}^{N}$ are discrete; for example, they could be images with pixel values in $\left\{0,1,\ldots,255\right\}$. Flow-based models are defined over continuous random variables (with the exception of discrete flows, Section 5.3), so they are not directly applicable to discrete data. To use flows with discrete data, we often dequantize $\left\{{\mathbf{x}}_{n}\right\}_{n=1}^{N}$ by adding continuous noise. The noise distribution can be fixed (e.g. uniform in $$ for the image example above), or learned, as for example in variational dequantization (Ho et al., 2019).

The first task is primarily quantitative: we can use the model to estimate densities, expectations, marginals, or other quantities of interest on never-before-seen data. Early work (Chen and Gopinath, 2000; Tabak and Turner, 2013) considered only synthetic low-dimensional cases, showing that normalizing flows could indeed represent skewed, multi-modal densities as well as kernel density estimators could. It was Laparra et al. (2011) that first applied Gaussianization to real data, using the density function to perform one-class classification to detect urban areas in satellite images. Rippel and Adams (2013) next showed that their deep flow model’s density could detect rotations and corruptions of images. Papamakarios et al. (2017) performed a systematic comparison of their masked autoregressive flow on unconditional and conditional density estimation tasks, showing that the composition enabled by their framework allows for better density estimation than other variants (namely MADE and Real NVP). Grathwohl et al. (2019) performed similar experiments to validate the effectiveness of continuous-time flows.

1.2 Generation

Image generation has been given serious effort since the earliest work on flows. Laparra et al. (2011), in the same work mentioned above, used Gaussianization to generate gray-scale images of faces. Rippel and Adams (2013) also demonstrated early success in generative performance as MNIST samples from their model looked rather compelling. Dinh et al. (2015), through the use of their coupling parameterization, showed further improvements including density estimation competitive with other high-capacity models (such as deep mixtures of factor analysers) and respectable generation of SVHN digits. In follow-up work, Dinh et al. (2017) increased the capacity of their model by including scale transformations (instead of just translations), being the first to demonstrate that flows could produce sharp, visually compelling full-color images. Specifically, Dinh et al. (2017) showed compelling samples from models trained on CelebA, ImageNet ($64\times 64$), CIFAR-10, and LSUN. Kingma and Dhariwal (2018), using a similar model but with additional convolutional layers, further improved upon Dinh et al. (2017)’s results in density estimation and generation of high-dimensional images. Continuous (Grathwohl et al., 2019) and residual (Behrmann et al., 2019; Chen et al., 2019) flows have been demonstrated to produce sharp, high-dimensional images as well. Finally, Kumar et al. (2019) propose a normalizing flow for modeling video data, adapting the Glow architecture (Kingma and Dhariwal, 2018) to synthesize raw RGB frames.

The autoregressive model WaveNet (van den Oord et al., 2016a) demonstrated impressive performance in audio synthesis. While WaveNet is not a normalizing flow, in follow-up work van den Oord et al. (2018) defined a proper flow for audio synthesis by distilling WaveNet into an inverse autoregressive flow so as to make test-time sampling more efficient. Prenger et al. (2019) and Kim et al. (2019) have since formulated WaveNet variants built from coupling layers to enable fast likelihood and sampling, in turn obviating the need for van den Oord et al. (2018)’s post-training distillation step.

The most direct way to apply normalizing flows to text data is to define a discrete flow over characters or a vocabulary. Tran et al. (2019) take this approach, showing performance in character-level language modeling competitive to RNNs while having superior generation runtime. An alternative approach that has found wider use is to define a latent variable model with a discrete likelihood but a continuous latent space. A normalizing flow can then be defined on the latent space as usual. Ziegler and Rush (2019) use such an approach for character-level language modeling. Zhou et al. (2019), He et al. (2018), and Jin et al. (2019) define normalizing flows on the continuous space of word embeddings as a subcomponent of models for translation, syntactic structure, and parsing respectively.

Extending flows to operate on other structured objects is a burgeoning area of work. So far, flows have been applied to graphs (Deng et al., 2019), molecules (Madhawa et al., 2019; Honda et al., 2019), point clouds (Yang et al., 2019), and part models for motion synthesis (Henter et al., 2019).

2 Inference

In the previous section our focus was on modeling data and recovering its underlying distribution. We now turn to inference: estimating unknown quantities within a model. The most common setting is the computation of high-dimensional, analytically intractable integrals of the form:

Bayesian inference usually runs into such an obstacle when computing the posterior’s normalizing constant or when computing expectations under the posterior. Below we summarize the use of flows for sampling, variational inference, and likelihood-free inference.

Importance sampling (IS) computes intractable integrals by converting them to an expectation under an auxiliary distribution $q({\bm{\eta}})$:

where $q({\bm{\eta}})$ is a user-specified density function and $\hat{{\bm{\eta}}}_{s}$ is a sample from $q({\bm{\eta}})$. Clearly, IS requires both sampling and density evaluation. Since both operations are tractable for many flows, they make for an attractive model from which to construct the proposal.

The related technique of rejection sampling (RS) aims to draw samples from $p({\bm{\eta}})=\pi({\bm{\eta}})/Z$, where $\pi({\bm{\eta}})$ is again an unnormalized density. Both density evaluation and sampling are required from the proposal in RS, again making normalizing flows well-suited. Bauer and Mnih (2019) use the Real NVP architecture (Dinh et al., 2017) to parameterize a proposal distribution for RS since the coupling layers allow for fast density evaluation and sampling.

2.2 Markov Chain Monte Carlo

The application of flows in Markov chain Monte Carlo (MCMC) precedes the appearance of flows in deep learning by at least a few decades. One prominent example is Hamiltonian Monte Carlo (HMC), also known as Hybrid Monte Carlo (Duane et al., 1987; Neal, 2010). HMC operates on the ‘phase space’ $({\bm{\eta}},{\mathbf{v}})$, where ${\bm{\eta}}$ are the variables of interest and ${\mathbf{v}}$ are additional ‘momentum’ variables of the same dimensionality as ${\bm{\eta}}$. HMC generates samples from a joint distribution $p({\bm{\eta}},{\mathbf{v}})$ constructed so that its marginal over ${\bm{\eta}}$ is the distribution of interest. Central to HMC is the Hamiltonian defined by $H({\bm{\eta}},{\mathbf{v}})=-\log p({\bm{\eta}},{\mathbf{v}})$. Given a state $({\bm{\eta}},{\mathbf{v}})$, HMC proposes a new state $({\bm{\eta}}^{\prime},{\mathbf{v}}^{\prime})=T({\bm{\eta}},{\mathbf{v}})$ deterministically, where $T$ is a Hamiltonian flow followed by negation of the momentum variables. The proposed state is then accepted/rejected using the usual Metropolis–Hastings step. The Hamiltonian flow is the continuous-time flow generated by the following ODE:

This flow is volume-preserving, meaning that its absolute Jacobian determinant is $1$ everywhere, which, in combination with the negation of the momentum variables, ensures that the proposal is symmetric and thus cancels in the Metropolis–Hastings ratio.

It is also possible to construct MCMC algorithms with flows other than the Hamiltonian flow described above. One example is A-NICE-MC (Song et al., 2017), which is similar to HMC but constructs the proposal using an arbitrary volume-preserving flow $T(\cdot;{\bm{\phi}})$ parameterized by ${\bm{\phi}}$. Song et al. (2017) use the NICE model of Dinh et al. (2015), but their method applies to any volume-preserving flow more generally. Given a state $({\bm{\eta}},{\mathbf{v}})$, a new state is proposed which is equal to either $T({\bm{\eta}},{\mathbf{v}};{\bm{\phi}})$ or $T^{-1}({\bm{\eta}},{\mathbf{v}};{\bm{\phi}})$ with equal probability. This proposal is symmetric and so it cancels in the Metripolis–Hastings ratio. The parameters ${\bm{\phi}}$ are tuned to the distribution of interest using adversarial training.

While the reparameterization above is relatively straightforward, there is still the crucial issue of how to set or optimize the parameters of $T$. Titsias (2017) interleaves runs of the chain with updates to the flow’s parameters, performing optimization by maximizing the unnormalized reparameterized target under the last sample from a given run:

However, this choice is a heuristic, and is not guaranteed to encourage the chain’s mixing. As Hoffman et al. (2019) point out, such a choice may emphasize mode finding. As an alternative, Hoffman et al. (2019) propose fitting the flow model to $p({\bm{\eta}})=\pi({\bm{\eta}})/Z$ first via variational inference and then running Hamiltonian Monte Carlo on the reparameterized model, using a sample from the flow to initialize the chain.

2.3 Variational Inference

We can also use normalizing flows to fit distributions over latent variables or model parameters. Specifically, flows can usefully serve as posterior approximations for local (Rezende and Mohamed, 2015; van den Berg et al., 2018; Kingma et al., 2016; Tomczak and Welling, 2016) and global (Louizos and Welling, 2017) variables.

For example, suppose we wish to infer variables ${\bm{\eta}}$ given some observation ${\mathbf{x}}$. In variational inference with normalizing flows, we use a (trained) flow-based model $q({\bm{\eta}};{\bm{\phi}})$ to approximate the posterior as follows:

Normalizing flows can be thought of as implementing a ‘generalized reparameterization trick’ (Kingma and Welling, 2014a; Rezende et al., 2014; Kingma and Welling, 2014b), as they leverage a transformation of a fixed distribution to draw samples from a distribution of interest. Flows therefore define flexible approximate posteriors that are readily reparameterizable by design.

2.4 Likelihood-Free Inference

Models are often implicit, meaning they are not defined in terms of a likelihood function $p({\mathbf{x}}\,|\,{\bm{\eta}})$ that describes how observable variables ${\mathbf{x}}$ depend on model parameters ${\bm{\eta}}$. Rather, they come in the form of a simulator that takes in parameters ${\bm{\eta}}$ and simulates variables ${\mathbf{x}}$ (Diggle and Gratton, 1984). Such simulator-based models are common in scientific fields such as cosmology (Alsing et al., 2018), high-energy physics (Brehmer et al., 2018), and computational neuroscience (Gonçalves et al., 2020). Inferring the parameters ${\bm{\eta}}$ of a simulator-based model given observed data ${\mathbf{x}}$ is often referred to as likelihood-free inference (Papamakarios, 2019), simulation-based inference (Cranmer et al., 2020), or approximate Bayesian computation (Beaumont et al., 2002; Beaumont, 2010). The typical assumption in likelihood-free inference is that it is easy to simulate variables ${\mathbf{x}}$ from the model given ${\bm{\eta}}$, but it is intractable to evaluate the likelihood $p({\mathbf{x}}\,|\,{\bm{\eta}})$.

Normalizing flows are a natural fit for likelihood-free inference, especially flows conditioned on side information (e.g. Winkler et al., 2019; Ardizzone et al., 2019). Assuming a tractable prior distribution $p({\bm{\eta}})$ over the parameters of interest, we can generate a data set $\left\{({\bm{\eta}}_{n},{\mathbf{x}}_{n})\right\}_{n=1}^{N}$ where ${\bm{\eta}}_{n}\sim p({\bm{\eta}})$ and ${\mathbf{x}}_{n}$ is simulated from the model with parameters ${\bm{\eta}}_{n}$. In other words, $({\bm{\eta}}_{n},{\mathbf{x}}_{n})$ is a joint sample from $p({\bm{\eta}},{\mathbf{x}})=p({\bm{\eta}})\,p({\mathbf{x}}\,|\,{\bm{\eta}})$. Then, using the techniques described in Section 6.1, we can fit a flow-based model $q({\bm{\eta}}\,|\,{\mathbf{x}})$ conditioned on ${\mathbf{x}}$ to the generated data set $\left\{({\bm{\eta}}_{n},{\mathbf{x}}_{n})\right\}_{n=1}^{N}$ in order to approximate the posterior $p({\bm{\eta}}\,|\,{\mathbf{x}})$ (Greenberg et al., 2019; Gonçalves et al., 2020). Alternatively, we can fit a flow-based model $q({\mathbf{x}}\,|\,{\bm{\eta}})$ conditioned on ${\bm{\eta}}$ in order to approximate the intractable likelihood $p({\mathbf{x}}\,|\,{\bm{\eta}})$ (Papamakarios et al., 2019). In either case, the trained flow-based model is useful for various downstream tasks that require density evaluation or sampling.

3 Using Flows for Representation Learning

Flows also have applications as building blocks for downstream tasks. We discuss two such cases, namely supervised learning and reinforcement learning.

Invertible ResNets—in addition to implementing residual flows, as discussed in Section 3.3—have been explored for classification (Jacobsen et al., 2018; Behrmann et al., 2019). This line of work exploits invertibility for other end-purposes than density computation. The first is for engineering improvements: to reduce the model’s memory footprint by obviating the need to store activations for backpropagation (Gomez et al., 2017). The second is for improved model interpretability and understanding of the mechanics of deep learning. Jacobsen et al. (2018) showed that an invertible ResNet could be trained to nearly the same ImageNet accuracy as a non-invertible ResNet. This achievement may help us understand to what degree discarding information is crucial to deep learning’s success (Tishby and Zaslavsky, 2015). Jacobsen et al. (2019) used invertible architectures to study the relationship between invariance and vulnerability to adversarial attacks.

Flows can be used for joint generative and predictive modeling by using them as the core component of a hybrid model (Nalisnick et al., 2019). Like in invertible ResNets, the flow is used as a deep, neural feature extractor, but unlike in ResNets, the architecture is chosen such that the Jacobian determinant is tractable. Leveraging the generative model ${\mathbf{x}}=T({\mathbf{u}})$ to reparameterize the joint density, we can write:

3.2 Reinforcement Learning

Finally, in this section we give two examples of how normalizing flows have been used thus far for reinforcement learning (RL).

Haarnoja et al. (2018) and Ward et al. (2019) use such a policy within the maximum entropy and soft actor-critic frameworks respectively.

Schroecker et al. (2019) use normalizing flows within the imitation-learning paradigm for control problems. Given the observed expert (continuous) state-action pairs $\mathopen{}\left(\bar{{\mathbf{s}}},\bar{{\mathbf{a}}}\right)\mathclose{}$, a core challenge to imitation learning is accounting for unobserved intermediate states. Schroecker et al. (2019) use conditional flows to simulate these intermediate states and their corresponding actions. Specifically, for a state $\bar{{\mathbf{s}}}_{t+j}$ ($j\geq 1$), predecessor state-action pairs are sampled from the model:

Both terms on the right-hand side are defined by MAFs (Papamakarios et al., 2017).

Conclusions

We have described normalizing flows and their use for probabilistic modeling and inference. We addressed key issues such as their expressive power (Section 2.2) and the fundamentals underlying their construction (both in discrete and continuous time). We also described the general principle of probability transformations (Section 5.1) and its implications for defining flows beyond Euclidean space. In particular, we showed that discrete domains, mixtures of flows, and extensions to Riemannian manifolds all follow from this generalized perspective. Lastly, we summarized the primary applications of flows (Section 6): tasks ranging from density estimation to likelihood-free inference to classification.

While many flow designs and specific implementations will inevitably become out-of-date as work on normalizing flows continues, we have attempted to isolate foundational ideas that will continue to guide the field well into the future. One of these keystone principles is the chain rule of probability and its relationship to transformations with a triangular Jacobian. Autoregressive flows stand on these two pillars, with the former underlying their expressive power and the latter providing their efficient implementation. Similarly, the Banach fixed-point theorem provides the mathematical foundation for contractive residual flows. While alternative parameterizations of the translation function $g_{\bm{\phi}}$ or normalization strategies may be developed, the underlying Lipschitz constraints cannot be deserted without violating the fixed-point theorem.

Throughout the text we emphasized crucial implementation notes that guide a successful application of flows. Perhaps of foremost importance is determining the computational constraints on evaluating the forward and inverse transformations. As we showed in Section 2, sampling and density evaluation place distinct demands on the transformation. Since flows have no inherent requirement as to whether $T$ should implement $\mathcal{U}\rightarrow\mathcal{X}$ or vice versa, we are free to choose which direction better suits our application. If either sampling or density estimation is the primary objective—but not both—then autoregressive flows present an attractive, flexible class of model. Yet if both sampling and density evaluation must be done often or quickly, then implementing the autoregressive flow with a coupling-based conditioner will make both operations efficient at the cost of expressive power. On the other hand, non-autoregressive flows such as linear and residual flows allow for interaction between all dimensions at each step in the flow. While these interactions can be useful at times—making linear flows good for permuting variables between successive autoregressive flows and residual flows highly expressive—other limitations arise. For instance, contractive residual flows typically require iterative algorithms for sampling and density evaluation. As all flow constructions present trade-offs of some form, we hope this article provides a coherent and accessible summary to guide practitioners through these choice points.

Looking forward, the obstacles preventing wider application of normalizing flows are similar in spirit to those faced by any probabilistic model. However, unlike other probabilistic models that require approximate inference as they scale, flows usually admit analytical calculations and exact sampling even in high dimensions. Rather, the difficulty is transferred to the construction of the flow’s transformation: how can we define ever more flexible transformations while keeping exact density evaluation and sampling computationally tractable? This is currently the focus of much work and will likely remain a core issue for some time. More study of the theoretical properties of flows is also needed. Understanding their approximation capabilities for finite sample and finite depth settings would help practitioners select which flow classes are best for a given application. Our discussion of generalizations in Section 5 will hopefully provide grounding as well as inspiration for this next wave of developments in the theory and application of normalizing flows.

We would like to thank Ivo Danihelka for his invaluable feedback on the manuscript, and the anonymous reviewers for their many improvement suggestions. We also thank Hyunjik Kim and Sébastien Racanière for useful discussions on a wide variety of flow-related topics.

A Proof of KL Dualities

B Constructing Linear Flows

A common approach to efficiently parameterizing invertible matrices is via matrix decompositions. The idea is to decompose the matrix ${\mathbf{W}}$ into a product of structured matrices, each of which can be easily constrained to be invertible and has a $\mathcal{O}\mathopen{}\left(D^{2}\right)\mathclose{}$ inverse and $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ determinant. Below we discuss a few such parameterizations.

Every $D\times D$ matrix ${\mathbf{W}}$ can be written as follows:

where ${\mathbf{P}}$ is a permutation matrix, ${\mathbf{L}}$ is lower triangular, ${\mathbf{U}}$ is upper triangular, and all three are of size $D\times D$. We can easily constrain ${\mathbf{W}}$ to be invertible by restricting ${\mathbf{L}}$ and ${\mathbf{U}}$ to have positive diagonal entries. In that case, the absolute determinant of ${\mathbf{W}}$ can be computed in $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ time by:

Moreover, the linear system ${\mathbf{P}}{\mathbf{L}}{\mathbf{U}}{\mathbf{z}}={\mathbf{z}}^{\prime}$ can be solved by (a) undoing the permutation, (b) solving a lower-triangular system, and (c) solving an upper-triangular system. Solving triangular systems can be done in time $\mathcal{O}\mathopen{}\left(D^{2}\right)\mathclose{}$ by forward/backward substitution, hence inverting the PLU flow can also be done in time $\mathcal{O}\mathopen{}\left(D^{2}\right)\mathclose{}$. In practice, ${\mathbf{P}}$ is typically fixed to a chosen or random permutation, and ${\mathbf{L}}$ and ${\mathbf{U}}$ are learned.

Multiplication with a triangular matrix can be viewed as a special case of an autoregressive flow with affine transformer and linear conditioner. Hence, the PLU flow is simply a composition of two such autoregressive flows with opposite order, followed by a permutation. Linear flows based on the PLU decomposition have been proposed by e.g. Oliva et al. (2018); Kingma and Dhariwal (2018). A similar decomposition in the case where ${\mathbf{W}}$ is a convolution operator was proposed by Hoogeboom et al. (2019b) and was termed emerging convolution.

Another way to decompose a $D\times D$ matrix is using the QR decomposition as follows:

where ${\mathbf{W}}$ is an orthogonal matrix and ${\mathbf{R}}$ is upper triangular, both of size $D\times D$. In fact, we can parameterize any invertible matrix using the QR decomposition, even if we restrict the diagonal entries of ${\mathbf{R}}$ to be positive. In that case, the absolute determinant of ${\mathbf{W}}$ is given in $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$ time by:

Inverting the linear system ${\mathbf{Q}}{\mathbf{R}}{\mathbf{z}}={\mathbf{z}}^{\prime}$ can be done in $\mathcal{O}\mathopen{}\left(D^{2}\right)\mathclose{}$ time by first multiplying by ${\mathbf{Q}}^{\top}$ and then solving a triangular system by forward/backward substitution. We can also think of the QR flow as a linear autoregressive flow followed by an orthogonal flow; in the next paragraph we will discuss in more detail how to parameterize an orthogonal flow. The QR flow was proposed by Hoogeboom et al. (2019b).

Nonetheless, an effective parameterization of orthogonal matrices can be challenging (see e.g. Shepard et al., 2015; Lezcano-Casado and Martínez-Rubio, 2019, for reviews on the subject). As we discussed earlier, either the orthogonal matrices with determinant $1$ or those with determinant $-1$ can be continuously parameterized, but not both. In the normalizing-flows literature, the following parameterizations have been explored, each with their own strengths an weaknesses.

The exponential map (Golinski et al., 2019). Given a skew-symmetric matrix ${\mathbf{A}}$, i.e. a $D\times D$ matrix such that ${\mathbf{A}}^{\top}=-{\mathbf{A}}$, the matrix exponential ${\mathbf{Q}}=\exp{\mathbf{A}}$ is always an orthogonal matrix with determinant $1$. Moreover, any orthogonal matrix with determinant $1$ can be written this way. However, computing the matrix exponential takes in general $\mathcal{O}\mathopen{}\left(D^{3}\right)\mathclose{}$ time, so this parameterization is only suitable for small-dimensional data.

The Cayley map (Golinski et al., 2019). Again given a skew-symmetric matrix ${\mathbf{A}}$ as above, the matrix ${\mathbf{Q}}=({\mathbf{I}}+{\mathbf{A}})({\mathbf{I}}-{\mathbf{A}})^{-1}$ is always orthogonal with determinant $1$. However, this parameterization also takes $\mathcal{O}\mathopen{}\left(D^{3}\right)\mathclose{}$ time to compute, and can only parameterize those orthogonal matrices that don’t have $-1$ as an eigenvalue.

Hence, an orthogonal matrix can be parameterized by a product of $K$ such matrices, i.e. ${\mathbf{Q}}=\prod_{k=1}^{K}{\mathbf{H}}_{k}$, where $K$ doesn’t need to be greater than $D$. Each Householder matrix has determinant $-1$, so the determinant of ${\mathbf{Q}}$ is $(-1)^{K}$. Each Householder transformation can be computed in time $\mathcal{O}\mathopen{}\left(D\right)\mathclose{}$, so an orthogonal transformation parameterized this way can be computed in time $\mathcal{O}\mathopen{}\left(KD\right)\mathclose{}$. The Householder parameterization is not unique and contains saddle points (for example, any permutation of the vectors $\left\{{\mathbf{v}}_{1},\ldots,{\mathbf{v}}_{K}\right\}$ gives the same orthogonal matrix) which can make optimization harder.

References