Course:CPSC522/NormalizingFlows

Title

This page covers Normalizing Flows, a class of generative model which transforms a latent distribution to match the data distribution using an invertible transformation.

Principal Author: Matthew Niedoba
Collaborators:

Abstract

Normalizing Flows are an attractive class of generative model which maps each element of a simple analytical distribution one to one to a complex data distribution through a bijective transformation. Notably, this construction allows for both sampling from the data distribution and exact likelihood computation. In this page, we introduce normalizing flows and detail their construction and application.

Builds on

An understanding of Probability and Probability general semantics is required to understand the construction of Normalizing Flows. The bijective transforms of normalizing flows are typically parameterized by Neural Networks.

Related Pages

Normalizing Flows are just member of the broad class of generative models. Other types of generative models include:

• Variational Auto-Encoders
• Generative Adversarial Networks
• Diffusion Probabilistic Model

Method

Source and Target Distributions

Normalizing Flows model the data distribution by transforming the latent space (right) into the data space (left) through an invertible mapping. Source: Dinh, L., Sohl-Dickstein, J., & Bengio, S. (2016)

The aim of normalizing flows is to model a target distribution $p(\mathbf{𝐱})$, $\mathbf{𝐱}\in {\mathbb{ℝ}}^D$. Since the analytic form of $p(\mathbf{𝐱})$ is not generally known, normalizing flows model it by transforming samples from a source distribution $p(\mathbf{𝐳})$ through a transformation $T$, generally parameterized with some parameters $\theta$.$$\mathbf{𝐱}=T(\mathbf{𝐳})\text{where}\mathbf{𝐳}\sim p(\mathbf{𝐳}).$$The goal of normalizing flows is to achieve a one to one, invertible mapping between $\mathbf{𝐱}$ and $\mathbf{𝐳}$. To this end, normalizing flows restrict the transformation $T$ to diffeomorphisms - transformations which are invertible and where both $T$ and $T^{-1}$ are differentiable. Since $T$ must be invertible, we require that $\mathbf{𝐳}$ must be the same dimensionality as $\mathbf{𝐱}$. Using the inverse transformation, we can transform data samples to samples from the source distribution$$\mathbf{𝐳}=T^{-1}(\mathbf{𝐱})\text{where}\mathbf{𝐱}\sim p(\mathbf{𝐱}).$$

Likelihood Computations

The previous section illustrates how to sample data from a normalizing flow by transforming samples from a source distribution through the transformation $T$. Another important task in generative modelling is computing the likelihood of data under the model. Due to the one to one mapping between the source and target distributions, normalizing flows allow exact computation of likelihoods through the change of variables formula$$p(\mathbf{𝐱})=p\left(T^{-1}\left(\mathbf{𝐱}\right)\right)\left|\text{det}\left[\frac{\partial T^{-1}\left(\mathbf{𝐱}\right)}{\partial \mathbf{𝐱}}\right]\right|.$$Here, $\text{det}\left[\frac{\partial T^{-1}\left(\mathbf{𝐱}\right)}{\partial \mathbf{𝐱}}\right]$ refers to the determinant of the Jacobian of $T^{-1}$ with respect to $\mathbf{𝐱}$. Since $\mathbf{𝐳}=T^{-1}(\mathbf{𝐱})$, and using the identity $\text{det}(J^{-1})=\text{det}(J{)}^{-1}$ we can also write$$p(\mathbf{𝐱})=p\left(\mathbf{𝐳}\right){\left|\text{det}\left[\frac{\partial T\left(\mathbf{𝐳}\right)}{\partial \mathbf{𝐳}}\right]\right|}^{-1}.$$

Training Objective

Like many generative models, the goal of training is to approximate the true data distribution $p(\mathbf{𝐱})$ with our model $p_{\theta }(\mathbf{𝐱})$ which transforms data from our source distribution to the target distribution through the parameterized transformation $T$. The objective of training is to minimize the KL divergence between these two distributions. That is to minimize$$L=\underset{\theta }{\min }{D}_{KL}(p(\mathbf{𝐱})||p_{\theta }(\mathbf{𝐱}))$$Note that this is the so-called forward KL divergence, which is an expectation over samples from the true data distribution $p(\mathbf{𝐱})$, usually in the form of a dataset of examples. Noting that $D_{KL}(P||Q)={\mathbb{𝔼}}_{x\sim P(x)}[\log \left(\frac{P(x)}{Q(x)}\right)]$, we can rewrite the loss as the sum of two expectations$$L=\underset{\theta }{\min }{\mathbb{𝔼}}_{\mathbf{𝐱}\sim p(\mathbf{𝐱})}[\log (p(\mathbf{𝐱}))]-{\mathbb{𝔼}}_{\mathbf{𝐱}\sim p(\mathbf{𝐱})}[\log (p_{\theta }(\mathbf{𝐱}))]$$Since the first term is constant with respect to the parameters, we can see that minimizing the KL divergence of a normalizing flow with the data distribution is equivalent to minimizing the negative log likelihood of the data under the normalizing flow model

$$L=\underset{\theta }{\min }-{\mathbb{𝔼}}_{\mathbf{𝐱}\sim p(\mathbf{𝐱})}[\log (p_{\theta }(\mathbf{𝐱}))]$$During training, we compute the exact likelihood $p_{\theta }(\mathbf{𝐱})$ by using the change of variables formula.

Finite Normalizing Flows

The key challenge in designing a normalizing flow is choosing the structure of the transformation $T$ which maps between the source in target distributions. Such a transformation must be complex enough to model the data distribution. However, generating such a complex transformation in one shot is difficult, especially when it must be invertible. Instead of computing the transformation in one shot, finite normalizing achieve the required complexity through the composition of a finite number of simpler transformations.$$T=T_N\circ \cdots \circ T_1$$Each transformation $T_n$ can be thought of as a miniature normalizing flow, transforming an intermediate source distribution $p({\mathbf{𝐳}}_{n-1})$ into an intermediate target distribution $p({\mathbf{𝐳}}_n)$ via the relation ${\mathbf{𝐳}}_n=T_n\left({\mathbf{𝐳}}_{n-1}\right)$ for $n=1,\dots ,N$. We set ${\mathbf{𝐳}}_0$ equal to the samples from our original source distribution (usually a multivariate Gaussian) and aim to have ${\mathbf{𝐳}}_N$ match the target distribution $p(\mathbf{𝐱})$. With the finite normalizing flow construction, the determinant of the Jacobian of the overall transformation is equal to the product of the determinants of the Jacobians of each transform$$\text{det}\left[\frac{\partial T\left(\mathbf{𝐳}\right)}{\partial \mathbf{𝐳}}\right]=\prod _{n=1}^N\text{det}\left[\frac{\partial T_n\left({\mathbf{z}}_{\mathbf{𝐧}\mathbf{-}\mathbf{𝟏}}\right)}{\partial {\mathbf{z}}_{\mathbf{𝐧}\mathbf{-}\mathbf{𝟏}}}\right]$$By selecting individual transforms which are invertible and for which the determinant of the Jacobian can be computed efficiently, we ensure that the overall normalizing flow is also invertible with an easy to compute Jacobian determinant. In the next sections, we discuss choices of transforms which have these properties and can be composed to construct more complex normalizing flows.

Coupling Flows

Coupling transformations, introduced by , aim to make the Jacobian matrix triangular by partitioning $\mathbf{𝐳}$ into two parts, ${\mathbf{𝐳}}_a$ and ${\mathbf{𝐳}}_b$. Then, we define the transformation of each partition seperately$$T({\mathbf{z}}_{\mathbf{𝐚}})={\mathbf{𝐳}}_{a}T({\mathbf{𝐳}}_b)={\mathbf{z}}_{\mathbf{𝐛}}+f_{\theta }({\mathbf{𝐳}}_a)$$Coupling transformations are easily invertible, and the Jacobian matrix has the form $$\frac{\partial T(\mathbf{𝐳})}{\partial \mathbf{𝐳}}=\left[\begin{array}{cc}\mathbf{𝐈}& \mathbf{𝟎}\\ \frac{\partial f}{\partial {\mathbf{z}}_{\mathbf{𝐚}}}& \mathbf{𝐈}\end{array}\right]$$Since the Jacobian is lower triangular, the determinant is equal to one.

The partitioning of $\mathbf{𝐳}$ is modified with each transformation layer such that all components of $\mathbf{𝐳}$ are transformed by the end of the flow.

RealNVP

RealNVP extends the coupling transformations introduced in by adding a scaling to the transformation. Specifically$$T({\mathbf{𝐳}}_b)={\mathbf{𝐳}}_b\cdot e^{s_{\theta }({\mathbf{𝐳}}_a)}+t_{\theta }({\mathbf{𝐳}}_a)$$Here, the transformation is parameterized by two neural networks: ${\mathbf{𝐬}}_{\theta }({\mathbf{𝐳}}_a)$ and ${\mathbf{𝐭}}_{\theta }({\mathbf{𝐳}}_a)$ which control the scale and shift of the transformation. With this modification, the transformation is still easily invertible, and the Jacobian is still lower triangular$$\frac{\partial T(\mathbf{𝐳})}{\partial \mathbf{𝐳}}=\left[\begin{array}{cc}\mathbf{𝐈}& \mathbf{𝟎}\\ \frac{\partial t}{\partial {\mathbf{z}}_{\mathbf{𝐚}}}& \text{diag}(\text{exp}(s_{\theta }({\mathbf{z}}_{\mathbf{𝐚}}))\end{array}\right]$$where diag indicates a diagonal matrix. Since the Jacobian is lower triangular, the determinant is still the trace of the matrix.

Planar Flows

Planar flows are a type of transformation which allow for linear computation of the determinant of the Jacobian. They have the form$$T(\mathbf{𝐳})=\mathbf{𝐳}+\mathbf{𝐮}h({\mathbf{𝐰}}^{\mathrm{\top }}\mathbf{𝐳}+b).$$Here, the parameters $\theta$ of the transformation are $\mathbf{𝐮}\in {\mathbb{ℝ}}^D,\mathbf{𝐰}\in {\mathbb{ℝ}}^D,b\in \mathbb{ℝ}$. The determinant of the Jacobian is easily computed as $$\left|\text{det}\left[\frac{\partial T\left(\mathbf{𝐳}\right)}{\partial \mathbf{𝐳}}\right]\right|=|1+{\mathbf{𝐮}}^{\mathrm{\top }}{h}^{\prime }({\mathbf{𝐰}}^{\mathrm{\top }}\mathbf{𝐳}+b)\mathbf{𝐰}|$$With planar flows, $T$ and its determinant are easy to compute, but are not easily invertible. As a result, the authors of this method train their flow using the reverse KL. In this setup, they minimize the divergence $D(p(T_{\theta }(z))||p(x))$ by drawing samples from the source distribution.

Continuous Normalizing Flows

Synthetic Celebrity portraits generated using Glow, a normalizing flow. Source: Kingma, D. P., & Dhariwal, P. (2018).

Continuous normalizing flows consider the case of extending finite normalizing flows to an infinite number if infinitesimal transformations. If we let ${\mathbf{𝐳}}_0$ be a variable from our source distribution and let ${\mathbf{𝐳}}_1$be a variable from our target distribution, then the continuous normalizing flow transforming ${\mathbf{𝐳}}_0$ to ${\mathbf{𝐳}}_1$ is given by $${\mathbf{𝐳}}_1={\mathbf{𝐳}}_0+\underset{0}{\overset{1}{\int }}{f}_{\theta }({\mathbf{𝐳}}_t,t)dt$$The log density of the resulting distribution is given by the instantaneous change of variables formula, another ODE:$$\log p({\mathbf{𝐳}}_1)=\log p({\mathbf{𝐳}}_0)-\underset{0}{\overset{1}{\int }}\text{Tr}(\frac{\partial f}{\partial {\mathbf{𝐳}}_t})dt$$Notably, unlike the finite normalizing flows, computing the log density only requires evaluating the trace of the jacobian, instead of the determinant. This allows for more freedom in selecting $f$, but at the cost of using an numerical ODE solver for sampling and likelihood evaluation. Training continuous normalizing flows is also challenging, as it requires backpropagating through the ODE solver.

Applications

Normalizing Flows have primarily been used to model image data, such as in . However, some other applications exist such as for text modelling or audio . Normalizing flows have become less popular recently, possibly because the requirements of invertibility and easy to compute Jacobian determinants place large restrictions on the class of transformations that can be used. Instead, many practitioners have shifted to using other generative models, such as Generative Adversarial Networks, Variational Auto Encoders or Diffusion Probabilistic Models. However, recent work has lead to renewed interest in normalizing flows due to a simplification of the training objective for continuous normalizing flows through flow matching , which may be used for a new generation of powerful normalizing flow models.