标准化流(Normalizing Flow)教程（一）

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

0.背景

0.0 交代故事

你是个机器学习从业者，正苦哈哈地研究着产生式模型、贝叶斯深度学习、深度强化学习之类的玩意儿，趁手的工具捉襟见肘。今天我就给你带一件儿不错的东西：标准化流，相信我，它会武装你的算法工具箱，带你飞起。标准化流能把简单的地摊货概率密度(比如高斯分布)形式转换成某种高大上的分布形式。它可以用在产生式模型、强化学习、变分推断之类的地方。至于实现，Tensorflow有一类不错的函数，它们可以搭建出标准化流，并且在真实的数据库上训练得熨熨帖帖。
这个教程分为两部分：

部分1：分布和行列式。我将解释概率分布的可逆转换是怎么回事儿，这个技术是怎么用在复杂概率分布的拟合上的，以及这些转换怎么连起来成为一个标准化的“流”。
部分2：现代标准化流技术。我将调研最近研究者如何发展了这个技术，并借此诠释一大家子产生式模型——自回归模型(autoregressive models)、MAF、IAV、NICE、Real-NVP、平行波网络(Parallel-Wavenet)——它们彼此关联。

要是你已经掌握了线性代数、概率论、神经网络、Tensorflow的相关知识，那就可以上车啦！能了解一点深度学习、产生式模型的最新进展，就更靠谱了，不知道也没多大关系。

0.1 统计机器学习

统计机器学习旨在获悉参数分布 $p(x;\theta)$ ，以揭示数据的结构。以此出发，你将能达成以下四个目标：
1.采样该结构，生成数据。这样就节省了完整的生成过程，若完整的过程纷繁浩杂，则节省巨大[1]。同样，现实数据的维度高，采样后的低维结果就可管中窥豹，节省计算。
2.评估测试数据的似然概率。在拒绝采样或评估模型时有用。【译者按：看到后面你就明白这一条了。】
3.获得变量间的条件依赖关系。如 $p(x_2|x_1)$ 可用来判别或回归。【译者按：这一条与上一条刚好相对应，一个是unconditional likelihood estimation，一个是conditional likelihood estimation。】
4.评判算法。如熵、互信息、高阶矩等指标，均可一试。

以上四个目标，第一个研究的最充分。合成图像、声音，已在谷歌商用。但第二三四个，研究寥寥。仅举几例，GAN的解码器，支撑集（映射到非零值的自变量）不可得；DRAW模型乃至VAE模型，概率密度不可得；哪怕已解析地了解某分布，解析的度量（如KL距离，earth-mover距离）还是不可得。【译者按：这句话说的简单点就是：深度产生式模型不容易甚至不可能进行测试和度量】

即使你能合成个把接近真实数据的样本，你还离目标差得远咧！我们想要了解它们多大程度上“接近真实数据”[2]，想要得到灵活的条件密度（如增强学习中复杂的规则），还想要在一大家子的先验分布和后验分布上做变分推断！

0.2 高斯分布

你思索片刻，想起了你的“亲密战友”高斯分布。它在你为复杂的概率分布理出头绪而焦头烂额的时候，雪中送炭：采样方便、解析的密度已知、KL距离容易计算，还有中心极限定理的保证——任何大的数据都趋近于高斯分布，所以你怎么用它几乎都对！更别提重参数化技巧啦，你还能梯度反传。这些易用的性质，让高斯分布在产生式模型和增强学习里风头无两。

不过先别高兴的太早！许多实在的场景，我们还不能用它。在增强学习里，常有连续控制机器的场景，控制规则得用多变量高斯、对角化方差矩阵（motivariate Gaussians with diagonal covariance matrices）才行呐~

相形之下，只用单模式高斯分布模型（uni-modal Gaussian distribution），在需要多模式分布（multi-modal distribution）采样的时候，乏善可陈。试举一例，要机器人绕过湖面，到达湖后的房子，规则该是二者选一：向左绕或向右绕。单模式高斯模型容不下两个规则，所以就线性地中和了它们：直直地走向湖心。糟了个糕！

高斯分布失之过简。它不能容下对立的假设；高维的情况下分布也不够集中，出现边缘效应；还不能应对罕见事物。有没有更好的分布模型，可以满足如下条件呢？

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

答案是：有！你可以做这几件事。

0.3 你需要这样的分布

以下几个办法，能帮你实现需求：
1.模型融合。一个模型对付一个子任务。
2.立足当下。一次只考虑一步，让复杂的选择在此刻退化，而仅存一端。
3.非定向图模型（undirected graphical models）。
4.标准化流。你将得到一个理想的模型：可逆、可计算分布变换体积、易模拟。

下面介绍标准化流。

1. 改变了自变量，你就改变了体积

为了建立直观，举例如下。让X服从均匀分布 $Uniform(0,1)$ , $Y= f(X) = 2X + 1$ 。Y是X的仿射变换。如下图所示。

两个方块代表概率在实数域的分布函数p(x)和p(y)，函数值代表概率密度的值。概率值的积分必须等于1，任何一个分布都是这样。所以自变量变化范围扩展为两倍后，函数值就要相应地缩减一半，从而二者面积都等于1.【译者按：这就是为什么模型叫标准化流，因为有加和等于1的约束~】
若考虑X上的极小变化

x+dx

, Y也相应发生变化

y+dy

, 如下所示：

左半边的图代表一个局部增函数（

dy/dx>0

），右边则是局部减函数（

dy/dx<0

）。p(x)因为dx带来的局部变化一定要等于p(y)随着dy的局部变化：

p(x)dx = p(y)dy

为了让概率的变化守恒，我们只关心变化的量，而不关心变化的方向。（实际上，不管f(x)随着x变化增或者减，都不影响，我们假设y也做出相应的变化。）那么，我们有

p(y)=p(x)|dx/dy|

。注意到在log空间，这等价于

logp(y) = logp(x) + log|dx/dy|

。在log空间中计算有利于保证计算可靠性。【注意到，这里硬生生就多了一项！】
现在考虑多变量的例子。先考虑两个变量的情况。同上面的情况，缩放定义域中的无穷小量，我们就得到一个二维的、边长为dx的小方块。
给这个小方块的顶点标注

(x_1, x_2, x_3, x_4)

。我们现在只关心dx相对于边长为1的方块的变化率，不妨假设这个小方块就是一个边长为1的原点处的方块，那么顶点为

(0,0), (1,0), (0,1), (1,1)

。
乘以矩阵

[[a, b];[c, d]]

,我们的方块就会变成一个平行四边形，如下图。(0,0)还在(0,0), (1,0)来到了(a,b),

(0,1)

被送往

(c,d)

(1,1)

到达

(a+c,b+d)

。

于是，一个X域边长为1的方块,转型为Y域的平行四边形，大小变成了

ad-bc

。这个平行四边形的面积，正是转换矩阵的行列式！
三维的情况时，“转换为平行四变形”就对应为“转换为平行六面体”，或者更高维的情况也是以此类推，“转换为平行n维体”。行列式的道理也还是如此，线性变换后的体积，正好对应于变换矩阵的行列式。
要是转换函数f是非线性的咋办？你就不该把它简单地看作全空间的一个简单的平行体，而是对空间的每个点，对应于一个无穷小平行体。数学形式上看，局部线性变换对应的体积变化为

|det(J(f^{-1}(x)))|

,这里

J(f^{-1}(x))

代表雅各比矩阵的逆——是

|dx/dy|

的高维变种。于是，

y=f(x)

p(y) = p(f^{-1}(y))⋅|detJ(f^{-1}(y))|

logp(y)=logp(f^{-1}(y))+log|det(J(f^{-1}(y)))|

敝人在中学就学了行列式，完全不理解行列式的定义，只知道计算手段。知道真相后，无语凝噎：行列式就是局部线性转换的体积变化率。

2.Tensorflow实现转换分布

Tensorflow有一个优雅的API来实现分布的转换。一个转换后的分布由以下二者唯一确定：1.基础分布，就是我们要转换的那个分布；2.一个双射函数,它由三部分构成 1）前向映射 $y=f(x)$ ,f从d维实空间映射到d维实空间；2)反向映射 $x=f^{-1}(y)$ ；3）雅各比矩阵的逆对数行列式（inverse log determinant of the Jacobian） $log|detJ(f^{-1}(y))|$ ,后文简称为ILDJ。
在以上约定下，实现前向采样易如反掌：

bijector.forward(base_dist.sample())

评估转换后分布的对数密度，只需要：

distribution.log_prob(bijector.inverse(x)) + bijector.inverse_log_det_jacobian(x)

【译者按：就是上一节的公式】
更进一步，要是bijector.forward是一个可微函数，那么Y = bijector.forward(x)可以通过x的分布重参数化：x = base_distribution.sample()。这就意味着标准化流可以用在VAE的变分推断的环节，将原来版本的高斯分布改为你需要的形式。【译者按：这里是我觉得该模型最牛的地方。】
这里列出几个常用的Tensorflow分布，它们的实现实际上是通过一下的途径：

基础分布	双射函数	转换后的分布
正态分布	$exp(x)$	对数正态分布
$exp(rate=1)$	$-log(x)$	$Gumbel(0,1)$
$Gumbel(0,1)$	$Softmax(x)$	$Gumbel-Softmax / Concrete$

从约定俗成的叫法，转换后的分布叫做双射的逆函数映射后的基础分布(Bijector^−1BaseDistribution)【译者按：我并不知道标准译法┓( ´∀` )┏】，所以指数双射(ExpBijector)后的正态分布是对数正态分布。当然，这个命名规则也有例外，Gumbel-Softmax分布被命名为约束松弛的幺零类别（RelaxedOneHotCategorical ）分布，它通过Softmax函数映射一个Gumbel分布。

3. 标准化流

用了一个双射就停手吗？停手是不会停手的，这辈子都不会停手的。你可以把一系列双射连起来，在神经网络里像链子一样把它们拴在一起[3]。这个结构就叫“标准化流”。要是双射函数在bijector.log_prob有可变的参数，你就可以优化这个参数，该双射就可以把基础分布转换成任意的分布。每个双射函数可以写成一个网络的层，你可以用一个优化器来学习参数，最终拟合真实数据。算法通过最大似然估计，把拟合真实数据的分布问题变成拟合变换后的概率的对数密度问题。用对数密度的原因是为了计算的稳定性。
这页幻灯片取自Shakir Mohamed和Danilo Rezende在UWA的演讲:

image.png

不过，计算一个

N \times N

的雅各比矩阵的复杂度是

O(N^3)

,对于神经网络代价颇高。还有一个麻烦，逼近任意一个函数的逆也挺难。最近的标准化流研究重点，就放在利用GPU并行计算的优势、设计出表现力强的双射函数上，也就是放在如何方便地计算ILDJ上。

4.代码在此

咱们一起用100行代码来实现一个简单的标准化流吧！这个代码用到：

TF Distribution ——一个Tensorflow通用API。你需要1.5版以上的Tensorflow。
TF Bijector ——一个创建分布函数的通用API。
Numpy, Matplotlib。

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
tfd = tf.contrib.distributions
tfb = tfd.bijectors

你要拟合的分布是这个
$p(x_1, x_2)=N(x_1|μ = 1/4 {x_2}^2,\sigma^2= 1)⋅N(x_2 | μ = 0,σ^2= 4)$
你先采样一些样本（用Tensorflow的GPU功能直接采样，避免从CPU到GPU复制一次）

batch_size=512
x2_dist = tfd.Normal(loc=0., scale=4.)
x2_samples = x2_dist.sample(batch_size)
x1 = tfd.Normal(loc=.25 * tf.square(x2_samples),
                scale=tf.ones(batch_size, dtype=tf.float32))
x1_samples = x1.sample()
x_samples = tf.stack([x1_samples, x2_samples], axis=1)

$x_1$ 长这样：

基础分布是各向同性的高斯分布：

base_dist = tfd.MultivariateNormalDiag(loc=tf.zeros([2], tf.float32))

下一步，建立一个双射，并从这个双射搭建一个转换后的分布。你应该建立一个规规矩矩的全连接网络，就是没有什么非线性激活函数的矩阵乘法操作。
仿射函数的雅各比矩阵的计算相当廉价，不过最糟的时候行列式的计算却需要 $O(n^3)$ 这么多，慢成了乌龟爬。好在Tensorflow提供了一个仿射变换的框架，计算行列式相对简单。这一仿射变换参数化为下三角矩阵M加上一个低秩更新矩阵：
$M+V·D·V^T$
要方便地计算 $det(M+V·D·V^T)$ ，你需要矩阵行列式引理。
下一步，你需要可逆的非线性操作来表达非线性（不然仿射函数的组合仍是仿射函数）。Sigmoid和Tanh函数似乎是不错的选择，不过求逆的时候不稳定——输出在-1和1两点附近的一个小小的差别对应于输入的巨大变化。在你的实验里，不大可能在连着两个饱和的非线性操作后还能在梯度反传时稳如泰山。ReLU虽然是稳定的，不过在x<0时不可逆。
所以你可以选择PReLU(带参数的ReLU)，和Leaky ReLU一样，在负数定义域内有可学习的梯度。PReLU简单明了，你就可以拿雅各比矩阵练手了：ILDJ等于0，当x>0（没有体积变化），否则等于1/α（对x乘以α就刚好补偿了体积的压缩）。【译者按：1）没有体积变化是因为x>0时PReLU的函数形式为f(x)=x。x<0时以此类推2）作者用的函数名是LeakyReLU，实际就是PReLU】

# quite easy to interpret - multiplying by alpha causes a contraction in volume.
class LeakyReLU(tfb.Bijector):
    def __init__(self, alpha=0.5, validate_args=False, name="leaky_relu"):
        super(LeakyReLU, self).__init__(
            event_ndims=1, validate_args=validate_args, name=name)
        self.alpha = alpha

    def _forward(self, x):
        return tf.where(tf.greater_equal(x, 0), x, self.alpha * x)

    def _inverse(self, y):
        return tf.where(tf.greater_equal(y, 0), y, 1. / self.alpha * y)

    def _inverse_log_det_jacobian(self, y):
        event_dims = self._event_dims_tensor(y)
        I = tf.ones_like(y)
        J_inv = tf.where(tf.greater_equal(y, 0), I, 1.0 / self.alpha * I)
        # abs is actually redundant here, since this det Jacobian is > 0
        log_abs_det_J_inv = tf.log(tf.abs(J_inv))
        return tf.reduce_sum(log_abs_det_J_inv, axis=event_dims)

PReLU是一个元素级别的变换，所以雅各比矩阵是对角的。对角矩阵的行列式等于对角线上元素的乘积，所以对应在对数空间就是直接加和，得到对数雅各比矩阵[4]。你通过tfb.Chain搭建MLP双射，并作用在基础分布上。

d, r = 2, 2
DTYPE = tf.float32
bijectors = []
num_layers = 6
for i in range(num_layers):
    with tf.variable_scope('bijector_%d' % i):
        V = tf.get_variable('V', [d, r], dtype=DTYPE)  # factor loading
        shift = tf.get_variable('shift', [d], dtype=DTYPE)  # affine shift
        L = tf.get_variable('L', [d * (d + 1) / 2],
                            dtype=DTYPE)  # lower triangular
        bijectors.append(tfb.Affine(
            scale_tril=tfd.fill_triangular(L),
            scale_perturb_factor=V,
            shift=shift,
        ))
        alpha = tf.abs(tf.get_variable('alpha', [], dtype=DTYPE)) + .01
        bijectors.append(LeakyReLU(alpha=alpha))
# Last layer is affine. Note that tfb.Chain takes a list of bijectors in the *reverse* order
# that they are applied.
mlp_bijector = tfb.Chain(
    list(reversed(bijectors[:-1])), name='2d_mlp_bijector')
dist = tfd.TransformedDistribution(
    distribution=base_dist,
    bijector=mlp_bijector
)

最后，你训练模型，用最大似然法则：最大化样本的对数概率期望。这些样本是你选择的模型下，从真实数据的分布采样出来的。

loss = -tf.reduce_mean(dist.log_prob(x_samples))
train_op = tf.train.AdamOptimizer(1e-3).minimize(loss)
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
NUM_STEPS = int(1e5)
global_step = []
np_losses = []
for i in range(NUM_STEPS):
    _, np_loss = sess.run([train_op, loss])
    if i % 1000 == 0:
        global_step.append(i)
        np_losses.append(np_loss)
    if i % int(1e4) == 0:
        print(i, np_loss)

你可以可视化学习过程（很慢唉），并且把不同象限的点涂上不同的颜色：

齐活!Tensorflow能自动汇集所有雅各比矩阵行列式，利索又易读。完整代码在Github里。
你可能已经发现矩阵变换得很慢，需要许多全连接层才能实现转换[5]。在教程下半部分，我会涉及标准化流的现代技术。

致谢

非常感谢Dustin Tran给我解释清楚标准化流， Luke Metz, Katherine Lee,还有 Samy Bengio预先校读本篇教程，并且 Ben Poole, Rif A. Saurous, Ian Langmore帮我修改代码。你们是最棒的！

脚注

[1]我对于是否可以从一个算法学到的有限的数据集来给数据库增加新信息是很困惑的。对于基于概率的机器学习算法是否能真正取代真实的生成过程，尚须更多证据。可能将来会出现反面证据当头棒喝，告诉我们所有的生成算法仅仅能够节省计算，或者拟合的训练\测试集只是幸运的巧合。
[2]在 A note on the evaluation of generative models一文里面讨论了一个恼人的结论，高的对数似然对于生成“最可能”的图像既非必要，也非充分。不过，对数似然仍是比什么都没有要好，在实践中仍是一个判断模型的指标。
[3]标准化流和GAN模型可以通过编码-解码GAN结构来实现，这个模型学习的是产生器的逆（如ALI/BiGAN）。由于一个单独的编码器来反推 $u=G^{-1}(X)$ 来保证 $X=G(u)$ ，产生器可以认为是均匀分布的流。不过，我们不知道该对X进行何种扩展/收缩，所以我们不知道GAN学到的具体概率密度。你也可以搭建某种线性时间的雅各比矩阵，来从数值上建模ILDJ。
[4]引理“对角矩阵的行列式等于对角线元素的积”是很直观的，在几何上相当于把每个独立维度的长度相乘，整体的体积变化就正好等于每个方向上的变化的积。正如我们计算高维的矩形多面体一样。
[5]全连接层的拟合能力不高，因为每个仿射变换是一个2x2的矩阵，还因为PReLU把分布“扭曲”的速度非常慢（所以需要很多PReLU才行）。对于低维的分布，全连接层是非常糟的选择，这里这么做仅是为了教学目的。

【Normalizing Flows尚无标准的中文译名。Flow指的是数据“流”过一系列双射（可逆映射），最终映射到合适的表征空间；Normalizing指的是，表征空间的变量积分为1，满足概率分布函数的定义。无疑“标准化流”这个名字不够古雅，也不足够有启发性，“归一化流”同样糟糕；“流形”好一些，然而已经被用过；“涓流”在电池学中用过，但体现不了积分为1的特点。或许可以叫做“曲规”。（语出《荀子劝学》“木直中绳，輮以为轮，其曲中规”，意思是弯曲以符合圆规的模子。）无论如何，这篇教程中仍叫做“标准化流”。】