《Normalizing Flows for Probabilistic Modeling and Inference》论文笔记

1. 前置知识

标准化流(Normalizing Flow)能够将简单的概率分布转换为极其复杂的概率分布,可以用在生成式模型、强化学习、变分推断等领域,构建它所需要的工具是:行列式(Determinant)、雅可比矩阵(Jacobi)、变量替换定理(Change of Variable Theorem),下面先简单介绍这三个工具。

1.1 行列式

行列式的求法不再赘述,我们主要需要理解的是行列式的物理意义。一个矩阵的行列式的值表示的是该矩阵对空间所做的变换,将原来的空间放大或缩小了多少倍,比如二维空间在原点有一个边长为1的正方形a,对它做变换得到新的正方形b,$b = Wa$,$W=\left[\begin{matrix} 2 & 0 \ 0 & 2\end{matrix}\right]$,新的正方形边长被放大为原来的2倍,面积为原来的4倍,$det(W) = 4$,三维以及更高维空间亦同理

image-20200728143736053

1.2 雅可比矩阵

设有一个二维向量$z=\left[\begin{matrix} z_{1} \ z_{2} \end{matrix}\right]$,给定变换$f$,$x=f(z)=\left[\begin{matrix} z_{1}+z_{2} \ 2z_{1}\end{matrix}\right]$,那么变换$f$的雅可比矩阵$J_{f}$

变换$f$的逆变换$f^{-1}$,$z=f^{-1}(x)=\left[\begin{matrix} \frac{1}{2}x_{2} \ x_{1}-\frac{1}{2}x_{2} \end{matrix}\right]$,$f^{-1}$的雅可比矩阵$J_{f^{-1}}$

可以发现$J_{f}$与$J_{f^{-1}}$互为逆矩阵,事实上互为逆变换的$f$与$f^{-1}$,其二者对应的雅可比矩阵也互为逆阵,因此又由行列式的性质可得它们的雅可比行列式互为倒数,即

变换$f$可以不仅仅是矩阵变换,也可以任意的函数,将D维的向量$z$变换为D’维的$x$,$f: \mathbb{R}^{D} \rightarrow \mathbb{R}^{D’}$。

1.3 Change of Variable Theorem

假设有一变量$\mathbf{u}$,服从分布$\mathbf{u} \sim p_{u}(\mathbf{u})$,有一变换$T$,$\mathbf{x}=T(\mathbf{u})$,$p_{u}(\mathbf{u})$是已知的一种简单分布,变换$T$可逆,且$T$与$T^{-1}$都可微分,现在要求$p_{x}(\mathbf{x})$,即随机变量$\mathbf{x}$的概率密度函数,因为概率之和相等

因为x与u一一对应,所以从du体积映射到dx体积时,体积内包含的概率之和不变,那么被积分的部分绝对值必定处处相等,由于概率p必大于等于0,可去掉其两边的绝对值号,即得

第(8)式也可写为

(5)~(8)式的直观理解:两边都是概率密度乘以空间的大小,得到是一个标量,即随机变量落在该空间的概率大小,将变换$T$写入$\left| \frac{d \mathbf{z}}{d \mathbf{x}} \right|$,即将其写为$\left| \frac {\partial T^{-1}(\mathbf{x})}{\partial \mathbf{x}} \right|$,但x为向量而非标量,这里$\left| \frac{d \mathbf{z}}{d \mathbf{x}} \right|$要表示的是空间变化的大小关系,我们由雅可比矩阵的定义可知$\frac {\partial T^{-1}(\mathbf{x})}{\partial \mathbf{x}} = J_{T^{-1}}(\mathbf{x})$,又由行列式的物理意义,知道$J_{T}(\mathbf{z})$的绝对值为$T$将z映射到x时,空间大小放缩的倍数,即为概率密度放缩的倍数倒数,又因$\mathbf{det}J_{T^{-1}}(\mathbf{x}) = \mathbf{det}J_{T}(\mathbf{z})^{-1}$,因此可得(8)式。

论文原文的解释大意是:我们可以认为T在通过expand或contract R^D空间来使得pz变得和px相近。雅可比行列式detT的绝对值量化了原向量z附近的体积由于经过T变换后,体积相对变化的大小,即当z附近一块无穷小的体积dz,经过T后被map到了x附近的一块无穷小的体积dx处,那么detT等于dx除以dz,即映射后的体积是原来的几倍,因为dz中包含的概率等于dx中包含的概率,因此如果dz的体积被放大了,那么dx里的概率密度应该缩小

举个例子,假设随机变量z属于0-1均匀分布,在取值空间$C_{1}$=(0, 1)上,p(z)=1,有变换T,T(z)=2z,令x=T(z),则x必是$C_{2}$=(0, 2)上的均匀分布,但此时p(x)不再是1了,否则在(0, 2)上都有p(x)=1,积分可得概率之和为2,明显错误,因为变换T将原空间中z可取值的范围放大了一倍,从(0, 1)变为了(0, 2),即可取值空间从$C_{1}$变为$C_{2}$,空间放大的倍数为$\left| \mathbf{det} J_{T}(\mathbf{z}) \right|=2$,那概率密度缩小的倍数为$\left| \mathbf{det} J_{T^{-1}}(\mathbf{x}) \right| = \frac{1}{2}$,即相应的x概率密度应该缩小一倍,因此

2. 标准化流的定义和基础

我们的目标是使用简单的概率分布来建立我们想要的更为复杂更有表达能力的概率分布,使用的方法就是Normalizing Flow,flow的字面意思是一长串的T,即很多的transformation。让简单的概率分布,通过这一系列的transformation,一步一步变成complex、expressive的概率分布,like a fluid flowing through a set of tubes,fluid就是说概率分布像水一样,是可塑的易变形的,我们把它通过一系列tubes,即变换T们,塑造成我们想要的样子——最终的概率分布。下面开始使用的符号尽量与原论文保持一致。

2.1 Normalizing Flow’s properties

  1. x与u必须维度相同,因为只有维度相同,下面的变换T才可能可逆

  2. 变换T必须可逆,且T和T的逆必须可导

  3. 变换T可以由多个符合条件2的变换Ti组合而成

从使用角度来说,一个flow-based model提供了两个操作,一是sampling,即从pu中sample出u,经过变换T得到x,$\mathbf{x}=T(\mathbf{u}) ~~where ~~\mathbf{u}\sim p_{u}(\mathbf{u})$,另一个是evaluating模型的概率分布,使用公式 $p_{\mathrm{x}}(\mathbf{x})=p_{\mathrm{u}}\left(T^{-1}(\mathbf{x})\right)\left|\operatorname{det} J_{T^{-1}}(\mathbf{x})\right|$。

两种操作有不同的计算要求,sampling需要能够sample from pu 以及计算变换T,evaluating需要能够计算T的逆与雅可比行列式,并evaluate pu,因此计算时的效率与难度对应用来说至关重要

2.2 Flow-based models有多强的表达能力?

我们知道p(u)是很简单的一个概率分布,那么通过flow,我们能将p(u)转换为任意的概率分布p(x)吗?假设x为D维向量,p(x)>0,$x_{i}$的概率分布只依赖i之前的元素$x_{< i}$,那么可以将$p_{x}(x)$分解为条件概率的乘积

假设变换F将x映射为z,zi的值由xi的累积分布函数(cdf)确定

很明显F是可微分的,其微分就等于$p_{x}(\mathbf{x}_{i}|\mathbf{x}_{< i})$,由于Fi对xj的偏微分当j > i时等于0,因此$J_{F}(\mathbf{x})$是一个下三角矩阵,那么其行列式就等于其对角线元素的乘积,即

由于p(x)>0,所以雅可比行列式也>0,那么变换F的逆必存在,由(9)式,将x与z对调,T改为F,可得

即z是D维空间中(0, 1)之间的均匀分布。

上述对p(x)的限制仅仅是$x_{i}$依赖于$x_{< i}$的条件概率对$(x_{i}, x_{< i})$可微,且$p_{x}(\mathbf{x})>0 ~\forall~\mathbf{x}\in \mathbb{R}^{D}$,我们就使用变换F将它变为了最简单的(0, 1)均匀分布,又因为F可逆,所以我们可以使用F的逆将p(z)转换为任意满足上述条件的概率分布p(x)。我们再推广到任意的base distribution,假设p(u)满足上述p(x)满足的条件,那么我们可以使用变换G将任意的概率分布p(u)转换为p(z),再用F逆将p(z)转换为任意的概率分布p(x),即 使用变换$T = F^{-1} \circ G$,可将$p_u{\mathbf{u}}$变为$p_{x}(\mathbf{x})$。

2.3 使用flows来建模和推断

为了拟合一个概率模型,我们要拟合一个flow-based model $p_{x}(\mathbf{x};\theta)$去近似目标分布$p_{x}^{*}(\mathbf{x})$,$\theta$代表$(\phi,\psi)$ ,$\phi$和$\psi$分别是T与p(u)的参数,可以通过最小化KL散度和最大似然估计做到

2.3.1 正向KL散度与最大似然估计

假设现在我们手上有N条真实的数据,即来自于$p_{x}^{*}(\mathbf{x})$的samples,那么可以使用蒙特卡罗法来近似上面的期望值

最小化上式等价于求模型在该N条数据上的最大似然估计,我们一般使用随机梯度下降来优化上式的参数。

可见,为了使用正向KL散度或最大似然估计来拟合目标分布,我们需要计算$T^{-1}$、它的雅可比行列式, evaluate base分布 $p_{u}(\mathbf{u};\psi)$,以及关于它们参数的导数。

2.3.2 反向KL散度

当我们能够计算T,它的雅可比行列式,evaluate 目标分布p*以及从p(u)中sample时,使用反向KL散度是合适的,事实上,即使我们只能evaluate 目标分布乘以某个正则化常数,也可以最小化上式,$p_{\mathrm{x}}^{*}(\mathbf{x})=\widetilde{p}_{\mathrm{x}}(\mathbf{x}) / C$,$\widetilde{p}_{\mathrm{x}}(\mathbf{x})$是一个更好处理的概率分布,重写上式为

当我们有N条来自于$p_{u}(\mathbf{u;\psi})$的samples,为了最小化上式,使用蒙特卡罗法,并对变换T的参数$\phi$求偏导,可得

3. 构建标准化流(一):有限组合的变换

我们通过组合K个transformation得到T变换,令$\mathbf{z}_{0} = \mathbf{u}$,$\mathbf{z}_{K} = \mathbf{x}$

我们可以将每个$T_{k}$或$T_{k}^{-1}$都设为一个参数为$\phi_{k}$的神经网络,下面用$f_{\phi_{k}}$来统一表示这两者,但是这带来的问题是,我们必须保证该神经网络是可逆的,并且能够容易计算,否则,上述的正向KL散度需要变换$T_{k}$来做sampling,逆向KL散度需要$T_{k}^{-1}$来evaluating desity,如果变换$f_{\phi_{k}}$的逆不存在或不易求,则density evaluation或sampling将是效率极低的,甚至无法处理的。至于是否要求$f_{\phi_{k}}$有高效的求逆方法、$f_{\phi_{k}}$到底是使用$T_{k}$还是$T_{k}^{-1}$,则由具体的使用目的决定。下面忽略下标k,使用$f_{\phi}$表示神经网络,$\mathbf{z}$表示输入,$\mathbf{z}^{\prime}$表示输出。

3.1 自回归流 Autoregressive flows

2.2节我们知道了在合理的情况下,我们可以使用一个下三角雅可比矩阵将任意的概率分布$p_{x}(\mathbf{x})$变换为均匀分布,自回归流就是这样的一种构建方法

$\tau$就是我们的transformer,$c_{i}$是第i个conditioner,该变换是严格单调函数,因此必可逆。变换$\tau$的参数为$h_{i}$,$h_{i}$由conditioner决定,描述当输入的$\mathbf{z}_{i}$变化时,输出$\mathbf{z}_{i}^{\prime}$如何变。conditioner唯一的限制就是它只能将$\mathbf{z}_{< i}$作为输入,因此它可以是任意一个复杂的神经网络,不必关心是否可逆等问题。因此,可以看出,$f_{\phi}$的参数$\phi$其实就是conditioner的参数,但有时变换$\tau$也有它自己的额外参数(除了$h_{i}$)。上述变换的逆变换为

在正向计算中,因为输入$\mathbf{z}$是完全已知的,那么所有的$h_{i}$可以同时一次性求出来,因此$\mathbf{z}^{\prime}$也可以同时求出来,但是在求逆变换的计算时,要计算$\mathbf{z}_{i}$前必须先把$\mathbf{z}_{< i}$都计算出来,因为$\mathbf{z}_{< i}$是$h_{i}$的输入。
很明显,自回归流的雅可比矩阵是下三角形的,因为任意的$\mathbf{z}_{i}^{\prime}$都不依赖$\mathbf{z}_{>i}$,那么$\mathbf{z}^{\prime}_{\leq i}$关于$\mathbf{z}_{>i}$的偏导都为0,因此雅可比矩阵可写为

它的行列式就是对角线元素乘积,因此也非常好求

3.1.1 各种 Transformer$~\tau~$的实现

仿射自回归流(Affine autoregressive flows)

令$\tau$为下式

只要$\alpha_{i}$不为0,$\tau$变换的逆就存在,我们可以令$\alpha_{i}=\text{exp}(\tilde{\alpha_{i}})$,这样$\tilde{\alpha_{i}}$就是一个不受限制的参数了,该变换的雅可比行列式为

仿射自回归流虽然很简单,但它的一大缺点是表达能力受限(limited expressivity),假如z属于高斯分布,那么z’也必属于高斯分布,但可以通过堆叠多个仿射变换来增强表达能力,但仿射自回归流是否是一个普遍的概率分布近似器就未知了。

非仿射神经网络变换(Non-affine neural transformers)

不像上面连续使用K次变换,而是直接使用K个单调增函数$\sigma(\cdot)$的锥组合(conic combinations),h中的参数皆大于0。其实就是给仿射变换加上一个激活函数,再线性组合K种不同参数下的结果。一般使用反向传播优化参数,缺点是该变换往往不可逆,或者只能不断迭代求逆。

积分变换(Integration-based transformers)

$g(\cdot;\alpha_{i})$可以是任意的正值神经网络,导数很好求,就是$g(\mathbf{z}_{i};\alpha_{i})$,但积分缺乏analytical tractability,一种解决方法是让$g$为一个2L次的正多项式,积分结果就是关于$\mathbf{z}_{i}$的2L+1次的多项式,由于任意的2L次多项式可以写为多个L次多项式的平方之和,我们可以定义$\tau$为K个L次多项式平方之和的积分,如下

参数$\alpha$不受限制,仿射变换为L=0时的特例。由于5次及以上的方程没有根式解,因此当L>=2时,2L+1>=5,则无法直接求$\tau^{-1}$,只能使用二分搜索迭代求解。

神经样条流(Neural spline flows)

为了克服非仿射变换不易求逆,可以使用K个分段函数作为transformer,每个区间$[\mathbf{z}_{i(k-1)},\mathbf{z}_{ik}],k=1:K$上$\tau$是一个简单的低次单调函数,要求是在每个区间的端点必须与其相邻的区间连续,整个大区间两端的端点$\mathbf{z}_{i0},\mathbf{z}_{iK}$可以随意。

3.1.2 各种 conditioner$~c~$的实现

虽然ci可以是任意复杂的函数,比如神经网络,但每个zi都有一个ci的话,计算量和内存占用太大,解决的办法是参数共享

循环自回归流(Recurrent autoregressive flows)

一种conditioner参数共享的方法是使用循环神经网络RNN/GRU/LSTM来实现,

这种方法的主要缺点是计算不再能并行化,因为计算$s_{i}$必须先计算$s_{i-1}$,当处理高维数据,比如图片、视频时会很慢

掩码自回归流(Masked autoregressive flows)

为了让${h}_{i}$不能依赖$\mathbf{z}_{>=i}$,可以通过将神经网络中这些连接给去掉,方式是给矩阵中这些位置乘上0,就像Transformer中的self-attention在计算softmax时用负无穷mask掉上三角的注意力权重,达到future blind的目的。但该方法的一大缺点是求逆时的计算量是正向计算的D倍

上面的计算过程是:开始不知道z,就随机化,但是z0是任意的,也就是我们可以直接得到h1,然后用z1’与它计算出z1,纠正随机初始化的z的第1个元素,之后以此类推,第D次迭代后,z被完全纠正。Masked conditioner 每次只能使用z(<i),即只能计算出hi。开始直接能得到h1,计算出z1后,再让z1通过conditioner函数得到h2,再用公式计算z2,以此类推。这里一共计算了D次c,而在正向计算时只用一次,求逆时计算代价是正向的D倍,对高维数据来说无法接受。一种解决办法是类似于牛顿法的更新公式(我们要求使 f(z) = z’成立的 z,即求g(z) = f(z) - z’的零点,那么更新公式 z = z - a * g(z)/g’(z) = z - a * (f(z)-z’)/J)

我们使用$\mathbf{z^\prime}$初始化$\mathbf{z_0}$,$f_{\phi}^{-1}(\mathbf{z^\prime})$是上式唯一的不动点,一般$0< \alpha< 2$的情况下$\mathbf{z}_{k}$最终会收敛到某个局部最优点,否则将发散。

掩码自回归流主要适用于不需要求逆或维度较低的情况。

Coupling layers

Coupling layers是将z一分为二,前d个元素原封不动,d+1~D的元素依赖于前d的元素,h1~hd是常数,不依赖于z,hd+1~hD依赖于z<=d,计算公式类似上述几种方法

其雅可比矩阵左上角dxd是单位阵,右下角(D-d)x(D-d)为对角矩阵,其行列式即为右下角对角矩阵的行列式,即对角线乘积。

虽然coupling layers的计算效率更高,但表达能力受限,但我们可以通过堆叠多个coupling layer并在每层对z的元素进行重新排列来解决这个问题,比如第一层前d个元素被原封不动地copy,后D-d个元素通过神经网络,那第二层就原封不动地copy后D-d个元素,把前d个元素通过神经网络,如此交替多层,

3.2 线性流 Linear flows

在自回归流中输出zi’只依赖于z<=i,但元素的排列顺序有时是对模型学习有影响的,一种解决输入元素排列组合的方法是线性流,简单来说就是一个矩阵变换

雅可比矩阵就是W自身,我们需要保证W是可逆的,并且易于求逆和求行列式。

3.3 残差流 Residual flows

$g_{\phi}$是一个输出D维向量的神经网络,我们需要对其加以合适的限制,以让其可逆

3.3.1 Contractive residual flows

构建满足Lipschitz连续条件,且Lipschitz常数小于1的变换F

3.3.2 Residual flows based on the matrix determinant lemma

A为DxD的可逆矩阵,V、W为DxM的矩阵,M<D,有矩阵行列式引理如下

如果A是对角阵的话计算量从左式的$\mathcal{O}(D^{3}+D^{2}M)$降到了$\mathcal{O}(M^{3}+M^{2}D)$

Planar flow

Planar flow是一个单层神经网络,只有一个神经元w

$\sigma$是一个可微的激活函数例如tanh,假设$\sigma^{\prime}$处处大于0,且有上界S,则当$\mathbf{w}^{\top} \mathbf{v}>-\frac{1}{S}$时,雅可比行列式大于0

Sylvester flow

将Planar flow推广到有M个神经元W就是Sylvester flow,$\mathbf{V} \in \mathbb{R}^{D \times M}, \mathbf{W} \in \mathbb{R}^{D \times M}$,$\mathbf{b} \in \mathbb{R}^{M}$,S(z)是对角矩阵,元素是$\sigma^{\prime}\left(\mathbf{W}^{\top} \mathbf{z}+\mathbf{b}\right)$的对角线上的元素

4. 构建标准化流(二):连续时间的变换

假如标准化流有无数个连续的变换,即经过了无穷多个step的transformer,我们把这个step叫做时间,连续时间流可以通过下式构建

神经网络g接受时间t和状态zt作为输入,输出zt关于t的导数,我们可以通过求积分来得到变换的结果

逆变换为

log概率密度的导数可写为

一种近似求迹的方法如下,v是均值为0,协方差矩阵为单位阵的向量

对log概率密度积分,可以得到px与pu的关系式