Normalizing flows as a generalization of variational autoencoders?

Question

Normalizing flows are often introduced as a way of restricting the rigid priors that are placed on the latent variables in Variational Autoencoders.

For example, from the Pyro docs:

In standard probabilistic modeling practice, we represent our beliefs over unknown continuous quantities with simple parametric distributions like the normal, exponential, and Laplacian distributions. However, using such simple forms, which are commonly symmetric and unimodal (or have a fixed number of modes when we take a mixture of them), restricts the performance and flexibility of our methods.

[...]

Normalizing Flows [1-4] are a family of methods for constructing flexible learnable probability distributions, often with neural networks, which allow us to surpass the limitations of simple parametric forms.

However, it is my understanding that the latent variables $\mathbf{z}$ in normalizing flows are still often modeled as standard Gaussians. For example, the same Pyro docs page gives the following example of a Normalizing Flow:

base_dist = dist.Normal(torch.zeros(2), torch.ones(2))
spline_transform = T.spline_coupling(2, count_bins=16)
flow_dist = dist.TransformedDistribution(base_dist, [spline_transform])

So the base distribution used is a Gaussian, same as in Variational Autoencoders?

---

To summarize, I have read the statement that normalizing flows somehow "relax" the limitations of Variational Autoencoders, and in particular the limited expressiveness of the latent variable priors that are used, but I am not able to understand why that is the case. Could someone provide a layman explanation and point me to some resources for further reading?

My understanding is that both VAEs and Normalizing Flows map the input $\mathbf{x}$ to some latent representation $\mathbf{z}$ and then attempt to reconstruct the input $\mathbf{\hat{x}}$ from the latent representation. Both VAEs and Normalizing Flows usually model the latent variables $\mathbf{z}$ as coming from independent univariate normal distributions (AFAIK?). Normalizing flows have many more restrictions on the types of neural networks that can be used as the "encoder" and "decoder" (i.e. the model has to be bijective and invertable). So how do normalizing flows relax some of the assumptions of VAEs?

Source: https://lilianweng.github.io/lil-log/2018/10/13/flow-based-deep-generative-models.html

Answer 1

Let's compare and contrast the two methods, VAE and normalizing flows.

Similarly, both techniques optimize (by maximizing) the log-likelihood. However, it is not so apparent (from the image) how they construct the latent space $z$ that defines the joint posterior distribution $p(z,x)$.

The straightforward answer to your question is normalizing flows construct the latent space as a chain of series of simple, separate and independent mappings in the following form

$x \: = \: f_{\phi} \: = \: f_{k} \cdot \ldots \cdot f_{2} \cdot f_{1} (z_{0}),$

where each $f_{\phi}$ is a differentiable and invertible flow layer function, also known as a bijection. The sequential and independent flow layers transform a random variable $z_{0}$ to a new random variable $x$ (during inference). By design, they leverage tractable density evaluation, which is in contrast to approximate inference of VAEs, therefore overcoming the deficiencies of learning a variational lower bound on the log-likelihood. The computation of the Jacobian log determinant for each basic bijection operator is analytically determinable using the change of variables formula:

$P_X(x) \: = \: P_Z(z) \cdot \big\vert \: \det \: \big( \frac {dz} {dx} \: \big) \: \big\vert \: \: where \: \: x \: = \: f(z)$

This ultimately is the property that enables the final flow to produce expressive, flexible posteriors. It clarifies why the design, construction, and coupling of flow layers are core research problems and why restrictions on the types of neural network architectures are necessary (to impose invertibility and easy computation of the Jacobian of the determinant). This construction differs from VAEs, which typically model the latent space as a single Gaussian distribution. Accordingly, the composition of the latent space of the two generative methods also affects how sampling from the posterior is performed.

For those curious to link the said techniques to more state-of-the-art generative algorithms, diffusion models can be transformed into continuous normalizing flows (CNFs) and interpreted as a specific form of a Markovian Hierarchical Variational Autoencoder.

The following excerpts are taken from my book on variational inference. Learn more on the topic by visiting https://www.thevariationalbook.com/