Flow-based generative model

Consider ${\displaystyle z_{1}}$ and ${\displaystyle z_{0}}$. Note that ${\displaystyle z_{0}=f_{1}^{-1}(z_{1})}$.

By the change of variable formula, the distribution of ${\displaystyle z_{1}}$ is:

${\displaystyle p_{1}(z_{1})=p_{0}(z_{0})\left|\det {\frac {df_{1}^{-1}(z_{1})}{dz_{1}}}\right|}$

Where ${\displaystyle \det {\frac {df_{1}^{-1}(z_{1})}{dz_{1}}}}$ is the determinant of the Jacobian matrix of ${\displaystyle f_{1}^{-1}}$.

By the inverse function theorem:

${\displaystyle p_{1}(z_{1})=p_{0}(z_{0})\left|\det \left({\frac {df_{1}(z_{0})}{dz_{0}}}\right)^{-1}\right|}$

By the identity ${\displaystyle \det(A^{-1})=\det(A)^{-1}}$ (where ${\displaystyle A}$ is an invertible matrix), we have:

${\displaystyle p_{1}(z_{1})=p_{0}(z_{0})\left|\det {\frac {df_{1}(z_{0})}{dz_{0}}}\right|^{-1}}$

The log likelihood is thus:

${\displaystyle \log p_{1}(z_{1})=\log p_{0}(z_{0})-\log \left|\det {\frac {df_{1}(z_{0})}{dz_{0}}}\right|}$

In general, the above applies to any ${\displaystyle z_{i}}$ and ${\displaystyle z_{i-1}}$. Since ${\displaystyle \log p_{i}(z_{i})}$ is equal to ${\displaystyle \log p_{i-1}(z_{i-1})}$ subtracted by a non-recursive term, we can infer by induction that:

${\displaystyle \log p_{K}(z_{K})=\log p_{0}(z_{0})-\sum _{i=1}^{K}\log \left|\det {\frac {df_{i}(z_{i-1})}{dz_{i-1}}}\right|}$

The earliest example.[8] Fix some activation function ${\displaystyle h}$, and let ${\displaystyle \theta =(u,w,b)}$ with th appropriate dimensions, then

$${\displaystyle x=f_{\theta }(z)=z+uh(\langle w,z\rangle +b)}$$

The inverse ${\displaystyle f_{\theta }^{-1}}$ has no closed-form solution in general.

The Jacobian is ${\displaystyle |\det(I+h'(\langle w,z\rangle +b)uw^{T})|=|1+h'(\langle w,z\rangle +b)\langle u,w\rangle |}$.

For it to be invertible everywhere, it must be nonzero everywhere. For example, ${\displaystyle h=\tanh }$ and ${\displaystyle \langle u,w\rangle >-1}$ satisfies the requirement.

Let ${\displaystyle x,z\in \mathbb {R} ^{2n}}$ be even-dimensional, and split them in the middle.[3] Then the normalizing flow functions are

$${\displaystyle x={\begin{bmatrix}x_{1}\\x_{2}\end{bmatrix}}=f_{\theta }(z)={\begin{bmatrix}z_{1}\\z_{2}\end{bmatrix}}+{\begin{bmatrix}0\\m_{\theta }(z_{1})\end{bmatrix}}}$$

where ${\displaystyle m_{\theta }}$ is any neural network with weights ${\displaystyle \theta }$.

${\displaystyle f_{\theta }^{-1}}$ is just ${\displaystyle z_{1}=x_{1},z_{2}=x_{2}-m_{\theta }(x_{1})}$, and the Jacobian is just 1, that is, the flow is volume-preserving.

When ${\displaystyle n=1}$, this is seen as a curvy shearing along the ${\displaystyle x_{2}}$ direction.

The Real Non-Volume Preserving model generalizes NICE model by:[4]

$${\displaystyle x={\begin{bmatrix}x_{1}\\x_{2}\end{bmatrix}}=f_{\theta }(z)={\begin{bmatrix}z_{1}\\e^{s_{\theta }(z_{1})}\odot z_{2}\end{bmatrix}}+{\begin{bmatrix}0\\m_{\theta }(z_{1})\end{bmatrix}}}$$

Its inverse is ${\displaystyle z_{1}=x_{1},z_{2}=e^{-s_{\theta }(x_{1})}\odot (x_{2}-m_{\theta }(x_{1}))}$, and its Jacobian is ${\displaystyle \prod _{i=1}^{n}e^{s_{\theta }(z_{1,})}}$. The NICE model is recovered by setting ${\displaystyle s_{\theta }=0}$. Since the Real NVP map keeps the first and second halves of the vector ${\displaystyle x}$ separate, it's usually required to add a permutation ${\displaystyle (x_{1},x_{2})\mapsto (x_{2},x_{1})}$ after every Real NVP layer.

In generative flow model,[5] each layer has 3 parts:

• channel-wise affine transform
  $${\displaystyle y_{cij}=s_{c}(x_{cij}+b_{c})}$$

  
  with Jacobian ${\displaystyle \prod _{c}s_{c}^{HW}}$.
• invertible 1x1 convolution
  $${\displaystyle z_{cij}=\sum _{c'}K_{cc'}y_{cij}}$$

  
  with Jacobian ${\displaystyle \det(K)^{HW}}$. Here ${\displaystyle K}$ 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.

An autoregressive model of a distribution on ${\displaystyle \mathbb {R} ^{n}}$ is defined as the following stochastic process:[9]

$${\displaystyle {\begin{aligned}x_{1}\sim &N(\mu _{1},\sigma _{1}^{2})\\x_{2}\sim &N(\mu _{2}(x_{1}),\sigma _{2}(x_{1})^{2})\\&\cdots \\x_{n}\sim &N(\mu _{n}(x_{1:n-1}),\sigma _{n}(x_{1:n-1})^{2})\\\end{aligned}}}$$

where ${\displaystyle \mu _{i}:\mathbb {R} ^{i-1}\to \mathbb {R} }$ and ${\displaystyle \sigma _{i}:\mathbb {R} ^{i-1}\to (0,\infty )}$ are fixed functions that define the autoregressive model. By the reparametrization trick, the autoregressive model is generalized to a normalizing flow:

$${\displaystyle {\begin{aligned}x_{1}=&\mu _{1}+\sigma _{1}z_{1}\\x_{2}=&\mu _{2}(x_{1})+\sigma _{2}(x_{1})z_{2}\\&\cdots \\x_{n}=&\mu _{n}(x_{1:n-1})+\sigma _{n}(x_{1:n-1})z_{n}\\\end{aligned}}}$$

The autoregressive model is recovered by setting ${\displaystyle z\sim N(0,I_{n})}$.

The forward mapping is slow (because it's sequential), but the backward mapping is fast (because it's parallel).

The Jacobian matrix is lower-diagonal, so the Jacobian is ${\displaystyle \sigma _{1}\sigma _{2}(x_{1})\cdots \sigma _{n}(x_{1:n-1})}$.

Reversing the two maps ${\displaystyle f_{\theta }}$ and ${\displaystyle f_{\theta }^{-1}}$ of MAF results in Inverse Autoregressive Flow (IAF), which has fast forward mapping and slow backward mapping.[10]

Instead of constructing flow by function composition, another approach is to formulate the flow as a continuous-time dynamic.[11][12] Let ${\displaystyle z_{0}}$ be the latent variable with distribution ${\displaystyle p(z_{0})}$. Map this latent variable to data space with the following flow function:

${\displaystyle x=F(z_{0})=z_{T}=z_{0}+\int _{0}^{T}f(z_{t},t)dt}$

where ${\displaystyle f}$ is an arbitrary function and can be modeled with e.g. neural networks.

The inverse function is then naturally:[11]

${\displaystyle z_{0}=F^{-1}(x)=z_{T}+\int _{T}^{0}-f(z_{t},t)dt}$

And the log-likelihood of ${\displaystyle x}$ can be found as:[11]

${\displaystyle \log(p(x))=\log(p(z_{0}))-\int _{0}^{T}{\text{Tr}}\left[{\frac {\partial f}{\partial z_{t}}}\right]dt}$

Since the trace depends only on the diagonal of the Jacobian ${\displaystyle \partial _{z_{t}}f}$, this allows "free-form" Jacobian.[13] 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":[14][15]

Given any matrix ${\displaystyle W\in \mathbb {R} ^{n\times n}}$, and any random ${\displaystyle u\in \mathbb {R} ^{n}}$ with ${\displaystyle E[uu^{T}]=I}$, we have ${\displaystyle E[u^{T}Wu]=tr(W)}$. (Proof: expand the expectation directly.)

Usually, the random vector is sampled from ${\displaystyle N(0,I)}$ (normal distribution) or ${\displaystyle \{\pm n^{-1/2}\}^{n}}$ (Radamacher distribution).

When ${\displaystyle f}$ is implemented as a neural network, neural ODE methods[16] 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 (for example, it's impossible to flip a left-hand to a right-hand by continuous deforming of space, and it's impossible to turn a sphere inside out, or undo a knot), and the other is that the learned flow ${\displaystyle f}$ might be ill-behaved, due to degeneracy (that is, there are an infinite number of possible ${\displaystyle f}$ that all solve the same problem).

By adding extra dimensions, the CNF gains enough freedom to reverse orientation and go beyond ambient isotopy (just like how one can pick up a polygon from a desk and flip it around in 3-space, or unknot a knot in 4-space), yielding the "augmented neural ODE".[17]

Any homeomorphism of ${\displaystyle \mathbb {R} ^{n}}$ can be approximated by a neural ODE operating on ${\displaystyle \mathbb {R} ^{2n+1}}$, proved by combining Whitney embedding theorem for manifolds and the universal approximation theorem for neural networks.[18]

To regularize the flow ${\displaystyle f}$, one can impose regularization losses. The paper [14] proposed the following regularization loss based on optimal transport theory:

$${\displaystyle \lambda _{K}\int _{0}^{T}\left\|f(z_{t},t)\right\|^{2}dt+\lambda _{J}\int _{0}^{T}\left\|\nabla _{z}f(z_{t},t)\right\|_{F}^{2}dt}$$

where ${\displaystyle \lambda _{K},\lambda _{J}>0}$ 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 (not "bumpy") over space and time.