0. TL; DR

本文提出了一种基于任意一对基础概率密度和目标概率密度之间的连续时间归一化流的生成模型。该流的速度场是从一个在有限时间内在这两种密度之间插值的时间依赖性密度的概率流中推断出来的。插值方法导出了一个针对速度场本身的简单二次损失函数，该损失函数以期望的形式表示，易于通过经验估计。这个流可以用于从基础分布或目标分布中生成样本，并估计插值路径上任何时间的似然。

此外，该流还可以通过优化来最小化插值密度路径的长度，从而为构建最优传输映射铺平了道路。在基础分布是高斯密度的情况下，该归一化流的速度也可以用于构建一个扩散模型来对目标进行采样并估计其分数。在密度估计任务上的基准测试表明，学习到的流能够以更低的成本匹配甚至超越传统的连续流，并且在CIFAR-10和ImageNet 32×32上的图像生成任务中与扩散模型相媲美。

1. 背景介绍

现代生成模型的核心任务是学习一个映射，将一个简单的分布（如高斯噪声）转换为一个复杂的数据分布。无论是以隐式映射为基础的GAN，还是以可逆变换为基础的流模型，其挑战都在于如何设计这个传输过程，使其既能有效地拟合复杂的目标分布，又能在计算上保持高效。

在连续时间的视角下，这个问题可以被描述为设计一个时变的映射$X_t(x)$，它将初始的基础分布$\rho_0$在时间$t=0$时，逐渐推向在$t=1$时的目标分布$\rho_1$。这个连续的映射过程可以由一个常微分方程（ODE）来描述：

\[\dot{X}_t(x) = v_t(X_t(x)), \quad X_{t=0}(x) = x\]

其中，$v_t(x)$是驱动这个传输的速度场。相应的，概率密度$\rho_t(x)$的演化则遵循一个连续性方程：

\[\partial_t \rho_t + \nabla \cdot (v_t \rho_t) = 0\]

因此，生成建模的核心问题就变成了如何从数据中推断出这个速度场$v_t(x)$。

传统的方法，如最大似然估计（Maximum Likelihood Estimation, MLE），需要通过ODE求解器进行反向传播来计算似然的梯度，这个过程计算成本极高，限制了模型的可扩展性。而分数匹配（score-based）方法虽然通过学习得分函数$\nabla \log \rho_t(x)$来间接定义速度场，但其理论基础依赖于一个无限时间的扩散过程，并且引入了SDE的复杂性。

为了解决这些问题，作者提出了一种全新的方法，通过引入一个“随机插值体”（stochastic interpolant），直接、高效地学习速度场$v_t(x)$。

2. InterFlow 方法

InterFlow的核心是构建一个随机过程$x_t$，其在任意时刻$t$的分布$\rho_t(x)$能够平滑地连接基础分布$\rho_0$和目标分布$\rho_1$。

2.1 随机插值体 (Stochastic Interpolant)

作者定义了一个“插值体”$I_t(x_0, x_1)$，它是一个可微函数，满足边界条件$I_{t=0}(x_0, x_1) = x_0$和$I_{t=1}(x_0, x_1) = x_1$。然后，通过从$\rho_0$中独立采样$x_0$和从$\rho_1$中独立采样$x_1$来构造一个随机过程$x_t$：

\[x_t = I_t(x_0, x_1), \quad x_0 \sim \rho_0, x_1 \sim \rho_1 \text{ independent.}\]

作者将这个$x_t$称为随机插值体。一个特别有用的插值体例子是三角插值：

\[I_t(x_0, x_1) = \cos(\frac{1}{2}\pi t)x_0 + \sin(\frac{1}{2}\pi t)x_1\]

这个插值体在单位圆上平滑地连接$x_0$和$x_1$。

2.2 关联的概率流与速度场

InterFlow证明了：由随机插值体$x_t$诱导的概率密度$\rho_t(x)$满足连续性方程，并且其对应的速度场$v_t(x)$是唯一一个能够最小化一个简单二次目标函数的解。这个目标函数$G(\hat{v})$为：

\[G(\hat{v}) = \mathbb{E} \left[ |\hat{v}_t(I_t(x_0, x_1))|^2 - 2 \partial_t I_t(x_0, x_1) \cdot \hat{v}_t(I_t(x_0, x_1)) \right]\]

其中，期望$\mathbb{E}$是针对$t \sim U[0,1]$, $x_0 \sim \rho_0$, $x_1 \sim \rho_1$进行计算的。

这个结论的意义是巨大的：

可计算性：这个目标函数$G(\hat{v})$完全由期望构成，因此可以通过从$t, x_0, x_1$中采样来通过蒙特卡洛方法进行高效的经验估计。
无模拟训练：与MLE不同，它的计算完全不依赖于在训练过程中求解ODE，从而避免了昂贵的反向传播。
直接学习速度场：它直接为速度场$v_t(x)$提供了一个回归目标。通过让神经网络$\hat{v}_\theta(x,t)$去最小化这个二次损失，就可以学到正确的速度场。

2.3 优化传输路径

目标函数$G(\hat{v})$在最优速度场$v$处的最小值为：

\[G(v) = - \int_0^1 \int_{\mathbb{R}^d} |v_t(x)|^2 \rho_t(x) dx dt\]

这个值恰好是Wasserstein-2路径长度的负值。这意味着，我们可以通过最大化$G(v)$来找到更短的传输路径，从而实现最优传输。这可以通过优化插值体$I_t$或基础分布$\rho_0$的参数来实现。

2.4 与分数匹配和扩散模型的联系

当基础分布$\rho_0$是标准高斯分布，并且使用三角插值体时，得分函数$\nabla \log \rho_t(x)$可以被速度场$v_t(x)$显式地表示出来：

\[\nabla \log \rho_t(x) = -x - \frac{2}{\pi} \tan(\frac{1}{2}\pi t) v_t(x)\]

这个关系表明，学习速度场$v_t(x)$和学习得分函数$\nabla \log \rho_t(x)$在某种程度上是等价的。然而，InterFlow直接在ODE的概率流层面上工作，完全绕过了SDE和扩散过程的复杂性，提供了一条更简单、更直接的建模路径。

3. 实验分析

作者在一系列数据集上对InterFlow进行了全面的实验验证。

3.1 二维密度估计与数据集插值

作者在多个复杂的二维玩具数据集上进行了测试，如八个高斯混合、棋盘格等。InterFlow能够完美地捕捉到目标分布的所有模式，生成的样本与真实分布几乎无法区分。InterFlow在两个任意的、没有解析形式的经验分布之间学习流映射。它成功地学习了一个将“漩涡”分布转换为“棋盘格”分布的流，并准确地展示了中间时刻的插值分布。