标准化流(Normalizing Flow)（二）：现代标准化流技术

翻译自https://blog.evjang.com/2018/01/nf2.html
原作者：Eric Jang
译者：尹肖贻

0. 交代故事

我在下面的教程里教你干一件很酷的事儿：

你将学到使用MAF、IAF和Real-NVP等标准化流来变形二维高斯分布，而形成复杂的图形“SIGGRAPH”。就像……拉扯橡皮泥

在教程的上半部分，你学到使用标准化流，把高斯分布这类地摊货分布，“变形”为高大上的概率分布。你还亲自使用PReLU搭建了小型可逆神经网络，实现了一个简单的链式二维仿射双射函数。
不过，上次的全连接网络只用了两层隐含层，搭建的流弱爆了。更糟的是，非线性激活函数是单调的、分段线性的。这样一来，网络只能在原点附近，稍微地扭捏数据分布的形状【译者按：1）原文用的是manifold流形，对于非从业者来说，只需知道数据在空间中形成一定形状就可以了2）因为激活函数在原点附近是非线性的】。这个流想要实现更高级的变换，就完全不给力。比如把各向同性的高斯变成下面双模式的“双月”数据。

幸好最近的研究发明了几个更强大的标准化流。下面我们就去探索其中的几个技术。

1. 标准化流之自回归（autoregressive）模型和MAF

自回归概率密度估计技术，比如 WaveNet和PixelRNN，是用来学习联合概率密度 $p(x_{1:D})$ 的。该技术将联合概率密度，分解为以为条件概率密度的乘积，这里的 $x_i$ 依赖前 $i-1$ 个数据的取值：
$p(x) = \prod_i{p(x_i \,|\, x_{1:i-1})}$

1.1 学习阶段

上式中的条件概率密度大都有可学习的参数。例如，常见的选择是单变量高斯分布，自回归的概率密度 $p(x_{1:D})$ 就有两个参数，均值和标准差。这两个参数取决于先前的变量 $x_{1:i-1}$ 。
$p(x_i \,|\, x_{1:i-1}) = \mathcal{N}(x_i \,|\,\mu_i, (\exp\alpha_i)^2)$
$\mu_i = f_{\mu_i}(x_{1:i-1})$
$\alpha_i = f_{\alpha_i}(x_{1:i-1})$
这种方法基于一个简单粗暴的假设：每个变量都依赖于先于其出现的变量，而不是后来的变量。拍脑门就能想明白这不（全）是真的：顶上的像素和底下的像素能有啥因果关系？不过（令许多研究者惊讶的是），这东西产生出的图像效果居然还不错!

1.2 采样阶段

采样阶段，我们从标准高斯 $N(0,1)$ 采D个"噪声变量" $u_{1:D}$ ，进而依次地生成所有样本：
$x_i = u_i\exp{\alpha_i} + \mu_i$
$u_i \sim \mathcal{N}(0, 1)$
自回归采样技术可以这么理解，把（从标准高斯分布 $\mathcal{N}(0, \mathbb{I})$ ）采样得到的噪声变量转换为一个新的分布形式，或者说新的分布形式是标准高斯分布的转换分布（TransformedDistribution）。

1.3 搭建一个流

有了这个认识，我们可以堆叠多个自回归模型，搭建一个标准化流。这样做的好处是，你可以更改流中任意的一个双射函数的输入顺序 $x_1,...x_D$ ，这样就可以在某一环节（因为不合时宜的排序）效果不佳，而在另一个环节弥补。
掩模自回归流(MAF）实现了一个条件高斯自回归模型。这里对于任何一个分布 $x_i$ :

灰色的单元是当前待计算的单元，蓝色的单元表示其依赖的单元。和通过计算传递给网络（分别是涂上了洋红色、橘色的圆圈）。即使变换只有尺度变换和平移变换，这些变换也错综复杂地依赖先前的变量。对于第一个单元，其和不依赖于x和u，而设置为可学习的变量。
注意，之所以如此这般地设计变换函数，是因为计算的时候，不必要再计算和的逆。由于变换仅有尺度变换和平移变换，我们只需要逆运算这两个变换：。双射函数的前传和反传，仅仅取决于和，这样我们就可以在和的网络内部使用非线性激活函数(如ReLU)和非方形矩阵操作。
MAF模型的梯度反传可以这样来评估密度函数:
distribution.log_prob(bijector.inverse(x)) + bijector.inverse_log_det_jacobian(x))

2. 时间复杂度和MADE

自回归模型和MAF训练较快，因为计算下面D个似然概率
$p(x_1), p(x_2\,|\, x_1), ... p(x_D\,|\, x_{1:D-1})$
可以利用GPU的并行技术，一次性地用D个线程计算。这里假设计算任务，如在计算SIMD向量化用CPU或GPU并行计算时那样，资源没有上限。
从另一方面看，自回归采样较慢，因为你必须等到 $x_{1:i-1}$ 都算出来，才能计算 $x_i$ 。如此只能利用一条线程采样D个序列，而无法发挥并行计算的优势。
另一个问题是，在并行的反传计算时，我们要不要使用两个不同的（具有不同的输入长度的）网络，来计算 $\alpha_i$ 和 $\mu_i$ ？这样做不够高效，尤其考虑到D个网络学到的表示（learned representation）需要共享（只要不违背自回归的依赖关系）。在“密度估计的掩模自编码器”（MADE）一文中，作者给这个问题提供了一个不错的解决方案：使用一个网络，把所有的 $\alpha$ 和 $\mu$ 同时计算出来，而用来掩模的权重保证了自回归的有效性。
这一技巧也大大简化了反传的计算，从所有的x，只需一步即可计算出所有的u（D个输入，D个输出）。这就比处理D个不同的网络容易的多（(D+1)*D/2个输入，D个输出）

3.Inverse Autoregressive Flow (IAF)

可逆自回归流（IAF）的非线性平移和尺度变换操作，其输入是前一时刻的噪声变量，而不是输入前一时刻的数据采样。
$x_i = u_i\exp{\alpha_i} + \mu_i$
$\mu_i = f_{\mu_i}(u_{1:i-1})$
$\alpha_i = f_{\alpha_i}(u_{1:i-1})$

前传（采样）速度很快：对于所有的，并行D个线程可以一次算出来。IAF也可以使用MADE框架高效地实现并行。
不过，要是你想估计某一新数据附近的概率密度，就需要恢复所有的u，过程很长：首先计算，然后逐次计算。好在计算这一趟以后，IAF生成的数据的（log）概率直接就得到了。因为我们已知u的情况下，不需要再通过反向计算x得到u。
细心的你可能已经察觉到了，要是把底下标注，顶层标注，IAF和MAF的反传计算没有差别。对应而言，IAF的反传计算恰好也就是MAF（把x和u交换一下）。所以在tensorflow实现的时候，双射函数的类没有区别，而且逆转特征来交换正反传的操作很容易：
iaf_bijector = tfb.Invert(maf_bijector)
IAF和MAF有互补的计算考虑——MAF训练快采样慢，IFA训练慢采样快。对于训练网络，需要做的密度估计操作更多，采样操作较少，所以通常在学习密度的任务中选择MAF更为合适。

4.平行声波网络（Parallel Wavenet）

显而易见的问题是，IAF和MAF能否合二为一，达到最优性能？比如快速的训练和采样。
答案是：当然能！DeepMind实验室公布的平行声波网络正是这么做的：MAF快速训练一个产生式网络，IAF在这个网络的“启发”下最大化采样的似然概率。回顾一下IAF，计算外来数据的概率密度很慢（比如训练集中的数据），但是计算采样数据 $u_{1:D}$ 的密度很容易，这样就省了反传计算。只要约束“学生”IAF和“老师”MAF的概率分布差异，就可以完成训练。

这项研究在标准化流的领域里非常重要——最终实现的效果是，实时的音频合成的速度比采样的方法快20倍，而且已经在谷歌助手之类的产品上线了。

5.NICE and Real-NVP

最后请看Real-NVP，你可以认为是IAF双射函数的特殊情况。
训练时，在NVP的“配对层”（coupling layer），依次地调整标号 $d=1,2,...,D-1$ 。像IAF一样， $x_{d+1}$ 的尺度变换和平移变换取决于 $u_d$ 的值。所不同的是， $x_{d+2}, x_{d+3}, … x_{D}$ 都只依赖于 $u_{d}$ ，所以单次的网络前传就可以得到 $\alpha_{d+1:D}$ 和 $\mu_{d+1:D}$ 。【译者按：这里可能的意思是，算出来x以后，求个逆就得到 $\alpha$ 和 $\mu$ 】
对于 $x_{1:d}$ ，这些单元“放行”了所有信号，并被设置为 $u_{1:d}$ 。所以可以认为Real-NVP是一种特殊的MAF双射（当 $\alpha(u_{1:d}) = \alpha(x_{1:d})$ ）。

因为尺度平移变换后，整个层的统计量可以从或者一次前传得到，NVP可以在一次前传和后传完成所有的计算（也就是说，采样和估计都很快）。MADE架构就不再需要了。
从实验效果来看，Real-NVP比MAF或IAF都要差，使用同样多的层数NVP在我的例子中（比如SIGGRAPH形）也更差。Real-NVP和IAF在二维图形的情况下几乎是等价的。唯一的区别是在第一个单元，IAF通过尺度平移变换得到，不依赖于，而Real-NVP对第一个单元不作处理。
Real-NVP是NICE双射的后续工作。NICE只有平移变换并假定。因为NICE不做尺度变换，ILDJ始终是个常数！

6.批正则化的双射函数（batch normalization bijector）

Real-NVP 的论文中提到许多新奇的分布，其中之一是批正则化的双射函数，用来稳定训练过程。在传统的观念中，批正则化应用在训练神经网络的场景中，前传数据服从统计意义上的中心集中、方差是对角阵的高斯分布，批数据是正态分布的统计特性（running mean,running variance)通过指数移动平均值积累。测试时，积累的统计数据用来做标准化处理。
在标准化流中，训练时批正则化用在双射函数的逆的计算中bijector.inverse，在测试时积累的统计数据要去标准化（bijector.forward）。具体而言，批正则化双射函数这样实现：
反传：
1.计算现有变量x的均值和方差
2.更新当前的均值和方差
3.用当前的均值方差批正则化当前数据。
前传：
1.用当前的均值方差去标准化，得到数据的分布。
归功于TF双射函数，这一过程可用这些代码来实现：

class BatchNorm(tfb.Bijector):
    def __init__(self, eps=1e-5, decay=0.95, validate_args=False, name="batch_norm"):
        super(BatchNorm, self).__init__(
            event_ndims=1, validate_args=validate_args, name=name)
        self._vars_created = False
        self.eps = eps
        self.decay = decay

    def _create_vars(self, x):
        n = x.get_shape().as_list()[1]
        with tf.variable_scope(self.name):
            self.beta = tf.get_variable('beta', [1, n], dtype=DTYPE)
            self.gamma = tf.get_variable('gamma', [1, n], dtype=DTYPE)
            self.train_m = tf.get_variable(
                'mean', [1, n], dtype=DTYPE, trainable=False)
            self.train_v = tf.get_variable(
                'var', [1, n], dtype=DTYPE, initializer=tf.ones_initializer, trainable=False)
        self._vars_created = True

    def _forward(self, u):
        if not self._vars_created:
            self._create_vars(u)
        return (u - self.beta) * tf.exp(-self.gamma) * tf.sqrt(self.train_v + self.eps) + self.train_m

    def _inverse(self, x):
        # Eq 22. Called during training of a normalizing flow.
        if not self._vars_created:
            self._create_vars(x)
        # statistics of current minibatch
        m, v = tf.nn.moments(x, axes=[0], keep_dims=True)
        # update train statistics via exponential moving average
        update_train_m = tf.assign_sub(
            self.train_m, self.decay * (self.train_m - m))
        update_train_v = tf.assign_sub(
            self.train_v, self.decay * (self.train_v - v))
        # normalize using current minibatch statistics, followed by BN scale and shift
        with tf.control_dependencies([update_train_m, update_train_v]):
            return (x - m) * 1. / tf.sqrt(v + self.eps) * tf.exp(self.gamma) + self.beta

    def _inverse_log_det_jacobian(self, x):
        # at training time, the log_det_jacobian is computed from statistics of the
        # current minibatch.
        if not self._vars_created:
            self._create_vars(x)
        _, v = tf.nn.moments(x, axes=[0], keep_dims=True)
        abs_log_det_J_inv = tf.reduce_sum(
            self.gamma - .5 * tf.log(v + self.eps))
        return abs_log_det_J_inv

ILDJ的数学形式可方便地从逆函数的对数形式导出（参考单变量的例子）。

7.代码在此

归功于JoshDillon和谷歌贝叶斯流研究组的努力，在MADE架构下的Masked Autoregressive Flow 来实现u的快速训练的代码已经写好了。
我创建了一个复杂的二维分布，是利用这个混合脚本点云“SIGGRAPH”的形状。我们搭建了数据库、双射函数、转换分布，方法和教程一中非常类似，所以就不重复了，你可以参看这个Jupyter notebook 文件。这个代码可以通过MAF、IAF、Real-NVP实现，使用或不使用批正则化，重构双月形和“SIGGRAPH”形。
一个细节很容易被忽略，如果你忘记重置变量顺序，这个代码就不能正常工作。这是因为不重置顺序的话，所有的自回归分解就只能学到 $p(x_1 | x_2)$ 。幸好，Tensorflow有一个换序双射函数，可以专门实现这一功能。
这里可视化了流的学习过程，从一而终。看到它，我想起一台太妃糖搅拌机

8.教程总结

Tensorflow让标准化流的实现变得简单，自动收集所有雅各比矩阵行列式的操作让代码变得简单易读。如果你要选择一个实现标准化流的方案，要同时考虑前传和后传的快捷性，以及平衡流的表达能力和ILJD的计算速度。
在该教程的第一部分，我介绍了标准化流的动机：我们需要更强大的分布函数，用于强化学习和产生式模型。从大背景来看，在变分推断和隐式密度估计模型成功应用的今天，能够追踪变换的标准化流模型是否是人工智能应用（如机器人、架构估计）中最适合的模型，尚不明确。即便如此，标准化流仍是你工具箱中非常好用的工具，它们已经在现实应用中发挥作用，比如谷歌助手团队开发的实时的产生音频的模型。
即便像标准化流这样的显式的密度估计模型能方便地通过最大似然法训练，这些模型的作用远不止于此。它们和VAE、GAN这类模型是互补的。任何模型中的高斯分布直接就换上一个标准化流，比如VAE的先验分布和GAN的隐变量。举个例子，这篇论文用标准化流灵活选择变分先验，Tensorflow分布的论文提到VAE中标准化流做PixelCNN的解码器。平行声波网络通过最小化KL散度训练一个IAF“学生”模型。
标准化流最具启发的性质之一是每一步计算都是可逆的（即具有函数的定义清楚的逆）。这就意味着，如果我们想做一次反传计算时，不必在前传时预先存储激活值，而是重新计算一遍前传激活值。（存储这些值有可能代价很大。）在任务分配的过程很长的情况下，我们可以使用计算的可逆性来“恢复”过去的每一个选择状态，从而限制了计算时的内存消耗。这个主意在论文 RevNets中有所体现，正是受NICE双射函数运算可逆的启发。我回忆起电影Memento中的情节，主角不能存贮记忆了，于是通过可逆计算来记住事情。
谢谢阅读！

Code on Github

9.致谢

非常感谢 Dustin Tran, Luke Metz, Jonathan Shen, Katherine Lee, Samy Bengio 等人预先检查阅读了该教程。

10.引用文献和推荐阅读

这篇博客的内容和格式都受到 Masked Autoregressive Flow for Density Estimation
一文很深的影响。该论文写的很好，我最初理解这一研究主题，或多或少起源于此。快去读一下吧！
较早涉及NFS的论文有： https://math.nyu.edu/faculty/tabak/publications/Tabak-Turner.pdf and https://arxiv.org/pdf/1302.5125.pdf 和 https://arxiv.org/abs/1505.05770。
Laurent Dinh的演讲以及与推特皮质细胞（Twitter Cortex）研究者的讨论。一些简明扼要的观点来源于此。
用PyMC包实现标准化流教程。
还有一些工作我至今还一知半解，那就是弄清从标准化流到Langevin Flow(郎之万流) 以及 Hamiltonian Flow(哈密顿流)之间的关系。研究标准化流中的双射函数的文献汗牛充栋，其中一支研究到连续时间流（Continous-Time Flow，显然其变换表达的能力更强大。

11.译者补充与提问解答

我很喜欢这个教程，清晰易懂，所以把它全文翻译成中文。考虑到中英文表达习惯的差异，译文并不完全是原文的原句（原词）对应，但意思（大概）还原了原文的意思。
这里补充几个问题与其解答，也欢迎读者向我（或原作者）提问！

问题1：这篇论文Parallel WaveNet:Fast High-Fidelity Speech Synthesis涉及到了并行化，MAF作为teacher模型，IVF作为student模型，我始终不是很理解这个方法的原理。

回答1：
这个问题拆解为两部分。第一，为什么要让MAF做教师，IAF做学生；第二，并行化的训练是怎样完成的。
============================================
第一，原文提到了MAF训练快采样慢，IAF训练慢采样快。所谓训练就是从数据算出模型参数的过程，采样就是从模型参数产生新数据的过程。
+++++++++++++++++++++++
对于MAF，训练的时候，把所有的数据送进两个网络， $f_\alpha$ 和 $f_\mu$ 。在允许并行的情况下，一步就可以算出来所有的 $\alpha_{2,...,n}$ 和 $\mu_{2,...,n}$ ，结合采样得到的所有的 $u_{1...n}$ ，就可以算出 $\hat{x}_{1,...,n}=u_{1...n}exp(\alpha_{1,...,n})+\mu_{1,...,n}$ 。（注意这里 $\alpha_1$ 和 $\mu_1$ 是超参，不是从 $x_1$ 计算出来的）然后就可以去训练 $f_\alpha$ 和 $f_\mu$ 的参数，这个过程非常快。
采样的时候，我们手里拿着所有 $f_\alpha$ 和 $f_\mu$ 的参数，并 $\alpha_1$ 和 $\mu_1$ ，从 $x_1=u_1exp(\alpha_1)+\mu_1$ ， $\alpha_2=f_\alpha(x_1)$ ， $\mu_2=f_\mu(x_1)$ ，开始算， $x_2=u_2exp(\alpha_2)+\mu_2$ ， $\alpha_3=f_\alpha(x_1,x_2)$ , $\mu_3=f_\mu(x_1,x_2)$ ...这样一直算到 $x_n$ 算完，这个过程是串行的而不能做到并行。
注意，采样阶段， $u_i = (x_i-f_\mu(x_i))/\exp(f_\alpha(x_i))$ 这个式子里， $u_i$ 其实就是 $p(x_i)$ 。
+++++++++++++++++++++++
对于IAF，训练的时候，从超参 $\alpha_1$ 和 $\mu_1$ 开始计算得到 $u_1 = (x-\mu_1) * \exp(-\alpha_1)$ ，把它们送进两个网络， $f_\alpha$ 和 $f_\mu$ ， $\mu_2 = f_{\mu}(u_{1})$ ， $\alpha_2 = f_{\alpha}(u_{1})$ 。继续计算 $u_2 = (x-\mu_2) * \exp(-\alpha_2)$ 。其余的以此类推 $u_i = (x-\mu(u_{1:i-1})) * \exp(-\alpha(u_{1:i-1}))$ 。用所有的 $u_{1...n}$ 、 $\alpha_{1...n}$ 和 $\mu_{1...n}$ ，就可以算出 $\hat{x}_{1,...,n}=u_{1...n}exp(\alpha_{1,...,n})+\mu_{1,...,n}$ 。然后就可以去训练 $f_\alpha$ 和 $f_\mu$ 的参数。这个过程不能并行。
采样的时候，我们手里拿着所有 $f_\alpha$ 和 $f_\mu$ 的参数，采样基础分部产生 $u_{1...n}$ ，然后一步就可以算出所有的 $\alpha_{1,...,n}$ 和 $\mu_{1,...,n}$ ，于是 $x_{1,...,n}=u_{1...n}exp(\alpha_{1,...,n})+\mu_{1,...,n}$ 。这个过程是非常快。
+++++++++++++++++++++++
============================================
Parallel WaveNet训练过程是这样的：
1.训练一个MAF，得到合理的 $f_{\alpha MAF}$ 、 $f_{\mu MAF}$ 。详细步骤见上文。
2.采样基础网络，得到的随机噪声 $u_i$ ，送进IAF里得到 $x^i=u_iexp\alpha_i+\mu_i$ （注意这个IAF没有训练好），和 $log p_s(x^i) = log(x^i - f_{\mu IAF}(x^i)) -f_{\alpha IAF}(x^i)$ 。
3.给定测试数据 $x^i$ ，利用MAF的反向传递，计算(log)似然概率 $log p_t(x^i) = log(x^i -f_{\mu MAF}(x^i)) - f_{\alpha MAF}(x^i)$ 。
4.计算KL距离，作为损失。
============================================
要彻底理解这个框架，要对最大似然法和EM算法有个大概了解。我贴一张图，从知乎拷贝的。侵权就删：