Normalizing Flow Models

We want to design a latent variable model with tractable likelihoods. Desirable properties of model distribution \( p_{\theta}(\mathbf{x}) \) include easy-to-evaluate (closed form density) which is useful for training and easy-to-sample which is useful for generation.

Many simple distributions satisfy the above property, but true data distributions are more complex. The idea behind flow models is to map simple distributions to complex distributions through an invertible transformation.

Similar to VAE, we start from a simple prior: \( \mathbf{z} \sim \mathcal{N}(0,I)=p(\mathbf{z}) \), then we transform via \( p(\mathbf{x}|\mathbf{z})=\mathcal{N}(\mu_{\theta}(\mathbf{z}), \Sigma_{\theta}(\mathbf{z})) \)

The problem is that, training requires \( p_{\theta}(\mathbf{x}) \), but computing this requires enumerating over \( \mathbf{z} \). Thus, we want to easily “invert” \( p(\mathbf{x}|\mathbf{z}) \) and compute \( p(\mathbf{z}|\mathbf{x}) \) by design. The idea here is to make \( \mathbf{x} = f_{\theta}(\mathbf{z}) \) a deterministic and invertible function, so for any \( \mathbf{x} \) there is a unique corresponding \( \mathbf{z} \), thus we eliminate enumeration.

Consider a directed, latent-variable model over observed variables \(X\) and latent variables \(Z\). In a normalizing flow model, the mapping between \(Z\) and \(X\), given by \( \mathbf{f}_{\theta}:\mathbb{R}^{n}\mapsto\mathbb{R}^{n} \), is deterministic and invertible such that \( X=\mathbf{f}_{\theta}(Z) \) and \( Z=\mathbf{f}_{\theta}^{-1}(X) \).

Using Change of Variable Formula, the marginal likelihood \( p(\mathbf{x}) \) is given by

\begin{equation} p_{X}(\mathbf{x}; \theta) = p_{Z}(\mathbf{f}^{-1}_{\theta}(\mathbf{x})) \left\vert \det \left( \frac{\partial \mathbf{f}^{-1}_{\theta}(\mathbf{x})}{\partial \mathbf{x}} \right) \right\vert \end{equation}

From the above formula, “normalizing” - change of variables gives a normalized density after applying an invertible transformation; “flow” - invertible transformation can be composed with each other, like

\[ \mathbf{z}_{m} = (\mathbf{f}_{\theta}^{m} \circ \mathbf{f}_{\theta}^{m-1} \circ \dots \circ \mathbf{f}_{\theta}^{1})(\mathbf{z}_{0}) := \mathbf{f}_{\theta}(\mathbf{z}_{0}) \]

Sampling

The sampling is simple that we first sample latent variable \( \mathbf{z}\sim p(\mathbf{z}) \), then transform it with \( \mathbf{x} = \mathbf{f}_{\theta}(\mathbf{z}) \)

Latent representations

To represent the latent variable, it’s also straightforward to use inverse transformation \( \mathbf{z} = \mathbf{f}^{-1}_{\theta}(\mathbf{x}) \)