miaom Deep Learning 学习笔记 05 – Normalizing Flows
Contents
上节课介绍了基于能量的模型(Energy-Based Model)，对于 $x$ 的概率 $P(x) = \frac1Z \exp(E(x;\theta))$。我们可以通过最大似然训练(Maximum Likelihood Training
)训练这个网络。但是我们并没有一个很好的办法来从模型中采样。
这一讲，我们需要解决从概率密度 $p(x)$ 中采样的问题。
一般的采样方法 (Sampling Methods)
问题：我们有一个任意的概率密度函数 $p(x)$，我们要从中采样。
对于能算累积分布函数(CDF)的分布，比如分类分布(Categorical distribution)。假设累计分布函数为 $F(x)$，我们可以先从均匀分布中采样 $a$，$F^{-1}(x)$ 就是符合概率密度函数的随机变量。
对于正态分布(Gaussian distribution)，累计分布函数没有闭形式。我们可以根据中心极限定理，通过足够多的伯努利分布来趋近正态分布。一种更好的办法是用Box–Muller变换，利用二维正态分布的模长的累积分布来采样。
重要性采样(Importance Sampling)
对于没有什么好的数学形式的 $p(x)$，我们可以用重要性采样。具体地讲，我们需要找一个提议分布(proposal distribution) $q(x)$，将 $\frac{p(x)}{q(x)}$设为权值。我们有 $\mathbb E_{x\sim p}[f(x)] = \mathbb E_{x\sim q} [\frac{p(x)}{q(x)}f(x)]$。
这样一来，如果我们找的这个 $q(x)$ 很容易采样，那我们就能做到容易地从 $p(x)$ 中采样了。
在基于能量的模型中 $P(x) = \frac1Z \exp(E(x;\theta))$ 的 $Z$ 难以计算。使用重要性采样可以相当于给所有权值乘上一个 $Z$ 来避免计算 $Z$。
但是这里需要注意，尽管这两个东西的期望一样，但是实际采样的方差并没有限制。举个例子，$p(x)$ 的概率集中在 $[0,10^{-10}]$，然后 $q(x)$ 是 $[0,1]$ 的均匀分布，那么我 至少得有 $10^{10}$ 级别次的采样，才期望能够得到一个比较大的值。
简单地讲，$q$ 和 $p$ 比较接近的话，方差就小；$q$ 和 $p$ 差得远的话方差就大。我们需要让$q$ 和 $p$ 比较接近。方差大就需要更多的采样次数。尤其是面对高维的随机变量时，我们无法接受如此巨大的采样次数。（k维的情况，采样次数大概可以是 $\exp(k)$ 的级别。）
马尔可夫链蒙特卡洛方法(Markov Chain Monte-Carlo)
马尔可夫链包括一个状态空间 $S$，以及一个转移概率 $P(s_j|s_i)=T_{ij}$。我们称 $T$ 为转移矩阵。
对于合适的 $T$，马尔可夫链存在稳定分布(stationary distribution) $\pi = T^{\infty}\pi_0$，其中 $\pi_0$ 是初始分布。
存在稳定分布的一个充分条件是Detailed Balance(细致平衡)，即 $\pi(s)T(s\rightarrow s')=\pi(s')T(s'\rightarrow s)$。想要保证任意的初始分布都能达到同一个稳定分布，我们还需要遍历性(ergotic)，即 $\min_z\min_{z':\pi(z')>0} \frac{T(z\rightarrow z')}{\pi(z')}>0$。
如果我们可以对构造一个马尔可夫链的稳定分布等于我们的概率密度，那么我们就能用MCMC来采样了。
Metropolis-Hastings Algorithm
Metropolis-Hastings算法构造如下：首先有一个proposal distribution $q(s'|s)$，对于每次从 $s$ 出发的转移，首先根据 $q(s'|s)$ 采样 $s'$，下一个状态有 $\alpha=\min(1,\frac{p(s')q(s'\rightarrow s)}{p(s)q(s\rightarrow s')})$ 的概率是 $s'$，否则是仍然是 $s$。其中，proposal distribution $q$ 的一个简单的选择是 $s' := s + \text{random noise}$。
Metropolis-Hastings算法满足Detailed Balance和ergotic，$\pi = p$。所以它可以最终收敛到我们想要的概率分布。
Metropolis-Hastings算法是一个非常一般性的构造马尔科夫链的算法，但是当 $s$ 的维度增加，会出现Curse of dimensionality。由于我们的 $q$ 总是没那么好，接受 $s'$ 的概率 $\alpha$ 可以变得非常小。这使得Metropolis-Hastings算法需要更多的迭代。
Gibbs Sampling
另一种常用的构造马尔科夫链的算法是Gibbs Sampling。
对于一个 $N$ 维向量 $s$，Gibbs Sampling 每次随机取一个维度 $i$，然后从分布 $P(s_i|s_1,s_2,...,s_{i-1},s_{i+1},...,s_{N})$ 中采样 $s_i'$。新的向量 $s' = (s_1,s_2,...,s_{i-1},s'_i,s_{i+1},...,s_N)$。
实践中，我们希望 $P(s_i|s_1,s_2,...,s_{i-1},s_{i+1},...,s_{N})$ 比较容易采样，比如Exponential family。我们也可以用神经网络学这个分布。
Gibbs Sampling 不需要复杂的 proposal distribution，但是仍然需要很长时间才能收敛。回顾上节课受限玻尔兹曼机就用了特殊限制下改进后的Gibbs Sampling。
MCMC with Langevin dynamics
我们还可以用类似SGD的方法来采样。设我们有 $X=(x_1,...,x_N)$，要采样的向量是 $\theta$，每一轮转移
$$
\theta' \leftarrow \theta + \frac\epsilon2 \left( \sum_i \nabla_\theta \log P(x_i|\theta) \right) + \mathcal N(0,\epsilon I)。
$$
当 $\epsilon$ 退火到 $0$ 时，$\theta$ 会收敛到分布 $P(\theta|X)$。实际操作可以调整梯度和高斯噪声前面的系数。
Non-Sampling Methods
除了MCMC这样需要迭代采样的方法之外，还有一些别的采样的方法。比如Approximate Bayesian Inference 和 Scoring Matching。
Normalizing Flow（标准化流）
上节课Energy-Based Model是让神经网络学习一个能量函数 $E(x)$。相应的概率密度函数 $p(x) = \frac1Z \exp(E(x))$。由于需要MCMC，Energy-Based Model很难采样。而训练时要用到MLE loss。我们没有很好的办法精确地算出归一化系数 $Z$，于是也只能用采样的方式近似。
那么有没有一个模型容易采样并且能算准确的MLE loss呢。于是我们就有了Normalizing Flow。
Intuition
我们希望找一个 $y=f(x:\theta)$，其中 $x$ 是合于某个分布（比如 $k$ 维高斯分布）的隐变量，$y$ 是生成的东西，$\theta$ 是模型的参数。我们希望 $f(x:\theta)$ 好算，并且 $y$ 的概率密度好算。
举个例子，假设 $x\in [0,1]$ 的概率密度为 $p(x)$，$y=ax+b$。可以算出 $y$ 的概率密度 $p(y) = \frac1a p(x)$。
一般地，如果 $y=f(x)$，$f$ 是可逆的，那么 $p(y) = p(x) (f'(x))^{-1}$。
对于高维的 $x,y$，我们也能得到类似的结论。对于 $y=Ax + b$，$A$ 可逆，$p(y)=|\det A|^{-1} p(x)$。
Normalizing Flow的思想就是先从一个比较简单的分布（比如高斯分布）中采样 $z_0$，然后过 $K$ 个可逆的函数 $z_i = f_i(z_{i-1})$，最后输出 $x = z_K$ 的概率密度函数的可算的。这里的可逆函数包括线性的函数和各种激活函数。
这样，我们想要生成图像只需要算 $f_K(...f_1(z_0))$。训练时数据 $x$ 的Log-Likelihood是 $\log p(z_0) - \sum_i \log |\det \frac{\partial f_i}{\partial z_{i-1}}|$。
随机矩阵几乎可逆。剩下的问题是怎么算这个行列式。
Planar Flow (Rezende & Mohamed, 2016)
暴算一个 $d$ 维矩阵的行列式的复杂度是 $O(d^3)$，无法接受。我们需要对这个函数做一些限制。
一个解决方法是让 $f_\theta(z) = z + u \tanh(w^Tz+b)$，其中 $\theta = [u,w,b]$，$u,w\in \mathbb R^d$，$b\in \mathbb R$。
注意到 $\det (A + uv^T) = (1+v^TA^{-1}u)\det A$，于是
$$
\det \frac{\partial f}{\partial z} = \det (I + h'(w^Tz+b)uW^T) = 1 + h'(w^Tz+b)u^TW。
$$
这样只需要 $O(d)$ 时间就可以算出行列式了。实现时为了确保函数可逆，$u^Tw>-1$。
另一种处理这个行列式的办法是，直接让 $J = \frac{\partial f}{\partial z_{i-1}}$ 是下三角矩阵。
NICE (Nonlinear Independent Components Estimation)
我们这样构造可逆函数：假设输入是 $z$，$z_{1:m}$ 是 $z$ 中某 $m$ 维，$z_{m+1:d}$ 是另外 $d-m$ 维。函数的输出是 $x$。
- $x_{1:m} = z_{1:m}$。
- $x_{m+1:d} = z_{m+1:d} - \mu_\theta (x_{1:m})$。
这里的 $\mu_\theta$ 是一个神经网络。
反过来也很简单：$z_{1:m} = x_{1:m}$，$z_{m+1:d} = x_{m+1:d} + \mu_\theta (x_{1:m})$。
这个函数的雅可比行列式 $\det J = 1$。
我们用好几层这样的可逆的函数，在不同层选取不同的 $m$ 维 或者randomly shuffle来保证每个位置都经过了变换。NICE的最后一层用了 re-scaling layer，$x_i=s_iz_i$，雅可比行列式为 $\prod_i s_i$。
NICE除了能够生成，也可以用来inpainting（图像修补）。对于图片 $x = (x_v,x_h)$，$x_v$ 固定，我们对 $x_h$ 梯度上升，最大化概率密度 $p(x)$。也可以用 stochastic gradient MCMC 对这个分布采样。
Real-NVP
注意到我们其实没有必要严格限制雅可比行列式为 $1$。稍加更改：
- $x_{1:m} = z_{1:m}$。
- $x_{m+1:d} = (z_{m+1:d} - \mu_\theta (x_{1:m})) \odot \exp(\alpha_\theta(x_{1:m}))$。
这里的 $\mu_\theta$ 和 $\alpha_\theta$ 都是神经网络。$\log \det J = \sum_{i=m+1}^d \alpha_\theta(x_{1:m})_i$。
GLOW
GLOW 的核心思想是对于Channel较少的 $z$，使用计算量较小的1x1卷积层。实现细节见 Kingma et al. 2018。
Autoregressive Flow (AF)
其实我们只需要找一个顺序一个一个生成就可逆了，类似sequence model。
首先采样 $z_i\sim \mathcal N(0,1)$，$i=1,2,...,d$。然后
- $x_1 = \exp(\alpha_1)z_1+\mu_1$
- $x_2 = \exp(\alpha_2(x_1))z_2+\mu_2(x_1)$
- $x_3 = \exp(\alpha_3(x_1,x_2))z_3+\mu_3(x_1,x_2)$
- ......
这东西从 $z$ 算 $x$ 需要 $O(d)$ 的时间，但是从 $x$ 算概率密度时有，$z_i = (x_i-\mu_i) / \exp(\alpha_i)$，可以并行。
Autoregressive Flow的一个扩展是Masked Autoregressive Flow (MAF)，我们可以把AF的结构堆很多层，把上一层的 $x$ 作为下一层的 $z$。为了让正向采样的速度更快，可以使用MADE(Masked Autoencoder for Distribution Estimation)的结构。
Quantization
一个问题，我们训练时候的数据，比如图片，每个像素的值是 $0$ 到 $255$ 的整数。感觉上，如果我们的模型非常完美，我们最后学出来的概率密度就会聚集在这些整点，在 $p(x_i=1.5)$ 这种位置会比较低。这问题很大。
事实上，我们用MLE训练的时候，假设训练集的某个数据是 $x$，最大化的对象不应该是 $p(x)$，而应该是 $\int_{[0,1)^d}p(x+u)du$。
一个简单的解决方法是数据增强的时候给 $x$ 加一个噪声 $u\sim Uni([0,1)^d)$。我们有
$$
\begin{aligned}
\mathbb E_{x'}[\log p(x')]
=& \sum_{x\sim \text{data}} \int_{[0,1)^d} \log p(x+u)du\\
\le& \sum_{x\sim \text{data}} \log \int_{[0,1)^d} p(x+u)du\\
=& \mathbb E_x [\log \int_{[0,1)^d} p(x+u)du]。
\end{aligned}
$$
事实上我们最优化的是准确的MLE的一个下界。下节课VAE中会更多地涉及这样的东西。
WaveNet (DeepMind, 2016)
Autoregressive Flow的一个具体的例子是WaveNet。
WaveNet的目标是声音合成。概率密度函数为 $p(x) = \prod_i p(x_i|x_{< i})$。WaveNet使用时域卷积来完成这个任务。具体的讲，我们的输入和输出都是长度维 $n$ 的1维向量，其中 $n$ 是声音的长度。用时域卷积的一个问题在于，想要让一个位置受到之前所有位置的影响，简单的卷积需要 $O(n)$ 的层数。WaveNet使用Dilated Temporal Convolution很好的解决了这个问题。第 $i$ 层第 $j$ 个节点的输入是第 $i-1$ 层第 $j-2^i$ 个和第 $j$ 个输出。我们只需要 $O(\log n)$ 层就足够了。WaveNet还使用了Quantization、Residual connection、Conditioned generation，它是第一个生成原始信号(raw signals)的深度生成模型。
PixelCNN (DeepMind, 2016)
Autoregressive Flow也可以用来生成图片，比如PixelCNN。我们可以按从上到下，从左到右的顺序生成 $N^2$ 个位置。
为了让每个像素只依靠之前的像素的值，我们需要用到Masked Convolution。以5x5卷积为例，我们的mask为 $\begin{pmatrix}1&1&1&1&1\\1&1&1&1&1\\1&1&0&0&0\\0&0&0&0&0\\0&0&0&0&0\end{pmatrix}$。
多层Masked Convolution之后，每个像素实际上并没有和所有 $z$ 有关。论文中使用了 Gated convolution等来解决这个问题。
Inverse Autoregressive Flow (IAF)
之前也提到，尽管AF计算反向计算likelihood可以并行，当它正向从 $z$ 生成 $x$ 时需要消耗 $O(d)$ 的时间。所以AF训练很快，采样比较慢。
注意到AF中是 $x_i=\exp(\alpha_{x_{< i}}) _ \mu(x_{< i})$，我们把它改成 $x_i=\exp(\alpha_{z_{< i}}) _ \mu(z_{< i})$。这就是Inverse Autoregressive Flow (IAF)。IAF的采样可以并行很快，但是计算likelihood无法并行很慢。那么我们可以让采样和训练都很快吗？
Parallel WaveNet (DeepMind, ICML 2018)
我们可以用teacher student框架，先用之前的方法训练一个AF作为老师。然后训练一个IAF最为学生，去拟合老师。
训练学生的时候，因为我们已经有了老师，我们可以随机采样 $z$ 输入学生，得到 $x$。然后把 $x$ 输入老师，得到 $z'$。用两个网络的 $x$ 的概率密度的KL散度作为Loss。
下节课VAE中会对这样的结构做更多讨论。
No Comments