基于流的(Flow-based)生成模型简介
生成任务
我们先回顾一下所谓的生成任务,究竟是做什么事情。我们认为,世界上所有的图片,是符合某种分布 pdata(x)p_{data}(x)pdata(x) 的。当然,这个分布肯定是个极其复杂的分布。而我们有一堆图片 x1,x2,…,xm{x_1,x_2,\dots,x_m}x1,x2,…,xm ,则可以认为是从这个分布中采样出来的 mmm 个样本。我们通过训练希望得到一个生成器网络 GGG ,该网络定义了一个分布 pG(x)p_G(x)pG(x),能够做到输入一个从正态分布 π(z)\pi(z)π(z) 中采样出来的 zzz ,输出一张看起来像真实世界的图片 x=G(z)x=G(z)x=G(z) 。我们希望 pG(x)p_G(x)pG(x) 与真实的数据分布 pdata(x)p_{data}(x)pdata(x) 越接近越好。
从概率模型的角度来看,想要做到上面说的这件事情,就要通过最大化对数似然,来优化生成器 GGG 的参数:
G∗=argmaxG∑i=1mlogpG(xi)G^*=\arg\max_{G}\sum_{i=1}^m\log p_G(x_i) G∗=argGmaxi=1∑mlogpG(xi)
可以证明,最大化这个对数似然,就相当于最小化生成器分布 pG(x)p_G(x)pG(x) 与目标分布 pdata(x)p_{data}(x)pdata(x) 的 KL散度,即让这两个分布尽量接近:
G∗≈argminGKL(pdata∣∣pG)G^*\approx\arg\min_GKL(p_{data}||p_G) G∗≈argGminKL(pdata∣∣pG)
另一种著名的生成模型 VAE 也是一种概率模型,它优化的是这个对数似然的证据下界(ELBO)。而本文要介绍的流模型,可以直接优化这个对数似然本身。当然为了能够直接对目标进行优化,流模型的生成器 GGG 本身会有一些数学上的限制。
数学基础
为了搞懂流模型,我们需要回顾一些数学基础,主要包括:雅可比矩阵(Jacobian Matrix)、行列式(deteminant)、变量变换定理(Change of Variable Theorem)。
雅可比矩阵
定义
假设某函数 f:Rn→Rm\mathbf{f}:\mathbb{R}^n\rightarrow\mathbb{R}^mf:Rn→Rm, 从 x∈Rn\mathbf{x}\in\mathbb{R}^nx∈Rn 映射到向量 f(x)∈Rm\mathbf{f}(\mathbf{x})\in\mathbb{R}^mf(x)∈Rm, 其雅可比矩阵是一个 m×nm\times nm×n 的矩阵,换句话讲也就是从 Rn\mathbb{R}^nRn 到 Rm\mathbb{R}^mRm 的线性映射。其重要意义在于它表现了一个多变量向量函数的最佳线性逼近。因此,雅可比矩阵类似于单变量函数的导数。
此函数 f\mathbf{f}f 的雅可比矩阵 J\mathbf{J}J 为 m×nm\times nm×n 的矩阵,一般由以下方式定义:
J=[∂f∂x1,∂f∂x2,…,∂f∂xn]=[∂f1∂x1⋯∂f1∂xn∂f2∂x1⋯∂f2∂xn⋮⋱⋮∂fm∂x1⋯∂fm∂xn]m×n\mathbf{J}=[\frac{\partial\mathbf{f}}{\partial x_1},\frac{\partial\mathbf{f}}{\partial x_2},\dots,\frac{\partial\mathbf{f}}{\partial x_n}]= \begin{bmatrix} {\frac{\partial f_1}{\partial x_1}}&{\cdots}&{\frac{\partial f_1}{\partial x_n}}\\{\frac{\partial f_2}{\partial x_1}}&{\cdots}&{\frac{\partial f_2}{\partial x_n}}\\{\vdots}&{\ddots}&{\vdots}\\{\frac{\partial f_m}{\partial x_1}}&{\cdots}&{\frac{\partial f_m}{\partial x_n}}\\ \end{bmatrix}_{m\times n} J=[∂x1∂f,∂x2∂f,…,∂xn∂f]=⎣⎡∂x1∂f1∂x1∂f2⋮∂x1∂fm⋯⋯⋱⋯∂xn∂f1∂xn∂f2⋮∂xn∂fm⎦⎤m×n
通俗来讲,就是某个 mmm 维输入,nnn 维输出的函数,把它的各个输出变量对各个输入变量的偏微分求出来,然后按照上面的定义排好,就组成了雅可比矩阵。这里我们暂时只关心输入输出同维度的情况,即 m=nm=nm=n。
性质
如果有 x=f(z)x=f(z)x=f(z) ,其雅可比矩阵记为 JfJ_fJf ,其反函数 z=f−1(x)z=f^{-1}(x)z=f−1(x) ,其雅可比矩阵记为 Jf−1J_{f^{-1}}Jf−1 ,则有:JfJf−1=IJ_fJ_{f^{-1}}=IJfJf−1=I 。即如果两个函数互为反函数,则他们的雅可比矩阵互逆。
行列式
行列式(deteminant)是线性代数中的一个基本概念,矩阵 AAA 的行列式记作 ∣A∣|A|∣A∣ 或 det(A)det(A)det(A)。相信大家都还有印象,即使不记得复杂的行列式计算公式,也应该知道他就是对一个矩阵进行一顿计算,得到一个标量值。
性质
这里回顾几个行列式的常用性质,后面在介绍流模型时也会用到:
- 矩阵 AAA 中某行(或列)用同一数 kkk 乘,其行列式值也乘 kkk。
- 某个矩阵的行列式等于其转置的行列式, ∣A∣=∣AT∣|A|=|A^T|∣A∣=∣AT∣ 。
物理意义
这里还要着重强调的一点是:行列式可以看做是有向面积或体积的概念在一般的欧几里得空间中的推广。或者说,在 nnn 维欧几里得空间中,行列式描述的是一个线性变换对“体积”所造成的影响。
行列式的物理意义就是高维空间中的体积。
变量变换定理
变量变换定理是一个有效描述长度、面积、体积和广义n维体积(内容)如何被可微函数所扭曲的定理。特别是,变量变换定理将弄清内容扭曲的整个问题简化为理解无穷小的扭曲,即由线性映射的行列式所给出的导数(一个线性映射)的扭曲。
如果我们有 π(z)\pi(z)π(z) 和 p(x)p(x)p(x) ,并且知道二者之间的关系 x=f(z)x=f(z)x=f(z) ,那么能不能写出 π(z)\pi(z)π(z) 和 p(x)p(x)p(x) 之间的关系?是可以的!
均匀分布
先举个简单的例子:考虑一个最简单的均匀分布:π(z)\pi(z)π(z) ,它的取值范围是 [0,1][0,1][0,1] 。而它作为一个分布,有 ∫π(z)dz=1\int\pi(z)dz=1∫π(z)dz=1 ,那么自然地, π(z)=1\pi(z)=1π(z)=1 ,即图中蓝色方块的高为1。我们通过变换函数为 x=f(z)=2z+1x=f(z)=2z+1x=f(z)=2z+1 得到 p(x)p(x)p(x) ,p(x)p(x)p(x) 同样是均匀分布,通过 ∫p(x)dx=1\int p(x)dx=1∫p(x)dx=1 ,得到 p(x)=12p(x)=\frac{1}{2}p(x)=21 ,即图中绿色方块的高为 12\frac{1}{2}21 。这样,我们就得到了两个均匀分布之间的关系:p(x′)=12π(z′)p(x')=\frac{1}{2}\pi(z')p(x′)=21π(z′) 。
一般情况
接下来我们来看一般情况。我们想要考察两个一般的分布 π(z)\pi(z)π(z) 与 p(x)p(x)p(x) 在经过变化 x=f(z)x=f(z)x=f(z) 变化前后的关系。我们可以将 zzz 移动 一个很小的距离 Δz\Delta zΔz ,对应的, xxx 也移动了一个很小的距离 Δx\Delta xΔx ,这时可以认为在 Δz\Delta zΔz 和 Δx\Delta xΔx 范围内都是一个均匀分布。蓝色、绿色两个方块的面积是一样的,可以按照我们上面均匀分布的方式来计算,得到:
p(x′)=π(z′)∣dzdx∣p(x')=\pi(z')|\frac{dz}{dx}| p(x′)=π(z′)∣dxdz∣
即变换前后的分布差了一个微分,而这个微分在我们知道 fff 的情况下是可以算出来的。这里要加绝对值是因为 z=f(x)z=f(x)z=f(x) 的变换可能是相反方向的,但这不影响我们的计算方式。
高维的情况
上面介绍的都是一维的情况,那么在高维的情况下是怎样的呢?
以二维的情况为例。我们还是对 zzz 在两个维度 z1,z2z_1,z_2z1,z2 上增加一个很小的距离 Δz1,Δz2\Delta z_1,\Delta z_2Δz1,Δz2 ,而对应的 x1x_1x1 和 x2x_2x2 ,Δxij\Delta x_{ij}Δxij 表示 ziz_izi 改变时,xjx_jxj 的改变量。与之前类似的,图中蓝色、绿色四边形所对应的体积(注意这里三维的概率密度函数没有画出来,它也是均匀的)是相同的。有:
p(x′)∣det([Δx11Δx21Δx12Δx22])∣=π(z′)Δz1Δz2p(x')| det( \begin{bmatrix} {\Delta x_{11}}&{\Delta x_{21}}\\ {\Delta x_{12}}&{\Delta x_{22}}\\ \end{bmatrix}) | =\pi(z')\Delta z_1\Delta z_2 p(x′)∣det([Δx11Δx12Δx21Δx22])∣=π(z′)Δz1Δz2
然后就是对这个式子进行一系列变形:
p(x′)∣det([Δx11Δx21Δx12Δx22])∣=π(z′)Δz1Δz2p(x′)∣1Δz1Δz2det([Δx11Δx21Δx12Δx22])∣=π(z′)p(x′)∣det([Δx11/Δz1Δx21/Δz1Δx12/Δz2Δx22/Δz2])∣=π(z′)p(x′)∣det([Δx11/Δz1Δx21/Δz1Δx12/Δz2Δx22/Δz2]T)∣=π(z′)p(x′)=π(z′)∣1det(Jf)∣p(x')| det( \begin{bmatrix} {\Delta x_{11}}&{\Delta x_{21}}\\ {\Delta x_{12}}&{\Delta x_{22}}\\ \end{bmatrix}) | =\pi(z')\Delta z_1\Delta z_2\\ p(x')| \frac{1}{\Delta z_1\Delta z_2} det( \begin{bmatrix} {\Delta x_{11}}&{\Delta x_{21}}\\ {\Delta x_{12}}&{\Delta x_{22}}\\ \end{bmatrix}) | =\pi(z')\\ p(x')| det( \begin{bmatrix} {\Delta x_{11}/\Delta z_1}&{\Delta x_{21}/\Delta z_1}\\ {\Delta x_{12}/\Delta z_2}&{\Delta x_{22}/\Delta z_2}\\ \end{bmatrix}) | =\pi(z')\\ p(x')| det( \begin{bmatrix} {\Delta x_{11}/\Delta z_1}&{\Delta x_{21}/\Delta z_1}\\ {\Delta x_{12}/\Delta z_2}&{\Delta x_{22}/\Delta z_2}\\ \end{bmatrix}^T) | =\pi(z')\\ p(x')=\pi(z')|\frac{1}{det(J_f)}| p(x′)∣det([Δx11Δx12Δx21Δx22])∣=π(z′)Δz1Δz2p(x′)∣Δz1Δz21det([Δx11Δx12Δx21Δx22])∣=π(z′)p(x′)∣det([Δx11/Δz1Δx12/Δz2Δx21/Δz1Δx22/Δz2])∣=π(z′)p(x′)∣det([Δx11/Δz1Δx12/Δz2Δx21/Δz1Δx22/Δz2]T)∣=π(z′)p(x′)=π(z′)∣det(Jf)1∣
结论
总之,在数学基础部分我们要记住的一个结论就是下面的式子:
p(x′)∣det(Jf)∣=π(z′)p(x′)=π(z′)∣det(Jf−1)∣p(x')|det(J_f)|=\pi(z')\\ p(x')=\pi(z')|det(J_{f^{-1}})| p(x′)∣det(Jf)∣=π(z′)p(x′)=π(z′)∣det(Jf−1)∣
流模型公式推导
我们回到生成的任务上来,上面提到,生成任务就是要通过最大化对数似然优化生成器:
G∗=argmaxG∑i=1mlogpG(xi)G^*=\arg\max_{G}\sum_{i=1}^m\log p_G(x_i) G∗=argGmaxi=1∑mlogpG(xi)
而根据上面数学基础得到的结论,有:
pG(xi)=π(zi)∣det(JG−1)∣,zi=G−1(xi)p_G(x_i)=\pi(z_i)|det(J_{G^{-1}})|,\ \ \ \ z_i=G^{-1}(x_i) pG(xi)=π(zi)∣det(JG−1)∣, zi=G−1(xi)
则对数似然:
logpG(xi)=logπ(G−1(xi))+log∣det(JG−1)∣\log p_G(x_i)=\log \pi(G^{-1}(x_i))+\log |det(J_{G^{-1}})| logpG(xi)=logπ(G−1(xi))+log∣det(JG−1)∣
要训练一个好的生成器,只要最大化上面这个式子,就可以了。
现在的问题就是怎么把这个式子算出来,具体来说,这个式子计算的关键在以下两点:
- 如何计算 det(JG)det(J_G)det(JG)
- 如何计算 G−1G^{-1}G−1
我们设计的生成器网络 GGG 需要满足上面这两个条件,这就是前面提到的流模型生成器数学上的限制。而 VAE、GAN 等生成模型则可以是任意的神经网络,没有限制。
另外要提一点,流模型的输入输出的尺寸必须是一致的。这是因为如果想要 GGG 可逆,它的输入输出维度一致是一个必要条件(非方阵不可能可逆)。比如要生成 100×100×3100\times 100\times 3100×100×3 的图像,那输入的随机噪声也是 100×100×3100\times 100\times 3100×100×3 的。这与 VAE、GAN 等生成模型很不一样,这些生成模型的输入维度通常远小于输出维度。
流模型的网络设计
多层设计
前面提到流模型的生成器 GGG 是要收到数学上的一些限制的,这导致 GGG 本身的表达能力可能是不足的。所以,一般需要堆叠多层网络来得到一个生成器,这也是 “流模型” 这个名称的由来。
不过虽然堆叠了很多层,在公式上也没有什么复杂的。无非就是把一堆 GiG_iGi 连乘起来,通过 log\loglog 之后,又变成连加。
实际训练
我们再观察一下要最大化的对数似然:
logpG(xi)=logπ(G−1(xi))+log∣det(JG−1)∣\log p_G(x_i)=\log \pi(G^{-1}(x_i))+\log |det(J_{G^{-1}})| logpG(xi)=logπ(G−1(xi))+log∣det(JG−1)∣
发现整个式子只与 G−1G^{-1}G−1 有关。实际上,我们在训练过程中,就是通过样本 xix_ixi 训练 G−1G^{-1}G−1 ,然后再推理的时候反过来用 GGG 来根据随机噪声进行生成就可以了。
观察第一项,logπ(zi)=logπ(G−1(xi))\log\pi(z_i)=\log\pi(G^{-1}(x_i))logπ(zi)=logπ(G−1(xi)) ,我们知道 π\piπ 是正态分布,因此要最大化这一项,最好让 ziz_izi 是零向量,这样能取到最大值。然而,如果真的将 ziz_izi 取成全0了,那么微分也就是0,从而雅克比行列式 det(JG−1)det(J_{G^{-1}})det(JG−1) 也会是 0,这样第二项 log∣det(JG−1)∣\log|det(J_{G^{-1}})|log∣det(JG−1)∣ 就会是负无穷,这样整个式子没办法最大。即,式子的两项会有一个 tradeoff,前一项让 ziz_izi 尽量靠近零向量,而后一项又会让 ziz_izi 不要是全零。
coupling layer
在经典的流模型方法 NICE 和 RealNVP 中,使用的具体网络结构是 coupling layer。
做法
coupling layer 的输入是 DDD 维的 z1,…,zDz_1,\dots,z_Dz1,…,zD ,输出也是 DDD 维的 x1,…,xDx_1,\dots,x_Dx1,…,xD 。将输入输出分别拆成两组,前 ddd 维一组,后面 D−dD-dD−d 维是另一组。
输出 xix_ixi 的前 ddd 个元素就是输入 ziz_izi 的前 ddd 个元素直接拷贝过来。
而输出 xix_ixi 的后 D−dD-dD−d 个元素是这样计算的:输入的前 ddd 个元素先通过两个函数 F,HF,HF,H 得到 D−dD-dD−d 维的 βi\beta_iβi 和 γi\gamma_iγi 。然后与输入的后 D−dD-dD−d 个元素通过一个线性计算 xi>d=βizi+γix_{i>d}=\beta_iz_i+\gamma_ixi>d=βizi+γi 得到输出 xix_ixi 的后 D−dD-dD−d 个元素。注意这里的 F,HF,HF,H 可以是任意的函数、任意的网络,不需要满足前面提到的限制。
第一次接触 coupling layer 的读者可能会很奇怪,这个层花里胡哨一顿操作是要干嘛?别急,下面我们就通过分析 coupling layer 与流模型生成器网络的两个数学限制,来解释这个层设计的巧妙之处。
如何计算G−1G^{-1}G−1
首先我们来看 coupling layer 如何计算它的 inverse: z=G−1(x)z=G^{-1}(x)z=G−1(x) 。即如何根据 xix_ixi 得到 ziz_izi 。
对于前 ddd 维,这是很容易的,因为 xix_ixi 和 ziz_izi 的前 ddd 维是全等的,因此也是直接拷贝回来即可。
对于后 D−dD-dD−d 维,首先我们已经得到 ziz_izi 的前 ddd 维了,再通过函数 F,HF,HF,H 即可计算出 βi\beta_iβi 和 γi\gamma_iγi (这里可以看到 F,HF,HF,H 无需是可逆的),这样就可以通过反转之前的线性计算,即通过 zi>d=xi−γiβiz_{i>d}=\frac{x_i-\gamma_i}{\beta_i}zi>d=βixi−γi,得到 ziz_izi 的后 D−dD-dD−d 维度。
如何计算 det(JG−1)det(J_{G^{-1}})det(JG−1)
然后我们再来看 coupling layer 如何计算雅可比行列式 det(JG−1)det(J_{G^{-1}})det(JG−1) 。
在计算 coupling layer 时,我们将输入输出都分成了两个部分,现在计算雅可比行列式,就分成四块来看。
- 左上角的一块就是单位矩阵 III, 因为这部分输出是完全拷贝输入的;
- 右上角的部分是零矩阵,因为输出的前 ddd 维,与输入的后 D−dD-dD−d 维完全没关系;
- 左下角的矩阵可能会非常复杂,因为 F,HF,HF,H 是任意的网络,但是由于上半部分是单位矩阵和零矩阵,因此这部分与整个矩阵行列式的计算无关,整个矩阵的行列式只取决于右下角的部分
- 右下角的部分是一个对角矩阵,因为这部分是通过线性计算 xi>d=βizi+γix_{i>d}=\beta_iz_i+\gamma_ixi>d=βizi+γi 得到的,只与相同下标的有关,并且,值就是 βi\beta_iβi
因此,整个矩阵的行列式:
det(JG−1)=βd+1βd+2…βDdet(J_{G^{-1}})=\beta_{d+1}\beta_{d+2}\dots\beta_D det(JG−1)=βd+1βd+2…βD
堆叠coupling layer
前面提到过,流模型要通过堆叠多层网络来强化整个生成器的表达能力。但是,在堆叠 coupling layer 的时候要注意一点:每次 copy 的那一半要进行交换,不能每次都 copy 上半部分。如果每次 copy 的都是上半部分,那到最后生成的图像的上半部分会和输入是完全一样的噪声。因此,需要在堆叠时注意更换 copy 的半部。
另外,对于图像生成任务来说,区分前后半部的方法通常有两种:按空间和按通道。分别是某些像素坐标 (hi,wi)(h_i,w_i)(hi,wi) 做前半部,另一些做后半部;和某些通道做前半部,另一些做后半部。当然,也可以选择两种混用。
1x1 Conv
近年一篇较新的流模型的方法是来自 OpenAI 的 GLOW。这篇工作提出的方法是使用 1x1 的卷积来作为流模型的层。具体来说,图像的每一个像素位置 1x3 个元素会与一个 3x3 的矩阵 WWW 做乘积,得到输出的 1x3 。在这个过程中,WWW 是要学习的参数,模型通过学习 WWW 的参数,可能可以自己对通道进行 shuffle(如下图例子所示)。这样与 coupling layer 配合时就不用手动选择 copy 的半部,而是可以由模型通过学习自己决定通道交换的方法。
当然,作为流模型的一层,1x1 Conv 也需要满足前面提到的两个限制。对于雅可比行列式 det(JG−1)det(J_{G^{-1}})det(JG−1) 的计算,1x1 的卷积层是完全可解析的,即可以直接通过公式计算出来的,就是 det(W)d×ddet(W)^{d\times d}det(W)d×d。但是对于 G−1G^{-1}G−1 可逆性,原文并没有给出保证。虽然原文在实现时将 WWW 初始化为可逆的,并且 3x3 的矩阵不可逆的概率也比较小(行列式恰好等于零),但是毕竟没有在原理上保证可逆,在一些网友训练的过程中也出现了 NaN 的情况。
最后推荐一下 openai 写的关于 glow 的 demo 和 blog:
https://openai.com/blog/glow/
有很多好玩的应用。比如通过统计计算 “笑的人脸” 和 “不笑的人脸” 的差,再加到一张不笑的人脸上,让他笑起来等。
Ref
- realnvp paper
- nice paper
- glow paper
- Flow-based Generative Model
- 雅可比矩阵-维基百科