Flow-based generative model


A flow-based generative model is a generative model used in machine learning that explicitly models a probability distribution by leveraging normalizing flow, which is a statistical method using the Probability density function#Function of [random variables and change of variables in the probability density function|change-of-variable] law of probabilities to transform a simple distribution into a complex one.
The direct modeling of likelihood provides many advantages. For example, the negative log-likelihood can be directly computed and minimized as the loss function. Additionally, novel samples can be generated by sampling from the initial distribution, and applying the flow transformation.
In contrast, many alternative generative modeling methods, such as variational autoencoders (VAEs), generative adversarial networks (GANs), or diffusion models, do not explicitly represent the likelihood function.

Method

Let be a random variable with distribution.
For, let be a sequence of random variables transformed from. The functions should be invertible, i.e. the inverse function exists. The final output models the target distribution.
The log likelihood of is :
Learning probability distributions by differentiating such log Jacobians originated in the Infomax approach to ICA, which forms a single-layer flow-based model. Relatedly, the single layer precursor of conditional generative flows appeared in.
To efficiently compute the log likelihood, the functions should be easily invertible, and the determinants of their Jacobians should be simple to compute. In practice, the functions are modeled using deep neural networks, and are trained to minimize the negative log-likelihood of data samples from the target distribution. These architectures are usually designed such that only the forward pass of the neural network is required in both the inverse and the Jacobian determinant calculations. Examples of such architectures include NICE, RealNVP, and Glow.

Derivation of log likelihood

Consider and. Note that.
By the change of variable formula, the distribution of is:
Where is the determinant of the Jacobian matrix of.
By the inverse function theorem:
By the identity, we have:
The log likelihood is thus:
In general, the above applies to any and. Since is equal to subtracted by a non-recursive term, we can infer by induction that:

Training method

As is generally done when training a deep learning model, the goal with normalizing flows is to minimize the Kullback–Leibler divergence between the model's likelihood and the target distribution to be estimated. Denoting the model's likelihood and the target distribution to learn, the KL-divergence is:
The second term on the right-hand side of the equation corresponds to the entropy of the target distribution and is independent of the parameter we want the model to learn, which only leaves the expectation of the negative log-likelihood to minimize under the target distribution. This intractable term can be approximated with a Monte-Carlo method by importance sampling. Indeed, if we have a dataset of samples each independently drawn from the target distribution, then this term can be estimated as:
Therefore, the learning objective
is replaced by
In other words, minimizing the Kullback–Leibler divergence between the model's likelihood and the target distribution is equivalent to maximizing the model likelihood under observed samples of the target distribution.
A pseudocode for training normalizing flows is as follows:
  • INPUT. dataset, normalizing flow model.
  • SOLVE. by gradient descent
  • RETURN.

Variants

Planar Flow

The earliest example. Fix some activation function, and let with the appropriate dimensions, then The inverse has no closed-form solution in general.
The Jacobian is.
For it to be invertible everywhere, it must be nonzero everywhere. For example, and satisfies the requirement.

Nonlinear Independent Components Estimation (NICE)

Let be even-dimensional, and split them in the middle. Then the normalizing flow functions are where is any neural network with weights.
is just, and the Jacobian is just 1, that is, the flow is volume-preserving.
When, this is seen as a curvy shearing along the direction.

Real Non-Volume Preserving (Real NVP)

The Real Non-Volume Preserving model generalizes NICE model by:
Its inverse is, and its Jacobian is. The NICE model is recovered by setting.
Since the Real NVP map keeps the first and second halves of the vector separate, it's usually required to add a permutation after every Real NVP layer.

Generative Flow (Glow)

In generative flow model, each layer has 3 parts:
  • channel-wise affine transform with Jacobian.
  • invertible 1x1 convolution with Jacobian. Here is any invertible matrix.
  • Real NVP, with Jacobian as described in Real NVP.
The idea of using the invertible 1x1 convolution is to permute all layers in general, instead of merely permuting the first and second half, as in Real NVP.

Masked Autoregressive Flow (MAF)

An autoregressive model of a distribution on is defined as the following stochastic process:
where and are fixed functions that define the autoregressive model.
By the reparameterization trick, the autoregressive model is generalized to a normalizing flow:The autoregressive model is recovered by setting.
The forward mapping is slow, but the backward mapping is fast.
The Jacobian matrix is lower-diagonal, so the Jacobian is.
Reversing the two maps and of MAF results in Inverse Autoregressive Flow, which has fast forward mapping and slow backward mapping.

Continuous Normalizing Flow (CNF)

Instead of constructing flow by function composition, another approach is to formulate the flow as a continuous-time dynamic. Let be the latent variable with distribution. Map this latent variable to data space with the following flow function:
where is an arbitrary function and can be modeled with e.g. neural networks.
The inverse function is then naturally:
And the log-likelihood of can be found as:
Since the trace depends only on the diagonal of the Jacobian, this allows "free-form" Jacobian. Here, "free-form" means that there is no restriction on the Jacobian's form. It is contrasted with previous discrete models of normalizing flow, where the Jacobian is carefully designed to be only upper- or lower-diagonal, so that the Jacobian can be evaluated efficiently.
The trace can be estimated by "Hutchinson's trick":
Given any matrix, and any random with, we have.
Usually, the random vector is sampled from or .
When is implemented as a neural network, neural ODE methods would be needed. Indeed, CNF was first proposed in the same paper that proposed neural ODE.
There are two main deficiencies of CNF, one is that a continuous flow must be a homeomorphism, thus preserve orientation and ambient isotopy, and the other is that the learned flow might be ill-behaved, due to degeneracy.
By adding extra dimensions, the CNF gains enough freedom to reverse orientation and go beyond ambient isotopy, yielding the "augmented neural ODE".
Any homeomorphism of can be approximated by a neural ODE operating on, proved by combining Whitney embedding theorem for manifolds and the universal approximation theorem for neural networks.
To regularize the flow, one can impose regularization losses. The paper proposed the following regularization loss based on optimal transport theory:where are hyperparameters. The first term punishes the model for oscillating the flow field over time, and the second term punishes it for oscillating the flow field over space. Both terms together guide the model into a flow that is smooth over space and time.

Flows on manifolds

When a probabilistic flow transforms a distribution on an -dimensional smooth manifold embedded in, where, and where the transformation is specified as a function,, the scaling factor between the source and transformed PDFs is not given by the naive computation of the determinant of the Jacobian, but instead by the determinant of one or more suitably defined matrices. This section is an interpretation of the tutorial in the appendix of Sorrenson et al., where the more general case of non-isometrically embedded Riemann manifolds is also treated. Here we restrict attention to isometrically embedded manifolds.
As running examples of manifolds with smooth, isometric embedding in we shall use:
As a first example of a spherical manifold flow transform, consider the normalized linear transform, which radially projects onto the unit sphere the output of an invertible linear transform, parametrized by the invertible matrix :
In full Euclidean space, is not invertible, but if we restrict the domain and co-domain to the unit sphere, then is invertible, with inverse. The Jacobian of, at is, which has rank and determinant of zero; while as explained here, the factor relating source and transformed densities is:.

Differential volume ratio

For, let be an -dimensional manifold with a smooth, isometric embedding into. Let be a smooth flow transform with range restricted to. Let be sampled from a distribution with density. Let, with resultant density. Let be a small, convex region containing and let be its image, which contains ; then by conservation of probability mass:
where volume is given by Lebesgue measure in -dimensional tangent space. By making the regions infinitessimally small, the factor relating the two densities is the ratio of volumes, which we term the differential volume ratio.
To obtain concrete formulas for volume on the -dimensional manifold, we construct by mapping an -dimensional rectangle in coordinate space to the manifold via a smooth embedding function:. At very small scale, the embedding function becomes essentially linear so that is a parallelotope. Similarly, the flow transform, becomes linear, so that the image, is also a parallelotope. In, we can represent an -dimensional parallelotope with an matrix whose column-vectors are a set of edges that span the paralellotope. The volume is given by the absolute value of the determinant of this matrix. If more generally, an -dimensional paralellotope is embedded in, it can be represented with a matrix, say. Denoting the parallelotope as, its volume is then given by the square root of the Gram determinant:
In the sections below, we show various ways to use this volume formula to derive the differential volume ratio.

Simplex flow

As a first example, we develop expressions for the differential volume ratio of a simplex flow,, where. Define the embedding function:
which maps a conveniently chosen, -dimensional representation,, to the embedded manifold. The Jacobian is
To define, the differential volume element at the transformation input, we start with a rectangle in -space, having differential side-lengths, from which we form the square diagonal matrix, the columns of which span the rectangle. At very small scale, we get, with:
To understand the geometric interpretation of the factor, see the example for the 1-simplex in the diagram at right.
The differential volume element at the transformation output, is the parallelotope,, where is the Jacobian of at. Its volume is:
so that the factor cancels in the volume ratio, which can now already be numerically evaluated. It can however be rewritten in a sometimes more convenient form by also introducing the representation function,, which simply extracts the first components. The Jacobian is. Observe that, since, the chain rule for function composition gives:. By plugging this expansion into the above Gram determinant and then refactoring it as a product of determinants of square matrices, we can extract the factor, which now also cancels in the ratio, which finally simpifies to the determinant of the Jacobian of the "sandwiched" flow transformation, :
which, if, can be used to derive the pushforward density after a change of variables, :
This formula is valid only because the simplex is flat and the Jacobian, is constant. The more general case for curved manifolds is discussed below, after we present two concrete examples of simplex flow transforms.

Simplex calibration transform

A calibration transform,, which is sometimes used in machine learning for post-processing of the outputs of a probabilistic -class classifier, uses the softmax function to renormalize categorical distributions after scaling and translation of the input distributions in log-probability space. For and with parameters, and the transform can be specified as:
where the log is applied elementwise. After some algebra the differential volume ratio can be expressed as:
  • This result can also be obtained by factoring the density of the SGB distribution, which is obtained by sending Dirichlet variates through.
While calibration transforms are most often trained as discriminative models, the reinterpretation here as a probabilistic flow allows also the design of generative calibration models based on this transform. When used for calibration, the restriction can be imposed to prevent direction reversal in log-probability space. With the additional restriction, this transform is known in machine learning as temperature scaling.

Generalized calibration transform

The above calibration transform can be generalized to, with parameters and invertible:
where the condition that has as an eigenvector ensures invertibility by sidestepping the information loss due to the invariance:. Note in particular that is the only allowed diagonal parametrization, in which case we recover, while generalization is possible with non-diagonal matrices. The inverse is:
The differential volume ratio is:
If is to be used as a calibration transform, further constraint could be imposed, for example that be positive definite, so that, which avoids direction reversals.
For, and positive definite, then and are equivalent in the sense that in both cases, is a straight line, the slope and offset of which are functions of the transform parameters. For does generalize.
It must however be noted that chaining multiple flow transformations does not give a further generalization, because:
In fact, the set of transformations form a group under function composition. The set of transformations form a subgroup.
Also see: Dirichlet calibration, which generalizes, by not placing any restriction on the matrix,, so that invertibility is not guaranteed. While Dirichlet calibration is trained as a discriminative model, can also be trained as part of a generative calibration model.

Differential volume ratio for curved manifolds

Consider a flow, on a curved manifold, for example which we equip with the embedding function, that maps a set of angular spherical coordinates to. The Jacobian of is non-constant and we have to evaluate it at both input and output. The same applies to, the representation function that recovers spherical coordinates from points on, for which we need the Jacobian at the output. The differential volume ratio now generalizes to:
For geometric insight, consider, where the spherical coordinates are co-latitude, and longitude,. At, we get, which gives the radius of the circle at that latitude. The differential volume is:.
The above [|derivation] for is fragile in the sense that when using fixed functions, there may be places where they are not well-defined, for example at the poles of the 2-sphere where longitude is arbitrary. This problem is sidestepped by generalizing to local coordinates, where in the vicinities of, we map from local -dimensional coordinates to and back using the respective function pairs and. We continue to use the same notation for the Jacobians of these functions, so that the above formula for remains valid.
We can however, choose our local coordinate system in a way that simplifies the expression for and indeed also its practical implementation. Let be a smooth idempotent projection from the projectible set,, onto the embedded manifold. For example:
  • The positive orthant of is projected onto the simplex as:
  • Non-zero vectors in are projected onto the unit sphere as:
For every, we require of that its Jacobian, has rank, in which case is an idempotent linear projection onto the local tangent space. The columns of span the -dimensional tangent space at. We use the notation, for any matrix with orthonormal columns that span the local tangent space. Also note:. We can now choose our local coordinate embedding function, :
Since the Jacobian is injective, a local left inverse, say with Jacobian, exists such that and. In practice we do not need the left inverse function itself, but we do need its Jacobian, for which the above equation does not give a unique solution. We can however enforce a unique solution for the Jacobian by choosing the left inverse as, :
We can now finally plug and into our previous expression for, the differential volume ratio, which because of the orthonormal Jacobians, simplifies to:

Practical implementation

For learning the parameters of a manifold flow transformation, we need access to the differential volume ratio,, or at least to its gradient w.r.t. the parameters. Moreover, for some inference tasks, we need access to itself. Practical solutions include:
  • Sorrenson et al. give a solution for computationally efficient stochastic parameter gradient approximation for
  • For some hand-designed flow transforms, can be analytically derived in closed form, for example the above-mentioned simplex calibration transforms. Further examples are given below in the section on simple spherical flows.
  • On a software platform equipped with linear algebra and automatic differentiation, can be automatically evaluated, given access to only. But this is expensive for high-dimensional data, with at least computational costs. Even then, the slow automatic solution can be invaluable as a tool for numerically verifying hand-designed closed-form solutions.

Simple spherical flows

In machine learning literature, various complex spherical flows formed by deep neural network architectures may be found. In contrast, this section compiles from statistics literature the details of three very simple spherical flow transforms, with simple closed-form expressions for inverses and differential volume ratios. These flows can be used individually, or chained, to generalize distributions on the unit sphere,. All three flows are compositions of an invertible affine transform in, followed by radial projection back onto the sphere. The flavours we consider for the affine transform are: pure translation, pure linear and general affine. To make these flows fully functional for learning, inference and sampling, the tasks are:
  • To derive the inverse transform, with suitable restrictions on the parameters to ensure invertibility.
  • To derive in simple closed form the differential volume ratio,.
An interesting property of these simple spherical flows is that they don't make use of any non-linearities apart from the radial projection. Even the simplest of them, the normalized translation flow, can be chained to form perhaps surprisingly flexible distributions.

Normalized translation flow

The normalized translation flow,, with parameter, is given by:
The inverse function may be derived by considering, for : and then using to get a quadratic equation to recover, which gives:
from which we see that we need to keep real and positive for all. The differential volume ratio is given by Boulerice & Ducharme as:
This can indeed be verified analytically:
  • By a laborious manipulation of.
  • By setting in, which is given below.
Finally, it is worth noting that and do not have the same functional form.

Normalized linear flow

The normalized linear flow,, where parameter is an invertible matrix, is given by:
The differential volume ratio is:
This result can be derived indirectly via the Angular central Gaussian distribution , which can be obtained via normalized linear transform of either Gaussian, or uniform spherical variates. The first relationship can be used to derive the ACG density by a marginalization integral over the radius; after which the second relationship can be used to factor out the differential volume ratio. For details, see ACG distribution.

Normalized affine flow

The normalized affine flow,, with parameters and, invertible, is given by:
The inverse function, derived in a similar way to the normalized translation inverse is:
where. The differential volume ratio is:
The final RHS numerator was expanded from by the matrix determinant lemma. Recalling, the equality between and holds because not only:
but also, by orthogonality of to the local tangent space:
where is the Jacobian of differentiated w.r.t. its input, but not also w.r.t. to its parameter.

Downsides

Despite normalizing flows success in estimating high-dimensional densities, some downsides still exist in their designs. First of all, their latent space where input data is projected onto is not a lower-dimensional space and therefore, flow-based models do not allow for compression of data by default and require a lot of computation. However, it is still possible to perform image compression with them.
Flow-based models are also notorious for failing in estimating the likelihood of out-of-distribution samples. Some hypotheses were formulated to explain this phenomenon, among which the typical set hypothesis, estimation issues when training models, or fundamental issues due to the entropy of the data distributions.
One of the most interesting properties of normalizing flows is the invertibility of their learned bijective map. This property is given by constraints in the design of the models which guarantee theoretical invertibility. The integrity of the inverse is important in order to ensure the applicability of the change-of-variable theorem, the computation of the Jacobian of the map as well as sampling with the model. However, in practice this invertibility is violated and the inverse map explodes because of numerical imprecision.

Applications

Flow-based generative models have been applied on a variety of modeling tasks, including: