Normalizing Flows

封面图片来自链接。

Brief Introduction

生成模型的目标是构建参数化模型 \(p_\theta(x)\) 近似真实数据分布 \(p_\text{data}(x)\)，达成该目标的一种思路就是最大化训练样本的对数似然： \[ \max_\theta\;\mathbb E_{p_\text{data}(x)}[\log p_\theta(x)] \] 然而，直接计算对数似然是十分困难的。为此，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 模型希望构建一个可逆非线性变换 \(f_\theta\)，将输入 \(x\) 映射到 \(z\)，使得 \(p_Z(z)\) 是一个简单分布（例如标准高斯分布）：

根据概率论的知识，有变量替换公式 (change of variables)： \[ p_\theta(x)=p_Z(f_\theta(x))\left|\det\left(\frac{\partial f_\theta(x)}{\partial x}\right)\right| \] 其中 \(\det(\partial f_\theta(x)/\partial x)\) 表示 \(f_\theta\) 的 Jacobian 行列式。于是对数似然为： \[ \log p_\theta(x)=\log p_Z(f_\theta(x))+\log\left|\det\left(\frac{\partial f_\theta(x)}{\partial x}\right)\right| \] 因此，如果 Jacobian 行列式是容易计算的，那么对数似然就可以直接计算出来，从而可以优化求解 \(f_\theta\) 了。

进一步地，当训练完成后，由于 \(f_\theta\) 可逆，所以只需要通过简单的 ancestral sampling 就可以生成新的数据： \[ z\sim p(z),\quad x=f_\theta^{-1}(z) \] 因此，现在问题的关键就在于如何设计变换 \(f_\theta\)，使之满足：1. 可逆；2. Jacobian 行列式容易计算；3. 有足够强大的非线性变换能力。

NICE

NICE 的核心设计是 additive coupling layer. 首先将数据 \(x\) 划分为两部分 \(x_1\in\mathbb R^{d_1},x_2\in\mathbb R^{d_2}\)，然后定义 \(z=(z_1,z_2)\) 为： \[ \begin{align} &z_1=x_1\\ &z_2=m_\theta(x_1)+x_2 \end{align} \] 其中 \(m_\theta:\mathbb R^{d_1}\to\mathbb R^{d_1}\) 可以是任意复杂函数（如一个神经网络），如图所示：

在这样的设计下，可逆性非常容易满足： \[ \begin{align} &x_1=z_1\\ &x_2=z_2-m_\theta(z_1) \end{align} \] 而 Jacobian 行列式为： \[ \det\left(\frac{\partial (z_1,z_2)}{\partial (x_1,x_2)}\right)= \det\begin{pmatrix} \frac{\partial z_1}{\partial x_1}&\frac{\partial z_1}{\partial x_2}\\ \frac{\partial z_2}{\partial x_1}&\frac{\partial z_2}{\partial x_2} \end{pmatrix}= \det\begin{pmatrix} I_{d_1}&0\\ \frac{\partial{z_2}}{\partial x_1}&I_{d_2} \end{pmatrix}=1 \] 因此前两点要求都得到了满足。但一层 coupling layer 的复杂程度有限，特别是其中有一部分是直接复制的，因此我们考虑堆叠多层 coupling layers. 不过，直接堆叠将得到下面的结构：

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

堆叠多层 coupling layers 显然并不改变可逆性，且 Jacobian 行列式是每一层 Jacobian 行列式的累乘： \[ \det\left(\frac{\partial(z_1^{(2)},z_2^{(2)})}{\partial(x_1,x_2)}\right)=\det\left(\frac{\partial(z_1^{(2)},z_2^{(2)})}{\partial(z_1^{(1)},z_2^{(1)})}\cdot\frac{\partial(z_1^{(1)},z_2^{(1)})}{\partial(x_1,x_2)}\right)=\det\left(\frac{\partial(z_1^{(2)},z_2^{(2)})}{\partial(z_1^{(1)},z_2^{(1)})}\right)\det\left(\frac{\partial(z_1^{(1)},z_2^{(1)})}{\partial(x_1,x_2)}\right)=1\cdot 1=1 \] 不过 Jacobian 行列式为 1 意味着目前的结构不具备尺度缩放能力。为此，我们在最后添加一个尺度缩放层： \[ z=Sz^{(n)}=\text{diag}(s)z^{(n)}=\begin{bmatrix}S_{11}z_1^{(n)}&S_{22}z_2^{(n)}&\cdots&S_{dd}z_d^{(n)}\end{bmatrix}^T \] 尺度缩放层显然可逆，且 Jacobian 行列式为： \[ \det\left(\frac{\partial z}{\partial z^{(n)}}\right)= \det\begin{pmatrix} S_{11}&0&\cdots&0\\ 0&S_{22}&\cdots&0\\ \vdots&\vdots&\ddots&\vdots\\ 0&0&\cdots&S_{dd} \end{pmatrix}=\prod_{i=1}^d S_{ii} \] 这就完成了整个非线性映射 \(f_\theta\) 的设计。于是，在 NICE 的设计下，模型的对数似然为： \[ \log p_\theta(x)=\log p_Z(f_\theta(x))+\sum_{i=1}^d\log(|S_{ii}|) \] 其中先验分布 \(p_Z(\cdot)\) 可取标准高斯分布： \[ \log p_Z(z)=-\frac{1}{2}\Vert z\Vert^2-\frac{d}{2}\log(2\pi) \] 也可以取标准 logistic 分布（一维下分布函数为 logistic 函数，多维下各维度独立）： \[ \log p_Z(z)=-\sum_{i=1}^d\big[\log(1+\exp(z_i))+\log(1+\exp(-z_i))\big] \] 作者倾向于选取 logistic 分布。

Real NVP

Real NVP 是 NICE 作者的续作，将 NICE 中的 additive coupling layer 改进为了 affine coupling layer. 首先将数据 \(x\) 划分为两部分 \(x_1\in\mathbb R^{d_1},x_2\in\mathbb R^{d_2}\)，然后定义 \(z=(z_1,z_2)\)​ 为： \[ \begin{align} &z_1=x_1\\ &z_2=x_2\odot\exp(s_\theta(x_1))+t_\theta(x_1) \end{align} \] 其中 \(\odot\) 表示 Hadamard 积（逐元素乘积），\(s_\theta,t_\theta:\mathbb R^{d_1}\to\mathbb R^{d_1}\)​ 可以是任意复杂函数（如卷积神经网络），如图所示：

在这样的设计下，容易知道其逆变换为： \[ \begin{align} &x_1=z_1\\ &x_2=(z_2-t_\theta(z_1))\odot\exp(-s_\theta(z_1)) \end{align} \] 而 Jacobian 行列式为： \[ \det\left(\frac{\partial (z_1,z_2)}{\partial (x_1,x_2)}\right)= \det\begin{pmatrix} \frac{\partial z_1}{\partial x_1}&\frac{\partial z_1}{\partial x_2}\\ \frac{\partial z_2}{\partial x_1}&\frac{\partial z_2}{\partial x_2} \end{pmatrix}= \det\begin{pmatrix} I_{d_1}&0\\ \frac{\partial{z_2}}{\partial x_1}&\text{diag}(\exp(s_\theta(x_1))) \end{pmatrix}=\prod_{j=1}^{d_1}\exp(s_\theta(x_1)_j)=\exp\left[\sum_{j=1}^{d_1}s_\theta(x_1)\right] \] 对比 NICE，Real NVP 的 Jacobian 行列式不再恒为 1，这也是其名称“non-volume preserving”的含义。

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

棋盘式 masking 以坐标交替的方式做划分，而通道式 masking 将前一半通道和后一般通道划分成两组。另外，上图还同时展示了压缩操作，将 \(s\times s\times c\) 的 tensor 下采样到 \(\frac{s}{2}\times\frac{s}{2}\times 4c\) 的大小。设计网络时，作者堆叠使用多层交替棋盘式 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 卷积层。对于输入的图像数据 \(x\)，actnorm 层逐通道地做如下仿射变换： \[ z_{i,j}=s\odot x_{i,j}+b,\quad i=1,\ldots,h,\,j=1,\ldots,w \] 而可逆 1x1 卷积层可以看作是每个位置上沿通道的全连接层，即： \[ z_{i,j}=Wx_{i,j},\quad i=1,\ldots,h,\,j=1,\ldots,w \] 值得一提的是，由于 \(\det(W)\) 计算复杂度较高，作者将 \(W\) 进行 LU 分解： \[ W=PL(U+\text{diag}(s)) \] 其中 \(P\) 是一个置换矩阵，\(L\) 是对角线为 1 的下三角矩阵，\(U\) 是对角线为 0 的上三角矩阵，这样 \(\det(W)\) 就是对角线上所有 \(s\) 的乘积。作者初始化时随机采样一个旋转矩阵 \(W\)，计算出 \(P\) 后固定不动，优化 \(L,U,s\). 此时 \(P\) 就起到了交叉信息流的作用，因此我们不必再手动交换分支了。

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). ↩︎