Normalizing Flows

封面图片来自链接。

Brief Introduction

生成模型的目标是构建参数化模型 $$ 近似真实数据分布 $$，达成该目标的一种思路就是最大化训练样本的对数似然： $$ 然而，直接计算对数似然是十分困难的。为此，VAE 优化的并不是是对数似然本身，而是它的下界 ELBO；GANs 则通过对抗训练的方式隐式地优化对数似然。与这些模型不同的是，Normalizing Flow 模型（下文简称作 Flow 模型）通过巧妙设计网络架构，使之能够直接计算似然函数，从而直接优化对数似然。

Flow 模型最早由 NICE^[1]提出，后续有许多改进，如 Real-NVP^[2], Glow^[3], Flow++^[4]等。值得一提的是，将 Flow 模型时间连续化可得到连续归一化流 CNF^[5]. CNF 可以用 ODE 描述，与 Score-based Models 和 Diffusion Models 有密切的联系，代表工作包括近年的 Rectified Flow^[6]和 Flow Matching^[7]等。不过本文的内容主要集中在离散时间的 NICE, Real-NVP，Glow 和 Flow++ 上。

Change of Variables

Flow 模型希望构建一个可逆非线性变换 $$，将输入 $$ 映射到 $$，使得 $$ 是一个简单分布（例如标准高斯分布）：

根据概率论的知识，有变量替换公式 (change of variables)： $$ 其中 $$ 表示 $$ 的 Jacobian 行列式。于是对数似然为： $$ 因此，如果 Jacobian 行列式是容易计算的，那么对数似然就可以直接计算出来，从而可以优化求解 $$ 了。

进一步地，当训练完成后，由于 $$ 可逆，所以只需要通过简单的 ancestral sampling 就可以生成新的数据： $$ 因此，现在问题的关键就在于如何设计变换 $$，使之满足：1. 可逆；2. Jacobian 行列式容易计算；3. 有足够强大的非线性变换能力。

NICE

NICE 的核心设计是 additive coupling layer. 首先将数据 $$ 划分为两部分 $$，然后定义 $$ 为： $$ 其中 $$ 可以是任意复杂函数（如一个神经网络），如图所示：

在这样的设计下，可逆性非常容易满足： $$ 而 Jacobian 行列式为： $$ 因此前两点要求都得到了满足。但一层 coupling layer 的复杂程度有限，特别是其中有一部分是直接复制的，因此我们考虑堆叠多层 coupling layers. 不过，直接堆叠将得到下面的结构：

可以看见，上分支始终是直接复制前一层，最终也没有发生改变，这不是我们希望的。为此，只需要相邻两层交换上下分支就可以解决这个问题：

堆叠多层 coupling layers 显然并不改变可逆性，且 Jacobian 行列式是每一层 Jacobian 行列式的累乘： $$ 不过 Jacobian 行列式为 1 意味着目前的结构不具备尺度缩放能力。为此，我们在最后添加一个尺度缩放层： $$ 尺度缩放层显然可逆，且 Jacobian 行列式为： $$ 这就完成了整个非线性映射 $$ 的设计。于是，在 NICE 的设计下，模型的对数似然为： $$ 其中先验分布 $$ 可取标准高斯分布： $$ 也可以取标准 logistic 分布（一维下分布函数为 logistic 函数，多维下各维度独立）： $$ 作者倾向于选取 logistic 分布。

Real NVP

Real NVP 是 NICE 作者的续作，将 NICE 中的 additive coupling layer 改进为了 affine coupling layer. 首先将数据 $$ 划分为两部分 $$，然后定义 $$​ 为： $$ 其中 $$ 表示 Hadamard 积（逐元素乘积），$$​ 可以是任意复杂函数（如卷积神经网络），如图所示：

在这样的设计下，容易知道其逆变换为： $$ 而 Jacobian 行列式为： $$ 对比 NICE，Real NVP 的 Jacobian 行列式不再恒为 1，这也是其名称“non-volume preserving”的含义。

数据划分操作可以由 mask 表示，作者设计了针对图像的两种 masking 方式——棋盘式和通道式，如图所示：

棋盘式 masking 以坐标交替的方式做划分，而通道式 masking 将前一半通道和后一般通道划分成两组。另外，上图还同时展示了压缩操作，将 $$ 的 tensor 下采样到 $$ 的大小。设计网络时，作者堆叠使用多层交替棋盘式 masking、压缩操作和多层交替通道式 masking，使得模型具有强大的非线性变换能力。

另外，为了减少计算开销，作者还引入了多尺度架构设计：

其中 stacked coupling layers 表示若干 affine coupling layer 的堆叠，每经过一个 stage 后保留一半输出，另一半继续变换，从而减少参数量和计算量。

Glow

尽管 NICE 和 Real NVP 理论非常吸引人，但是生成效果其实不是很好，而 OpenAI 提出的 Glow 大大提高了 Flow 模型的生成效果。

Glow 的网络架构如下图所示：

可以看见，Glow 整体遵循了与 Real NVP 类似的多尺度架构；另外，在每一个 flow 块中，除了 affine coupling layer 以外，Glow 新引入了两个层—— actnorm 层和可逆 1x1 卷积层。对于输入的图像数据 $$，actnorm 层逐通道地做如下仿射变换： $$ 而可逆 1x1 卷积层可以看作是每个位置上沿通道的全连接层，即： $$ 值得一提的是，由于 $$ 计算复杂度较高，作者将 $$ 进行 LU 分解： $$ 其中 $$ 是一个置换矩阵，$$ 是对角线为 1 的下三角矩阵，$$ 是对角线为 0 的上三角矩阵，这样 $$ 就是对角线上所有 $$ 的乘积。作者初始化时随机采样一个旋转矩阵 $$，计算出 $$ 后固定不动，优化 $$. 此时 $$ 就起到了交叉信息流的作用，因此我们不必再手动交换分支了。

Glow 的三个核心组件的逆变换和对数 Jacobian 行列式如下表所示：

虽然 Glow 的效果非常不错，但是据说所需计算量极其巨大，OpenAI 用了 40 个 GPU，对于 CelebA-HQ 数据集 (256x256) 训练了 4000 个 epochs……

References

1. Dinh, Laurent, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516 (2014). ↩︎
2. Dinh, Laurent, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605.08803 (2016). ↩︎
3. Kingma, Durk P., and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems 31 (2018). ↩︎
4. Ho, Jonathan, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow-based generative models with variational dequantization and architecture design. In International conference on machine learning, pp. 2722-2730. PMLR, 2019. ↩︎
5. Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems 31 (2018). ↩︎
6. Liu, Xingchao, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 (2022). ↩︎
7. Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 (2022). ↩︎