Flow-based generative model

A flow-based generative model is a generative model used in machine learning that explicitly models a probability distribution by leveraging normalizing flow,^[1] which is a statistical method using the change-of-variable law of probabilities to transform a simple distribution into a complex one.

The direct modeling of likelihood provides many advantages. For example, the negative log-likelihood can be directly computed and minimized as the loss function. Additionally, novel samples can be generated by sampling from the initial distribution, and applying the flow transformation.

In contrast, many alternative generative modeling methods such as variational autoencoder (VAE) and generative adversarial network do not explicitly represent the likelihood function.

Method

 

Scheme for normalizing flows

Let ${\displaystyle z_{0}}$  be a (possibly multivariate) random variable with distribution ${\displaystyle p_{0}(z_{0})}$ .

For ${\displaystyle i=1,...,K}$ , let ${\displaystyle z_{i}=f_{i}(z_{i-1})}$  be a sequence of random variables transformed from ${\displaystyle z_{0}}$ . The functions ${\displaystyle f_{1},...,f_{K}}$  should be invertible, i.e. the inverse function ${\displaystyle f_{i}^{-1}}$  exists. The final output ${\displaystyle z_{K}}$  models the target distribution.

The log likelihood of ${\displaystyle z_{K}}$  is (see derivation):

${\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|}$ 

To efficiently compute the log likelihood, the functions ${\displaystyle f_{1},...,f_{K}}$  should be 1. easy to invert, and 2. easy to compute the determinant of its Jacobian. In practice, the functions ${\displaystyle f_{1},...,f_{K}}$  are modeled using deep neural networks, and are trained to minimize the negative log-likelihood of data samples from the target distribution. These architectures are usually designed such that only the forward pass of the neural network is required in both the inverse and the Jacobian determinant calculations. Examples of such architectures include NICE,^[2] RealNVP,^[3] and Glow.^[4]

Derivation of log likelihood

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|}$ 

Training method

Flow-based models are generally trained by maximum likelihood. A pseudocode is as follows:^[5]

• INPUT. dataset ${\displaystyle x_{1:n}}$ , normalizing flow model ${\displaystyle f_{\theta }(\cdot ),p_{0}}$ .
• SOLVE. ${\displaystyle \max _{\theta }\sum _{j}\ln p_{\theta }(x_{j})}$  by gradient descent
• RETURN. ${\displaystyle {\hat {\theta }}}$ 

Variants

Planar Flow

The earliest example.^[1] 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.

Nonlinear Independent Components Estimation (NICE)

Let ${\displaystyle x,z\in \mathbb {R} ^{2n}}$  be even-dimensional, and split them in the middle.^[2] 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.

Real Non-Volume Preserving (Real NVP)

^[3]$${\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,})}}$ .

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.

Generative Flow (Glow)

^[4] Each layer of Glow 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 an 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.

masked autoregressive flow (MAF)

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

$${\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 (it's sequential), but the backward mapping is fast (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)^[7].

Continuous Normalizing Flow (CNF)

Instead of constructing flow by function composition, another approach is to formulate the flow as a continuous-time dynamic.^[8] 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:^[8]

${\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:^[8]

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

The trace can be estimated by "Hutchinson's trick"^[9]:

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).

Because of the use of integration, techniques such as Neural ODE^[10] may be needed in practice.

Applications

Flow-based generative models have been applied on a variety of modeling tasks, including:

• Audio generation^[11]
• Image generation^[4]
• Molecular graph generation^[12]
• Point-cloud modeling^[13]
• Video generation^[14]

References

1. Danilo Jimenez Rezende; Mohamed, Shakir (2015). "Variational Inference with Normalizing Flows". arXiv:1505.05770 [stat.ML].
2. Dinh, Laurent; Krueger, David; Bengio, Yoshua (2014). "NICE: Non-linear Independent Components Estimation". arXiv:1410.8516 [cs.LG].
3. Dinh, Laurent; Sohl-Dickstein, Jascha; Bengio, Samy (2016). "Density estimation using Real NVP". arXiv:1605.08803 [cs.LG].
4. Kingma, Diederik P.; Dhariwal, Prafulla (2018). "Glow: Generative Flow with Invertible 1x1 Convolutions". arXiv:1807.03039 [stat.ML].
5. Kobyzev, Ivan; Prince, Simon J.D.; Brubaker, Marcus A. (2021-11). "Normalizing Flows: An Introduction and Review of Current Methods". IEEE Transactions on Pattern Analysis and Machine Intelligence. 43 (11): 3964–3979. doi:10.1109/TPAMI.2020.2992934. ISSN 1939-3539. {{cite journal}}: Check date values in: |date= (help)
6. Papamakarios, George; Pavlakou, Theo; Murray, Iain (2017). "Masked Autoregressive Flow for Density Estimation". Advances in Neural Information Processing Systems. 30. Curran Associates, Inc.
7. Kingma, Durk P; Salimans, Tim; Jozefowicz, Rafal; Chen, Xi; Sutskever, Ilya; Welling, Max (2016). "Improved Variational Inference with Inverse Autoregressive Flow". Advances in Neural Information Processing Systems. 29. Curran Associates, Inc.
8. Grathwohl, Will; Chen, Ricky T. Q.; Bettencourt, Jesse; Sutskever, Ilya; Duvenaud, David (2018). "FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models". arXiv:1810.01367 [cs.LG].
9. Hutchinson, M.F. (1989-01). "A Stochastic Estimator of the Trace of the Influence Matrix for Laplacian Smoothing Splines". Communications in Statistics - Simulation and Computation. 18 (3): 1059–1076. doi:10.1080/03610918908812806. ISSN 0361-0918. {{cite journal}}: Check date values in: |date= (help)
10. Chen, Ricky T. Q.; Rubanova, Yulia; Bettencourt, Jesse; Duvenaud, David (2018). "Neural Ordinary Differential Equations". arXiv:1806.07366 [cs.LG].
11. Ping, Wei; Peng, Kainan; Zhao, Kexin; Song, Zhao (2019). "WaveFlow: A Compact Flow-based Model for Raw Audio". arXiv:1912.01219 [cs.SD].
12. Shi, Chence; Xu, Minkai; Zhu, Zhaocheng; Zhang, Weinan; Zhang, Ming; Tang, Jian (2020). "GraphAF: A Flow-based Autoregressive Model for Molecular Graph Generation". arXiv:2001.09382 [cs.LG].
13. Yang, Guandao; Huang, Xun; Hao, Zekun; Liu, Ming-Yu; Belongie, Serge; Hariharan, Bharath (2019). "PointFlow: 3D Point Cloud Generation with Continuous Normalizing Flows". arXiv:1906.12320 [cs.CV].
14. Kumar, Manoj; Babaeizadeh, Mohammad; Erhan, Dumitru; Finn, Chelsea; Levine, Sergey; Dinh, Laurent; Kingma, Durk (2019). "VideoFlow: A Conditional Flow-Based Model for Stochastic Video Generation". arXiv:1903.01434 [cs.CV].

External links

• Flow-based Deep Generative Models
• Normalizing flow models