Basic Normalizing Flows

Basic Normalizing Flows

In this post, we are going to explore a very basic implementation of normalizing flows, the goal of which is to extend to more intersting scneario’s with more advanced architectures. Ultimately, we’d like to reproduce the results of Pawlowski, J., Urban, J. (2022).

Density estimation

In this section, we follow Shen et al. (2019). Suppose you have some data, $\mathcal{D} = { y^{(i)} }_{i=1}^N$ with $y^{(i)} \in \mathbb{R}^d$, whose underlying density, $p_y you’d like to know. The core idea is to find a transformation (neccesarily a bijection), $y=g(z)$, which maps a simple distribution, $p_z$, to the target distribution $p_y$. The term normalizing flows seems to stem from the flow reducing the target distrbution to a normal distribution.

Let us derive a relation between $p_z$ and $p_z$ using $f=g^{-1}$: \(\begin{aligned} p_y(y) dy &= p_z(z) dz\\ &= p_z( f(y)) | \text{det} \, Dg\,(f(y))|^{-1} dy \\ &= p_z(f(y))| \text{det} \, Df(y)| dy \end{aligned}\) where $Df= \frac{\partial f}{\partial x}$ is the Jacobian. We take the magnitude of the jacobian so as to ensure the pdfs remain positive. In the last line, we used the relation of Jacobians of inverse functions.

If we have a set of bijections which are easily inverted, we can simply compose together functions of this set to produce a more expressive transformation.

\(\begin{aligned} g &= g_N \circ g_{N-1} \circ ... \circ g_1 \\ f &= f_1 \circ f_{2} \circ ... \circ f_N \\ \end{aligned}\) and so for the Jacobians (using y=g(z)) \(\begin{aligned} Df(y) &= \prod_{i=1}^N D f_i(x_i) \\ \det Df(y) &= \prod_{i=1}^N \det D f_i(x_i) \end{aligned}\) where $x_1 = z$, $x_N=y$, and \(\begin{aligned} x_N &= y \\ x_{N-1} &= f_N(y) \\ x_i &= f_{i+1} \circ ... \circ f_N(y) \\ x_i &= g_i \circ ... \circ g_1(z) \\ x_1 &= g_1(z) \end{aligned}\)

Or for the other direction: \(\begin{aligned} \det Dg(z) &= \prod_{i=1}^N \det D g_i(x_i) \\ x_1 &= z \\ x_2 &= g_1(z) \\ x_i &= g_{i-1} \circ g_{i-2} \circ ... \circ g_1(z) \end{aligned}\)

To the transform, we seek to maximize the log-likelihood of the data given the transform parameters $\theta$ (we fix the base/latent distribution):

\[\begin{aligned} \log \, p(D|\theta) &= \sum_{i=1}^N \, \log \,p_y(y^{(i)}|\theta) \\ &= \sum_{i=1}^N \,\log \, p_z( f(y|\theta)) + \log \,| \text{det} \, D(f(y|\theta))| \\ &= \sum_{i=1}^N \,\log \, p_z( f(y|\theta)) - \log \,| \text{det} \, Dg_\theta\,(f(y|\theta))| \\ \end{aligned}\]

Implementing the $g$’s a layers in a neural network, we’d have something like this to compute the total log_prob:

def log_prob(flow, target_data)
  y = target_data.sample((512,))
  z = flow.inverse(y) # f(y)
  p_z = flow.base_dist
  log_pz = p_z.log_prob(f_y)

  logdetJ = []
  for layer in model.layers:
    logdetJ += layer.log_abs_det_j(z)
    z = layer(z)
    return log_pz - logdetJ

References

• Rezende, D. J., Mohamed, S. (2015). Variational Inference with Normalizing Flows. arXiv preprint arXiv:1410.8516.
• Shen, X., et al. (2019). Normalizing Flows for Probabilistic Modeling and Inference. arXiv:1908.09257.
• Pawlowski, J., Urban, J. (2022). Flow-based density of states for complex actions. arXiv:2203.01243.