Flow-based Generative Models.

本文目录：

  1. 流模型的结构
  2. 流模型的优化目标
  3. 流模型的各种变体
  4. 使用nflows构造流模型

1. 流模型的结构

流模型(flow-based model)是生成模型(generative model)的一个分支，旨在通过一系列可逆变换(双射函数$f$)建立较为简单的先验分布$p_Z(z)$与较为复杂的实际数据分布$p_X(x)$之间的映射关系。

若定义变换$x=f(z)$，根据概率密度的变量替换公式，实际概率分布$p_X(x)$可以由先验分布$p_Z(z)$表示：

\[p_X(x) = p_Z(z)\cdot |\det J_{f^{-1}}(x)| = p_Z(z)\cdot |\det J_{f}(z)|^{-1}\]

其中$z=f^{-1}(x)$是$x=f(z)$的反函数，$J_{f}(z)$是函数$f(z)$关于变量$z$的Jacobian矩阵，$\det J_{f}(z)$是Jacobian行列式。此时不需要显式地计算分布$p_X(x)$的概率密度函数，而是通过先验分布$p_Z(z)$的概率密度以及映射过程产生的Jacobian行列式计算即可。

对双射函数$x=f(z)$有两点要求：

  1. $x=f(z)$可逆，并且易于求逆函数。
  2. 对应的Jacobian行列式$\det J_{f}(z)$容易计算。

在实现时通过复合多个双射函数$z_i=f_{i}(z_{i-1})$以增强非线性拟合能力。模型整体像水管中的流水一样汇聚，因此称为流(flow)模型。

\[z_0 \leftrightarrow z_1 \leftrightarrow z_2 \leftrightarrow \cdots \leftrightarrow z_{K-1} \leftrightarrow z_{K} =x\]

其中Jacobian矩阵和Jacobian行列式可以根据链式法则计算：

\[\begin{aligned} J_{f}(z) &= [\frac{\partial x}{\partial z_0}] = [\frac{\partial z_{K}}{\partial z_0}]=[\frac{\partial z_{1}}{\partial z_{0}}][\frac{\partial z_{2}}{\partial z_{1}}]\cdots [\frac{\partial z_K}{\partial z_{K-1}}] = \prod_{k=1}^{K} J_{f_k}(z_{k-1}) \\ \det J_{f}(z) &= \det \prod_{k=1}^{K} J_{f_k}(z_{k-1})= \prod_{k=1}^{K} \det J_{f_k}(z_{k-1}) \end{aligned}\]

⚪ 讨论：流模型与自编码器

流(flow-based)模型学习真实数据分布$\hat{p}_X$与隐变量$p_Z$之间可逆、稳定的映射关系。在推断时从数据分布中采样样本$x$，通过函数$f^{-1}(x)$将其映射为隐空间中的近似样本$z$，隐空间的分布可以人为指定为已知的简单分布形式。在生成时从隐空间分布中采样样本$z$，通过反函数$f(z)$将其映射为数据分布中的近似样本$x$。

从结构上看，流模型也是一类特殊的自编码器模型。编码器$z=f^{-1}(x)$将输入$x$编码为隐变量$z$，并且使得$z$服从较为简单的分布形式(如标准正态分布)。由于编码器是可逆的，一旦编码器训练完成，便可以立即得到相对应的解码器$x=f(z)$。

⚪ 讨论：流模型与隐变量模型

事实上，流模型也是一类特殊的隐变量模型。隐变量模型是指由于原概率分布$p_{\theta}(x)$的形式通常是未知的，引入隐变量(latent variable) $z$ 间接地构造概率分布$p_{\theta}(x)$：

\[p_{\theta}(x) = \int p_{\theta}(x,z) dz = \int p_Z(z)p_{\theta}(x|z) dz\]

上式包含积分运算，直接求解比较困难。流模型通过巧妙的假设使得上式的计算变得可行。具体地，流模型把生成分布$p_{\theta}(x|z)$设定为Dirac函数 $\delta(x-f(z))$，其中$f(z)$应为可逆函数，记其反函数为$f^{-1}(x)$。Dirac函数$\delta(x)$具有如下性质：

\[f(x) = \int f(y)\delta(x-y) dy\]

则概率分布$p_{\theta}(x)$等价于：

\[p_{\theta}(x) = \int p_Z(z)p_{\theta}(x|z) dz = \int p_Z(z)\delta(x-f(z)) dz =p_Z(f^{-1}(x))\]

因此计算概率分布$p(x)$的表达式相当于对隐变量分布$p(z)$的表达式进行积分变换$z=f^{-1}(x)$。

2. 流模型的优化目标

流模型显式地给出了概率分布$p_X(x)$的表达式，因此求解过程可采用极大似然估计：

\[\begin{aligned} \log p_X(x) &= \log p_Z(z)\cdot |\prod_{k=1}^{K} \det J_{f_k}(z_{k-1})|^{-1} \\ &= \log p_Z(z) - \sum_{k=1}^{K}\log | \det J_{f_k}(z_{k-1})| \end{aligned}\]

其中先验分布$p_Z(z)$可以指定为形式较为简单的已知分布，比如各分量独立的标准正态分布：

\[p_Z(z) = \frac{1}{(\sqrt{2\pi})^D}e^{-\frac{1}{2}||z||^2}\]

此时优化目标可简化为：

\[\log p_X(x) \text{ ~ } -\frac{1}{2}||z||^2 - \sum_{k=1}^{K}\log |\det J_{f_k}(z_{k-1})|\]

前者表示输出的平方和，后者表示所有变换贡献的Jacobian行列式的对数。为了构造可解的目标函数，要求设计合适的双射函数$x=f(z)$使得Jacobian行列式容易计算。

3. 流模型的各种变体

从优化目标中可以看出，流模型是由先验分布$p_Z(z)$和双射函数$x=f(z)$唯一确定的。根据双射函数的不同设计思路，笔者将流模型分为以下几类：

方法 双射函数$x=f(z)$ 引入Jacobian行列式的特殊结构
Normalizing Flow \(\begin{aligned} \text{Planar Flow : } & x = z+uh(w^Tz+b) \\ \text{Radial Flow : } & x = z+\frac{b}{a+||z-z_0||^2}(z-z_0)\end{aligned}\) -
i-ResNet \(z = x + g(x), \mathop{\max}_{x_1 \ne x_2} \frac{||g(x_1)-g(x_2)||_2}{||x_1-x_2||_2} \lt 1\) -
IAF \(\begin{aligned} x_i &= z_i \cdot \exp(\alpha_i) + \mu_i \\ \mu_i &= g_{\mu_i}(z_{1:i-1}),\alpha_i = g_{\alpha_i}(z_{1:i-1}) \end{aligned}\) -
MAF \(\begin{aligned} x_i &= z_i \cdot \exp(\alpha_i) + \mu_i \\ \mu_i &= g_{\mu_i}(x_{1:i-1}),\alpha_i = g_{\alpha_i}(x_{1:i-1}) \end{aligned}\) -
NICE \(\begin{aligned} \text{加性}&\text{耦合层：}\\ z_{1:d}&= x_{1:d} \\ z_{d+1:D}&=x_{d+1:D}+m(x_{1:d}) \end{aligned}\) 尺度变换层：$z=s \otimes f^{-1}(x)$
Real NVP \(\begin{aligned} \text{仿射}&\text{耦合层：}\\ z_{1:d}&= x_{1:d} \\ z_{d+1:D}&=x_{d+1:D}\otimes \exp(s(x_{1:d}))+t(x_{1:d}) \end{aligned}\) 标准化层：batchnorm
多尺度结构：\(p(z_1,z_3,z_5) = p(z_1|z_2)p(z_3|z_4)p(z_5)\)
Glow \(\begin{aligned} \text{仿射}&\text{耦合层：}\\ z_{1:d}&= x_{1:d} \\ z_{d+1:D}&=x_{d+1:D}\otimes \exp(s(x_{1:d}))+t(x_{1:d}) \end{aligned}\) 标准化层：actnorm
通道置换操作：1x1可逆卷积
多尺度结构：\(p(z_1,z_3,z_5) = p(z_1|z_2)p(z_3|z_4)p(z_5)\)
Flow++ \(\begin{aligned} \text{逻辑混合}&\text{CDF耦合层：}\\ z_{1:d}= &x_{1:d} \\ z_{d+1:D}=&\sigma^{-1}(\text{MixLogCDF}(x_{d+1:D}; \\ &\pi(x_{1:d}),\mu(x_{1:d}),s(x_{1:d}))) \\ & \cdot \exp(a(x_{1:d})) + b(x_{1:d}) \end{aligned}\) 变分解量化：\(\Bbb{E}_{x \text{~} p_{data},\epsilon \text{~} \mathcal{N}(0,I)} [\log \frac{p_{model}(x+q_x(\epsilon))}{p(\epsilon)\cdot |\partial q_x/\partial \epsilon|^{-1}}]\)

4. 使用nflows构造流模型

nflows是一个基于PyTorch的第三方库，提供流模型的快捷搭建与使用。首先安装nflows库：

pip install nflows

在使用nflows库构造流模型时，首先需要定义先验分布$p_Z(z)$ (即distribution)和双射函数$x=f(z)$ (即transform)，然后构造流模型：

from nflows import transforms, distributions, flows

# Define an invertible transformation.
transforms = []
for _ in range(num_layers):
    transforms.append(transforms.ReversePermutation(features=2))
    transforms.append(transforms.MaskedAffineAutoregressiveTransform(features=2, 
                                                                     hidden_features=4))
transform = CompositeTransform(transforms)

# Define a base distribution.
base_distribution = distributions.StandardNormal(shape=[2])

# Combine into a flow.
flow = flows.Flow(transform=transform, distribution=base_distribution)

目标函数为负对数似然：

optimizer = optim.Adam(flow.parameters())

for i in range(num_iter):
    # Calculate loss and update gradient.
    optimizer.zero_grad()
    loss = -flow.log_prob(inputs).mean()
    loss.backward()
    optimizer.step()

从流模型中采样：

samples = flow.sample(num_samples)

先验分布$p_Z(z)$和双射函数$x=f(z)$的选择可参考distribution, transform。

⚪ 参考文献