Normalizing Flow

Posted on Edited on In

Normalizing flow

Research Gap：

现有研究（GAN，VAE， DRAW等deep generative models）不容易进行测试解析（如概率密度不可得）和度量（如KL距离，earth-mover距离不可得） 。

目标：分布既能表达真实样本又有传统统计学习所具备的的较好的解释性

Guassian distribtion：采样方便，解析密度可知，KL距离容易计算，“中心极限定理”（任何大数据最终区域Guassian distribution），参数化技巧，提督反转等优点。但真实样本的PDF往往与categorical distribution或者Guassian distribution的差距很大，分布不够集中出现边缘效应，不能应对罕见事物等缺点。

要求：

  • 足够复杂，容得下多个模式，比如增强学习中的图像和评分函数；
  • 足够简单，能采样，能估计密度，能重参数化。

Idea：

使用一个简单的概率分布，从它采样，再对samples作变换（等价于改不变概率分布）。如果让这个变换满足某些milde condition，应该可以得到一个关于变换后的随机变量的非常复杂的PDF。–>得到一个可逆，可计算分布变换，易模拟的模型 –> normalizing flow

意义：

  • 把简单分布（如Gaussian分布）的PDF转化成某种复杂分布。这里的flow值的是data glow经过一系列双射（可逆映射），最终可映射到合适的representative space；normalizing值的事representative space的变量积分为1（这也满足PDF的概念）。
  • 可产生复杂的分布函数，用于reinforcement learning、generative model、variational inference等。

Normalizing flow介绍：

Geerative models with tractable explicit density 值的是$p(x)$可以直接计算出来，包括了大部分auto regressive model (NADE, MADE, PixelRNN/CNN)，还包括flow-based model（通过变量替换，用$p(z)$通过变换得到$p(x)$，即$x=T(z)$。

Flow-based model:

  • likelihood based
  • change of variable (采用变量替换的方式解决因变量问题)
  • tractable（可精确infer $p(x;\theta)$

具体推导：

flow-based model：需要多个简单变换的复合来增强拟合能力$T = T_k \circ \cdots \circ T_1$，且简单变换$T_i$的逆与微分都满足可乘性。得到的分布由base distribution和双射函数得到：

  1. base distribution：可以是gaussian 分布, base_dist()
  2. 双射函数：
    1. 前向映射 $x=T(u)$, X=bijector.forward(base_dist.saple())
    2. 反向映射 $u=T^{-1}(x)$, J=new_dist.log_prob(bijector.inverse(x))+bijector.inverse_log_det_jacobian(x), x是样本。
    3. Jacobian矩阵的逆对数行列式 $\log(|\operatorname{det}(J(T^{-1}(x)))|)$. 用于评估转换后分布的对数密度。算法通过最大似然估计，把拟合真实数据的分布问题变成拟合变换后的概率的对数密度问题

在神经网络中，取决于模型任务的不同，部分双射函数可以只实现前向映射$T_k$或后向映射$T_k^{-1}$中的一个。我们用$f_{\phi_k}$表示带有参数$\phi_k$的双射函数层，多个这类简单变换构成的双射函数层复合，成为流模型。

目标函数：$\log(p_X(x_n)) = \log(p_u(T^{-1}(x_n))) + \log(|\operatorname{det}(J(T^{-1}(x_n)))|^{-1})$. 可以用KL divergence和Monte Carlo方法的无偏估计得到：

  1. 正向KL divergence
    $$
    \begin{aligned}
    \mathcal{L}(\boldsymbol{\theta}) &=D_{\mathrm{KL}}\left[p_{\mathrm{x}}^{}(\mathrm{x}) | p_{\mathrm{x}}(\mathrm{x} ; \theta)\right] \
    &=-\mathbb{E}{p{\mathrm{x}}^{    }(\mathrm{x})}\left[\log p_{\mathrm{x}}(\mathrm{x} ; \theta)\right]+\mathrm{const.} \
    &=-\mathbb{E}{p{\mathrm{x}}^{}(\mathrm{x})}\left[\log p_{\mathrm{u}}\left(T^{-1}(\mathrm{x} ; \phi) ; \psi\right)+\log \left|\operatorname{det} J_{T^{-1}}(\mathrm{x} ; \phi)\right|\right]+\mathrm{const} .
    \end{aligned}
    $$
    如果我们拥有真实样本$x_n\sim p^    (x)$，用Monte Carlo简化后，有
    $$
    \mathcal{L}(\boldsymbol{\theta}) \approx-\frac{1}{N} \sum_{n=1}^{N} \log p_{\mathrm{u}}\left(T^{-1}\left(\mathbf{x}{n} ; \boldsymbol{\phi}\right) ; \boldsymbol{\psi}\right)+\log \left|\operatorname{det} J{T^{-1}}\left(\mathbf{x}_{n} ; \boldsymbol{\phi}\right)\right|+\text { const. }
    $$
    当目标分布$p^*(x)$的真实样本可以得到但PDF无法计算出来时，evaluate这个model只需要用到$T^{-1}$和$\operatorname{det}(J(T^{-1}))$，而不用计算$T$与$p(u)$也可以训练model，但是如果我们想在训练好的模型上生成新的样本，$T$与$p(u)$也仍然需要.

    在具体的实践中，我们可以通过采用梯度下降的方法最小化损失函数$\mathcal{L}(\boldsymbol{\theta})$:
    $$
    \begin{array}{l}
    \nabla_{\phi} \mathcal{L}(\theta) \approx-\frac{1}{N} \sum_{n=1}^{N} \nabla_{\phi} \log p_{u}\left(T^{-1}(\mathbf{x} n ; \phi) ; \psi\right)+\nabla \phi \log \left|\operatorname{det} J_{T^{-1}}(\mathbf{x} n ; \phi)\right| \
    \nabla \psi \mathcal{L}(\theta) \approx-\frac{1}{N} \sum_{n=1}^{N} \nabla_{\psi} \log p_{u}\left(T^{-1}\left(\mathbf{x}_{n} ; \phi\right) ; \psi\right)
    \end{array}
    $$

  2. 反向KL divergence

    当我们拥有真实样本$p^*$的分布，但很难得到真实样本分布的采样$x_n$时，可以通过反向KL divergence来评估模型，具体的到做事通过变量替换，对$x$的操作转换成对$u$的操作：
    $$
    \begin{aligned}
    \mathcal{L}(\boldsymbol{\theta}) &=D_{\mathrm{KL}}\left[p_{\mathrm{x}}(\mathbf{x} ; \boldsymbol{\theta}) | p_{\mathrm{x}}^{}(\mathbf{x})\right] \
    &=\mathbb{E}{p{\mathrm{x}}(\mathbf{x} ; \boldsymbol{\theta})}\left[\log p_{\mathrm{x}}(\mathbf{x} ; \boldsymbol{\theta})-\log p_{\mathrm{x}}^{    }(\mathbf{x})\right] \
    &=\mathbb{E}{p{\mathrm{u}}(\mathbf{u} ; \psi)}\left[\log p_{\mathrm{u}}(\mathbf{u} ; \boldsymbol{\psi})-\log \left|\operatorname{det} J_{T}(\mathbf{u} ; \boldsymbol{\phi})\right|-\log p_{\mathrm{x}}^{}(T(\mathbf{u} ; \boldsymbol{\phi}))\right]
    \end{aligned}
    $$
    当我们能够计算$T$，它的雅可比行列式，evaluate 目标分布$p^    $以及从$$p(u)$$中sample时，使用反向KL divergence是合适的，事实上，即使我们只能evaluate 目标分布乘以某个正则化常数，也可以最小化上式, $p^*_{\mathrm x}(\mathrm x) = \tilde p_{\mathrm x}(\mathrm x)/C$，其中，$\tilde p_{\mathrm x}(\mathrm x)$是一个更好处理的分布，因此，可以得到
    $$
    \mathcal{L}(\boldsymbol{\theta})=\mathbb{E}{p{u}(\mathbf{u} ; \psi)}\left[\log p_{\mathrm{u}}(\mathbf{u} ; \boldsymbol{\psi})-\log \left|\operatorname{det} J_{T}(\mathbf{u} ; \boldsymbol{\phi})\right|-\log \tilde{p}{\mathrm{x}}(T(\mathbf{u} ; \boldsymbol{\phi}))\right]+\text { const. }
    $$
    当我们有N条来自于$p    {u}(\mathbf{u} ; \psi)$的samples，为了最小化上式，使用蒙特卡罗法，并对变换$T$的参数$\phi$求偏导，可得
    $$
    \nabla_{\phi} \mathcal{L}(\boldsymbol{\theta}) \approx-\frac{1}{N} \sum_{n=1}^{N} \nabla_{\phi} \log \left|\operatorname{det} J_{T}\left(\mathbf{u}{n} ; \boldsymbol{\phi}\right)\right|+\nabla{\phi} \log \tilde{p}{\mathbf{x}}\left(T\left(\mathbf{u}{n} ; \boldsymbol{\phi}\right)\right)
    $$
    由于$u_n\sim p_u(u)$，且$\phi$不更新因为$u_n$的采样$\phi$是固定的，具体优化如下：
    $$
    \nabla_{\phi} \mathcal{L}(\boldsymbol{\theta}) \approx-\frac{1}{N} \sum_{n=1}^{N} \nabla_{\phi} \log \left|\operatorname{det} J_{T}(\mathbf{u}_n ; \boldsymbol{\phi})\right|+\nabla \boldsymbol{\phi} \log \tilde{p}_\mathbf{x}(T(\mathbf{u}_n ; \boldsymbol{\phi}))
    $$
    反向KL散度，是通过对$p(u)$采样通过变量替换以避开$x$来训练评估我们的模型。其常见的应用有变分推断和模型蒸馏。变分推断指的是，先验概率或似然函数$p(x)$ is intractable ，变分推断想求后验 $p(z|x)$，只能通过构造新的分布去逼近。那么分布是已知的，对隐变量的采样是未知的，符合我们的条件。在模型蒸馏中， $p^*$ 符合同样可求，但采样不方便的境地。因此采用反向KL散度评估模型去训练蒸馏，使新的模型具有快速采样的能力。

如何找到映射函数T既是可逆的而且逆变换计算容易，这个$T$又可以简便计算雅克比行列式？对于逆变换计算这一点，如果无法高效进行逆计算，那么评估训练模型以及采样都是困难的，取决于$f_{\phi_k}$实现的是$T_k$还是$T_k^{-1}$，因此逆变换的高效计算是个因模型而异的难题。

basic trick: 计算雅可比行列式最好优化成三角矩阵 diagonal matrix，because the determinant of a triangular matrix is just the product of the diagonal elements, 因此将原先的$O(D^3)$复杂度降为$O(D)$，能够解决部分计算问题！

与VAE和GAN的比较

normalizing flow：

  • exact log-likelihood evaluation: $\log p_\theta(x) = \log p_\theta(z) + \sum_{i=1}^K \log \left|\operatorname{det}\left( \frac{\part f_i*{-1}}{\part z_i} \right)\right|$
  • exact posterior inference : $z=f^{-1}(x)$

VAE:

  • Lower bound on log-likelihood (ELBO)
  • approximate posterior $q_\phi(z|x)$ (encoder network) [because the actural posterior is intractable]

GAN:

  • No log-likelihood evaluation (optimize a minx objective involving the discriminator and generator)
  • no latent variable inference (generator cannot ensure full support over data)

模型1-Finite compositions

Finite compositions model的优点在于，增加简单变换$T$的个数为$K$，仅仅增加$O(K)$的计算复杂度。

Finite compositions model: construct a flow with transformation $T$ by composing a finite number of simple transformations $T_k$. 有限个简单变换$T$构成的flow mode

  1. Autgressive flows
  2. Linear flows
  3. Residual flows

耦合技术 Coupling layers

NICE (Dihn et al. 2015), Real NVP (Dihn et al. 2017), Glow (Kingma and Dhariwal, 2018), WaveGlow (Prenger et al. 2019), FloWaveNet (kim et al. 2019), Flow++ (Ho et al. 2019).

优点：

  • 采用了计算对称性，模型评估和求逆计算的速度快（efficiently）
  • 模型评估和采样可以在一次神经网络传递中进行（allow both density evaluation and sampling to be performed in a single neural network pass）

缺点：

  • 流模型的表达力可能不足。因此需要考虑计算和表达力的动态平衡。

自回归技术

应用

概率建模

样本生成

模型推断

表示学习（构建产生式模型和预测模型的混合模型，如监督学习和强化学习）

概率模型的建模，推断，监督学习和强化学习这些领域。