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++<span class="hint--top hint--rounded" aria-label="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.">[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
- Dinh, Laurent, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516 (2014). ↩︎
- Dinh, Laurent, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605.08803 (2016). ↩︎
- Kingma, Durk P., and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems 31 (2018). ↩︎
- 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. ↩︎
- Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems 31 (2018). ↩︎
- 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). ↩︎
- Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 (2022). ↩︎