Flow Matching For Generative Modeling

Flow Matching For Generative Modeling

一、基于流的（Flow based）生成模型

生成模型

我们先回顾一下所谓的生成任务，究竟是想要做什么事情。我们认为，世界上所有的图片，是符合某种分布 $p_{data}(x)$ 的。当然，这个分布肯定是个极其复杂的分布。而我们有一堆图片 $x_1,x_2,\dots ,x_m$ ，则可以认为是从这个分布中采样出来的 $m$ 个样本。我们通过训练希望得到一个生成器网络 $G$ ，该网络能够做到输入一个从正态分布 $\pi (z)$ 中采样出来的 $z$ ，输出一张看起来像真实世界的图片 $x=G(z)\sim p_G(x)$ 。我们希望采样并生成出的数据分布 $p_G(x)$ 与真实的数据分布 $p_{data}(x)$ 越接近越好。

从概率模型的角度来看，想要做到上面说的这件事情，可以通过最大化对数似然 $\log p_G$，来优化生成器 $G$ 的参数：
$$G^{\ast }=\arg \underset{G}{\max }\sum _{i=1}^m\log p_G(x_i)$$
可以证明，最大化这个对数似然，就相当于最小化生成器分布 $p_G(x)$ 与目标分布 $p_{data}(x)$ 的 KL散度，即让这两个分布尽量接近：
$$G^{\ast }\approx \arg \underset{G}{\min }KL(p_{data}||p_G)$$

概率密度的变量变换定理

给定一个随机变量 $z$ 及其概率密度函数 $z\sim \pi (z)$ ，通过一个一对一的映射函数 $f$ 构造一个新的随机变量 $x=f(z)$。如果存在逆函数 $f^{-1}$ 满足 $z=f^{-1}(x)$，那么新变量 $x$ 的概率密度函数 $p(x)$ 计算如下：
$$\begin{gathered} p(x)=\pi (z)|\frac{dz}{dx}|=\pi (f^{-1}(x))|\frac{d{f}^{-1}}{dx}|=\pi (f^{-1}(x))|(f^{\prime -1}(x)|若z为随机变量 \\ p(\mathbf{x})=\pi (\mathbf{z})|\det (\frac{d\mathbf{z}}{d\mathbf{x}})|=\pi (f^{-1}(\mathbf{x}))|\det ({\mathbf{J}}_{f^{-1}})|若\mathbf{z}为随机向量 \end{gathered}$$
其中 $\det (\cdot )$ 表示行列式，$\mathbf{J}$ 表示雅可比矩阵，是向量函数中因变量各维度关于自变量各维度的偏导数组成的矩阵，可类比为单变量函数的导数。

流模型推导

现在流行的生成模型五花八门，各显神通。VAE 优化变分下界 ELBO、GAN 通过对抗训练来隐式地逼近数据分布。流模型则可以直接优化对数似然。

流模型通过最大化对数似然，来优化生成器 $G$：
$$G^{\ast }=\arg \underset{G}{\max }\sum _{i=1}^m\log p_G(x_i)$$
而根据变量变换定理，有：
$$p_G(x_i)=\pi (z_i)|\det (J_{G^{-1}})|,z_i=G^{-1}(x_i)$$
则对数似然：
$$\log p_G(x_i)=\log \pi (G^{-1}(x_i))+\log |\det (J_{G^{-1}})|$$
要训练一个好的生成器 $G$ 我们只需要训练一个（或一系列）网络完成从噪声分布 $\pi (z)$ 到数据分布 $p_{\text{data}}(x)$ 的变换就可以了。在采样生成时，求出生成器的逆向网络 $G^{-1}$，再将随机采样的噪声输入，即可生成新的符合数据分布的样本。只要最大化上面这个式子，就可以了。

现在的问题就是怎么把这个式子算出来，具体来说，这个式子计算的关键在以下两点：

• 如何计算行列式 $\det (J_G)$
• 如何求逆矩阵 $G^{-1}$

我们设计的生成器网络 $G$ 需要满足上面这两个条件，这就是流模型生成器数学上的限制。在之前的流模型研究中，研究者们提出了许多设计精巧的网络（如 decoupling layer），可以巧妙地使得网络满足上述两点便利计算的要求。

另外要提一点，流模型的输入输出的尺寸必须是一致的。这是因为如果想要 $G$ 可逆，它的输入输出维度一致是一个必要条件（非方阵不可能可逆）。比如要生成 $100\times 100\times 3$ 的图像，那输入的随机噪声也是 $100\times 100\times 3$​​ 的。这与 VAE、GAN 等生成模型很不一样，这些生成模型的输入维度通常远小于输出维度。

堆叠多个网络

在实际中，由于可逆神经网络存在数学上的诸多限制，其单个网络的表达能力有限，我们一般需要堆叠多层网络来得到一个生成器，这也是 “流模型” 这个名称的由来。不过虽然堆叠了很多层，在公式上也没有什么复杂的。无非就是把一堆 $G_i$ 连乘起来，通过 $\log$​ 之后，又变成连加。

比如我们有 $K$ 个网络 { fi}i=1K\{f_i\}_{ {i=1}}^K{ fi​}i=1K​，对噪声分布 $\pi ({\mathbf{z}}_0)$ 进行 $K$ 步变换，得到数据 $\mathbf{x}$，即有：
$$\mathbf{x}={\mathbf{z}}_k=f_K(f_{K-1}...f_1({\mathbf{z}}_0))$$
对于其中第 $i$ 步有：
$$\begin{gathered} {\mathbf{z}}_i\sim p_i({\mathbf{z}}_i) \\ {\mathbf{z}}_i=f_i({\mathbf{z}}_{i-1}),{\mathbf{z}}_{i-1}=f^{-1}({\mathbf{z}}_i) \end{gathered}$$
根据变量变换定理，相邻两步之间的隐变量分布的关系为：
$$\begin{array}{rl}{p}_i({\mathbf{z}}_i)& =p_{i-1}(f_i^{-1}({\mathbf{z}}_i))|\det {\mathbf{J}}_{f_i^{-1}}|\\ & =p_{i-1}({\mathbf{z}}_{i-1})|\det {\mathbf{J}}_{f_i}{|}^{-1}\end{array}$$
每一步的对数似然为：
log⁡pi(zi)=log⁡pi−1(zi−1)−log⁡(det⁡(Jfi)) \log p_i(\mathbf{z}_i)=\log p_{i-1}(\mathbf{z}_{i-1})-\log(\det(\mathbf{J}_{f_i})) logp