0%

Flow-based Generative Model学习笔记

参考

Flow-based Generative Model - YouTube

How I Understand Flow Matching

Flow-based Deep Generative Models | Lil’Log

[1605.08803] Density estimation using Real NVP

笔记中的截图、部分公式来自以上资源。

笔记

数学回顾

雅可比行列式

上学期,我刷高等数学题目时,常常利用雅可比行列式,对重积分的变量进行换元。

我的理解是把一个坐标系下的小方块变形成另一个坐标系下的小方块时,雅可比行列式告诉我们新的小方块体积是原来小方块的多少倍

我回顾一下:

变量替换定理

一个简单的引入:因为区域包含的概率为1,所以概率密度变为了一半:

1

由于$\mathbf{x}=G(\mathbf{z})$,所以对任意的概率密度函数,在$\mathbf{z’} \sim \mathbf{z’}+\Delta \mathbf{z}$这个很小的范围对应的$\mathbf{x’} \sim \mathbf{x’}+\Delta \mathbf{x}$里,也是和上例一样的小矩形,所以有:

2

把右上角的式子变形:(注意要加绝对值)

3

一维的情况就是这样。引申到二维:

4

同理,有

5

变形,

6

放入行列式,

7

这是偏导数,所以转置后就是雅可比行列式,

8

变形,

9

这就是变量替换公式。

Normalizing Flows

要找到 loss 函数,还要看最大似然。这一次模型的目的是,将一个标准正态分布的 z 通过 generator 对应到我们想训练出来的分布上,如图,

10

因此,最大似然应该是所有我们能提供的训练数据的样本在 generator 生成的分布下的联合概率。取个负对数似然就是 loss 函数了,是:

要训练 generator 的参数,就要:

其中 $x^i$ 是我们提供的原始数据。现在我们考虑如何计算。刚刚有:

代入$\mathbf{z}^i=G^{-1}(\mathbf{x}^i)$,取对数,有:

这就是我们要 maximize 的式子。为了能够计算,能看出来我们需要:

  • 能知道 $G^{-1}$,即 G 是可逆的(双射函数);
  • 能(较高效地)计算 $\det(J_{G})$。

由此能看出来,G 有很多限制条件,不能用任意复杂的神经网络去建立,需要设计特定的神经网络结构。为增加模型的表达能力,我们常用多个结构简单、可逆的 G 模块逐层组合构造最终的 G。

11

这么多 G,loss 函数该如何计算呢?刚才知道:

又知道:

所以有:

12

这就可以递推到$\mathbf{z}_0$了,就是:

将多个可逆变换串联起来形成的分布变换路径称为归一化流(normalizing flow)。它的 loss 函数可以通过上面精确计算对数似然然后取相反数得到。之后,人们开始各种设计 G。比如RealNVP、NICE、Glow、Autoregressive flow等等。我们接下来会讲解。

训练

先考虑一个 G 的例子作为引入:

观察,发现只有 $G^{-1}$。我们会训练$G^{-1}$,用 $G$ 生成(输入 z 产生 x)。

我们来看上面这个式子。

第一项

我们希望它越大越好。我们知道 $\pi$ 是正态分布,那只有接近零向量的时候,输出概率最大。如果只看这一项,机器会学着让 $G^{-1}$ 输出接近零向量。

第二项

如果刚刚 $G^{-1}$ 输出的 $z^i$ 都是零向量,那么$J_{G^{-1}}$就是零矩阵。行列式的值是 0。代入对数,就是负无穷。与我们 maximize 的目标相悖。

Coupling Layer(耦合层)

它被用在 NICE,RealNVP 和 Glow 上。

RealNVP

13

RealNVP(Real-valued Non-Volume Preserving)通过堆叠一系列可逆双射变换函数来实现归一化流。在每个双射$f: \mathbf{x} \mapsto \mathbf{y}$中,仿射耦合层(Affine Coupling Layer)的输入维度分为两个部分:

  • 前 d 个维度保持不变;
  • 从第 d + 1 到第 D 个维度,会经历放射变换,具体就是 scale 和 shift。scale 和 shift 的参数怎么来?它们都是前 d 个维度的函数。前 d 个维度决定了第 d + 1 到第 D 个维度会怎么变换。

就是:

其中,s、t 是 scale 和 translation 函数。

这显然是符合我们的要求的,就是:

  • 变换是可逆的(双射函数);
  • 能(较高效地)计算雅可比行列式。

它的雅可比行列式左上角的 d × d 区域是单位矩阵,右上角是0,右下角是对角矩阵。

由于 t 与 s 不用求逆,所以它们可以用任意复杂深度神经网络得到。

不过这样的操作会让一些维度保持不变。为了确保所有输入都有机会被改变,可以反转每一层中的顺寻,以便不同的组件保持不变。

用这种交替模式,在一个转换层中保持相同的单元集在下一个转换层中总是被修改。

但深层堆叠仿射耦合层带来了两个问题:

  • 信息传播不均衡:有些通道(维度)可能在浅层被更新了很多次,而另一些通道直到很深的地方才开始参与变换,容易导致梯度传播不均衡、训练困难。

  • 前层输出分布漂移(distribution shift):每层的输出分布会逐步变化,导致后面的网络看到的数据分布不稳定,使训练变得不稳定、收敛慢。

批量归一化(Batch Normalization)可以这些问题。

  • 因为每一层的输出会被标准化为均值为0、方差为1,这让下一层的输入始终处于一个稳定的尺度范围,防止分布漂移,便于优化器更有效地更新参数。
  • 缓解更新不均问题(梯度稳定)。让被耦合层跳过不参与仿射变换的维度保持标准化的状态。

RealNVP 可以在一个多尺度结构(multi-scale architecture)中运行,这里不展开。多尺度架构用了一些方法用于处理图像,例如像素掩码策略(Spatial Checkerboard Pattern Masking),像国际象棋的黑白格子一样划分像素,黑格子上的像素值保持不变,白格子上的像素值进行仿射变换,下一层再反过来处理黑格子,保证信息全部都被处理过,避免总是处理图像的同一部分;Channel-wise Masking(通道掩码)选择前一半通道作为条件输入,对后一半通道进行仿射变换,下一层再反过来,适合低分辨率。

NICE(非线性独立成分估计)

NICE (Non-linear Independent Component Estimation; Dinh, et al. 2015) 就是去掉了 scale 的RealNVP:

Glow(使用 1 × 1 Convolution(1 × 1 卷积))

这被用在Glow (Kingma and Dhariwal, 2018) 里。

Glow 用一个机器学习出来的 3 × 3 的矩阵 W,去乘上 1 × 1 × 3 的一个点的通道向量,得到输出。每个点都操作后,就起到了打乱(shuffle)的作用。

例如,如果 W 被学习成下图右下角这样,就是对顺序的打乱:

14

W 需要是可逆的。雅可比矩阵就是 W 本身:

15

全部的 z 与 x 的 generator 的雅可比行列式就是$(\det(W))^{d×d}$:

16

Glow 中的一层Flow Step顺序是:

  • Actnorm(不用batch norm)
    • 对每个通道做仿射归一化(Affine normalization),$s$ 和 $b$ 是可学习参数,初始化时用 batch statistics 设置。:
  • Invertible 1×1 Convolution
    • 可以看成是全通道的线性变换。
  • Affine Coupling Layer
    • 与 RealNVP 一样进行仿射变换。

Glow 中的 Multi-scale architecture 是:

模块 作用 具体
🟨 Squeeze 把空间维度挤到通道维,减少空间分辨率、增加通道数 2×2 的 patch 变成一个 vector,提升表达力
🟩 Flow Step × K 应用多次 Flow Step Actnorm、Invertible 1×1 Convolution、Affine Coupling Layer
🟦 Split 分出一部分 latent variable $z_i$,剩下的继续变换 降低后续计算量。把通道维切成一半,然后:一半作为 latent variable $z_i$ 被作为准备的输出;另一半继续往下传递,进入下一层
🔁 重复以上结构 L 次 多层堆叠,每一层输出一个 $z_i$ 最后输出 $[z_1, …, z_L]$ 作为最终 latent 表达

18

Autoregressive flow(自回归流)

17

如图,每一维 $x_i$ 的生成,都依赖于前面所有的 $z_1, …, z_{i-1}$ 输入 Conditioner 来输出的一个隐藏向量 $h_i$。$\tau$函数是可逆的。

符合要求:可逆、可算雅可比行列式。

Residual flows(残差流)

19

Residual flows 的结构如图:

20

判定一下必要的条件:

G 可逆:

已知:

可逆就是,给出 x,能否得到 z?

虽然 u 是神经网络,但是如果 u 函数是Contractive Map(压缩映射),它就可以使可逆的。

压缩映射的定义是:对于任意两个输入 $\mathbf{z}_a, \mathbf{z}_b$,都有:

如图:

21

如果 u 是Contractive Map,根据巴拿赫不动点定理(Banach fixed-point theorem),存在一个不动点:

如果 u 是Contractive Map,那么 $x-u(\hat{z})$也是Contractive Map,其中$\mathbf{x}$是我们已经给出的,我们正在寻找符合

的$\mathbf{z}$。

由于 $x-u(\hat{z})$是 Contractive Map,所以存在一个不动点:

移项,得

所以$\mathbf{z}_*$就是我们要找的。

根据巴拿赫不动点定理,$\mathbf{z}_*$是唯一的,所以这个

是可逆的。

可以用迭代算法来求得$\mathbf{z}_*$,就是 algorithm for inversion。由于 f(z)=x−u(z) 是一个压缩映射,我们可以从任意初始猜测 $z_0$ 开始,通过重复应用 f 来逼近不动点:

这个迭代过程会收敛到唯一的解$z_*$。

雅可比行列式:

残差流的变换是

它的雅可比矩阵是

其中,$I$是单位矩阵,$\frac{\partial u(\mathbf{z})}{\partial \mathbf{z}}$是$u(\mathbf{z})$的雅可比矩阵。

接下来,经过一些数学推导,可以用无穷级数计算雅可比行列式,

接下来使用蒙特卡洛方法(Hutchinson’s Trace Estimator)来估计迹