基于流的生成模型

基于流的生成模型（flow-based generative model）是机器学习中的一类生成模型，利用归一化流（normalizing flow）显式建模概率分布。这是一种使用概率密度变量变换法将简单分布转换为复杂分布的统计方法。

直接建模似然函数具有很多优点。例如，可以直接计算得到负对数似然并将其作为损失函数最小化。此外，通过从初始分布中采样并应用流变换可以生成新的样本。

相比之下，变分自编码器、生成对抗网络等其他生成模型无法显式地表示似然函数。

## 方法

<image>

考虑随机变量$z_1$和$z_0$，其中$z_0 = f^{-1}_1(z_1)$。对于$i = 1, ..., K$，定义一系列由$z_0$变换得到的随机变量$z_i = f_i(z_{i-1})$。其中函数$f_1, ..., f_K$满足可逆性，即存在反函数$f^{-1}_i$。最终输出变量$z_K$用于对目标分布进行建模。

$z_K$的对数似然为（参见下方推导过程）：

$$\log p_K(z_K) = \log p_0(z_0) - \sum_{i=1}^{K} \log \left|\det \frac{df_i(z_{i-1})}{dz_{i-1}}\right|$$

为了高效计算对数似然，函数$f_1, ..., f_K$应既易于求逆，也易于计算其雅可比矩阵的行列式。在实践中，这些函数通常使用深度神经网络建模，并通过训练以最小化目标分布数据样本的负对数似然。这些架构一般设计为只需神经网络的正向传播即可完成逆运算与雅可比行列式的计算，例如NICE、RealNVP、Glow等。

### 对数似然的推导

考虑$z_1$与$z_0$，两者之间满足$z_0 = f^{-1}_1(z_1)$。通过概率密度变量变换公式，$z_1$的分布为：

$$p_1(z_1) = p_0(z_0)\left|\det \frac{df_1^{-1}(z_1)}{dz_1}\right|$$

其中$\det \frac{df_1^{-1}(z_1)}{dz_1}$是$f^{-1}_1$的雅可比矩阵的行列式。

由反函数定理可以得到

$$p_1(z_1) = p_0(z_0)\left|\det \left(\frac{df_1(z_0)}{dz_0}\right)^{-1}\right|$$

使用行列式性质$\det(A^{-1}) = \det(A)^{-1}$（其中$A$是可逆矩阵），则有：

$$p_1(z_1) = p_0(z_0)\left|\det \frac{df_1(z_0)}{dz_0}\right|^{-1}$$

对上式取对数后得到对数似然：

$$\log p_1(z_1) = \log p_0(z_0) - \log \left|\det \frac{df_1(z_0)}{dz_0}\right|$$

对于任意$z_i$和$z_{i-1}$都能得到类似结论。最终由递归关系可以得到：

$$\log p_K(z_K) = \log p_0(z_0) - \sum_{i=1}^{K} \log \left|\det \frac{df_i(z_{i-1})}{dz_{i-1}}\right|$$

## 训练方法

与训练其他一些深度学习模型类似，归一化流的目标是最小化模型的似然分布与目标分布之间的KL散度。将模型的似然分布记为$p_\theta$，要学习的目标分布记为$p^*$，则（正向）KL散度为：

$$D_{KL}[p^{*}(x)||p_{\theta}(x)] = -\mathbb{E}_{p^{*}(x)}[\log(p_{\theta}(x))] + \mathbb{E}_{p^{*}(x)}[\log(p^{*}(x))]$$

上式右边第二项表示目标分布的熵，与模型参数$\theta$无关，因此在优化时可以忽略，留下需要最小化的项是目标分布下的负对数似然的期望。由于此项难以直接计算，可以通过重要性采样的蒙特卡洛方法来近似。若已从目标分布$p^*(x)$中独立采样得到的数据集$\{x_{i}\}_{i=1:N}$从的样本，则该项可近似估计为：

$$-\hat{\mathbb{E}}_{p^{*}(x)}[\log(p_{\theta}(x))] = -\frac{1}{N} \sum_{i=0}^{N} \log(p_{\theta}(x_{i}))$$

因此，可以将学习目标

$$\underset{\theta}{\operatorname{arg\,min}}\ D_{KL}[p^{*}(x)||p_{\theta}(x)]$$

替换为

$$\underset{\theta}{\operatorname{arg\,max}}\ \sum_{i=0}^{N} \log(p_{\theta}(x_{i}))$$

换句话说，最小化KL相度相当于最大化模型在观测样本下的似然。

训练归一化流的伪代码如下：

- 输入：数据集$x_{1:n}$，归一化流模型$f_\theta(\cdot), p_0$
- 求解：通过梯度下降法最化$\max_\theta \sum_j \ln p_\theta(x_j)$
- 输出：优化后的参数$\hat\theta$

## 变体

### 平面流

平面流（planar flow）是最早的归一化流方法。给定某个激活函数$h$，以及具有适当维度的参数$\theta = (u, w, b)$，可以定义$$x = f_\theta(z) = z + u h(\langle w, z \rangle + b)$$一般而言，逆映射$f_\theta^{-1}$没有解析解。

相应的雅可比行列式为$|\det (I + h'(\langle w, z \rangle + b) uw^T)| = |1 + h'(\langle w, z \rangle + b) \langle u, w\rangle|$。

为保证其处处可逆，行列式必须在整个定义域内非零。例如当$h = \tanh$、$\langle u, w \rangle > -1$时可以满足可逆性要求。

### 非线性独立成分估计（NICE）

非线性独立成分估计（nonlinear independent components estimation，简称NICE）假设$x, z\in \R^{2n}$是偶数维变量，并将其从中间分成两部分。此时归一化流的定义为

$$x = \begin{bmatrix} <pre> x_1 \\ x_2 \end{bmatrix}= f_\theta(z) = \begin{bmatrix} z_1 \\z_2 \end{bmatrix} + \begin{bmatrix} 0 \\ m_\theta(z_1) \end{bmatrix}$$

其中$m_\theta$是任何带有权重$\theta$的神经网络。

其逆映射$f_\theta^{-1}$为$z_1 = x_1, z_2 = x_2 - m_\theta(x_1)$，雅可比行列式为1，即该归一化流动是体积保持（volume-preserving）的。

当$n=1$时，这一映身可以视为沿$x_2$方向的一种曲线剪切。

### 实值非体积保持（Real NVP）

实值非体积保持（real non-volume preserving，简称Real NVP）是NICE模型的一种推广，定义为：

$$x = \begin{bmatrix} <pre> x_1 \\ x_2 \end{bmatrix}= f_\theta(z) = \begin{bmatrix} z_1 \\ e^{s_\theta(z_1)} \odot z_2 \end{bmatrix} + \begin{bmatrix} 0 \\ m_\theta(z_1) \end{bmatrix}$$

其逆映射是$z_1 = x_1, z_2 = e^{-s_\theta (x_1)}\odot (x_2 - m_\theta (x_1))$，相应的雅可比行列式为$\prod^n_{i=1} e^{s_\theta(z_{1, })}$。当$s_\theta = 0$时，退化为NICE模型。由于Real NVP映射将向量$x$的两部分分开处理，通常需要在每一层后添加一个置换操作$(x_1, x_2) \mapsto (x_2, x_1)$。

### 生成流（Glow）

生成流（generative flow，简称Glow）模型中每层由三个部分组成：

- 通道方向的仿射变换$$y_{cij} = s_c(x_{cij} + b_c)$$相应的雅可比行列式为$\prod_c s_c^{HW}$。
- 可逆1x1卷积$$z_{cij} = \sum_{c'} K_{cc'} y_{cij}$$相应的雅可比行列式为$\det(K)^{HW}$，其中$K$是任意可逆矩阵。
- Real NVP部分，其雅可比行列式如前所述。

生成流通过引入可逆1x1卷积，改进了Real NVP中仅仅置换前后两部分的方式，而是对所有层的通道进行总体上的置换。

### 掩码自回归流（MAF）

掩码自回归流（masked autoregresssive flow，简称MAF）基于自回归模型，其定义了一个分布在$\R^n$上的随机过程：

$$\begin{align} <pre> x_1 \sim& N(\mu_1, \sigma_1^2)\\ x_2 \sim& N(\mu_2(x_1), \sigma_2(x_1)^2)\\ &\cdots \\ x_n \sim& N(\mu_n(x_{1:n-1}), \sigma_n(x_{1:n-1})^2)\\ </pre> \end{align}$$

其中$\mu_i: \R^{i-1} \to \R$和$\sigma_i: \R^{i-1} \to (0, \infty)$是定义自回归模型的固定函数。

使用重参数化技巧，该自回归模型可以被推广为归一化流：

$$\begin{align} <pre> x_1 =& \mu_1 + \sigma_1 z_1\\ x_2 =& \mu_2(x_1) + \sigma_2(x_1) z_2\\ &\cdots \\ x_n =& \mu_n(x_{1:n-1}) + \sigma_n(x_{1:n-1}) z_n\\ </pre> \end{align}$$

令$z \sim N(0, I_{n})$ 可以重新得到自回归模型。

正向映射由于是顺序性的因而会很慢，但反向映射因为可以并列而会比较很快。

相应的雅可比矩阵是下对角矩阵，其行列式为$\sigma_1 \sigma_2(x_1)\cdots \sigma_n(x_{1:n-1})$。

通过反转$f_\theta$和$f_\theta^{-1}$这两个映射，可以得到逆自回归流 (inverse autoregressive flow，简称IAF)。与MAF相反，IAF的正向映射较快、反向映射较慢。

### 连续归一化流（CNF）

除了通过函数组合来构建流的方法，另一种方法是将流表示为连续时间动力学，得到连续归一化流（continuous normalizing flow，简称CNF）。设$z_0$为具有分布$p(z_0)$的潜变量，使用以下流函数将此潜变量映射到数据空间：

$$x = F(z_0) = z_T = z_0 + \int_0^T f(z_t, t) dt$$

其中$f$是任意函数，可以使用神经网络等进行建模。其反函数为：

$$z_0 = F^{-1}(x) = z_T + \int_T^0 f(z_t, t) dt = z_T - \int_0^T f(z_t,t) dt$$

于是可以得到$x$的对数似然：

$$\log(p(x)) = \log(p(z_0)) - \int_0^T \text{Tr}\left[\frac{\partial f}{\partial z_t} \right] dt$$

由于上式中的迹仅取决于雅可比矩阵$\partial_{z_t} f$的对角线，这意味着对雅可比矩阵的形式没有任何限制。这与先前的离散归一化模型不同，后者将雅可比矩阵设计为上对角或下对角形式，以便更高效地计算其行列式。

其中的迹可以使用“哈钦森技巧”（Hutchinson's trick）来估计：给定任意矩阵$W\in \R^{n\times n}$以及满足$E[uu^T] = I$的任意随机向量$u\in \R^n$，可以得到$E[u^T W u] = tr(W)$。

通常，随机向量$u$可以从正态分布$N(0, I)$ 或拉德马赫分布$\{\pm n^{-1/2}\}^n$中进行采样。

当$f$由神经网络实现时，需要使用神经常微分方程。实际上，CNF最早是与神经常微分方程在同一篇论文中提出的。

CNF主要存在两个缺陷：一是连续流必须是同胚的，从而保持方向性和环境同痕（例如，不可能通过空间连续变形将左手翻转为右手，也不可能将球面外翻或解开一个结），二是学习到的流$f$可能会由于退化而表现不佳（有无数个可能的$f$都能解决同一问题)。

通过增加额外的维度，CNF可以获得足够的自由度来反转方向并突破环境同痕（类似于可以在三维空间中翻转一个多边形，或在四维空间中解开一个结），从而得到“增强神经常微分方程”。

通过结合流形的惠特尼嵌入定理和神经网络的通用近似定理，能够证明$\R^n$的任何同胚可以由$\R^{2n+1}$上的神经常微分方程近似。

还可以为流$f$引入正则化损失，例如基于最优传输理论的正则化损失：

$$\lambda_{K} \int_{0}^{T}\left\|f(z_t, t)\right\|^{2} dt +\lambda_{J} \int_{0}^{T}\left\|\nabla_z f(z_t, t)\right\|_F^{2} dt$$

其中$\lambda_K, \lambda_J >0$是超参数。第一项惩罚模型随时间变化流场的振荡，第二项则惩罚模型随空间变化流场的振荡。这两项共同引导模型生成在空间和时间上平滑的流。

## 缺点

尽管归一化流在估计高维概率密度函数方面取得了成功，但其设计仍然存在一些缺陷。首先，归一化流的潜空间并不是一个低维空间，因此基于流的模型默认情况下不支持数据压缩，需要很大的计算量。不过，仍有办法可以用它们进行图像压缩。

此外，基于流的模型在估计分布外样本（即非训练集分布中抽取的样本）的似然值时通常表现不佳。学者提出了一些假设来解释这一现象，其中包括典型集假设、模型训练中的估计问题或由于数据分布熵引起的基础性问题。

归一化流最有趣的特性之一是其学习到的双射映射的可逆性。这一特性通过模型设计中的约束得以保证，从而确保理论上的可逆性。逆映射对确保变量变换定理的适用性、雅可比行列式的计算以及模型采样都至关重要。然而在实践中，由于数值不精确性，这种可逆性可能被破坏，进而导致逆映射爆炸。

## 应用

基于流的生成模型已应用于多种场景，例如：

- 音频生成
- 图像生成
- 分子图生成
- 点云建模
- 视频生成
- 有损图像压缩
- 异常检测