图生成的变分流匹配.
0. TL; DR
本文提出了流匹配的一个变分推断（variational inference）公式，称之为变分流匹配（variational flow matching, VFM）。基于这个公式开发了CatFlow，一种用于分类数据的流匹配方法。CatFlow易于实现，计算高效，并在图生成任务上取得了强大的结果。
VFM的关键观察是，可以将流的向量场参数化为后验概率路径（posterior probability path）的变分近似，该后验概率路径是轨迹可能终点的分布。作者表明，这种变分解释同时包含了CatFlow目标和原始流匹配目标作为其特例。作者还探讨了VFM与基于分数的模型（其动态是随机的而非确定性的）之间的关系，并基于一个重加权的VFM目标推导出了模型似然的一个界。
作者在三个图生成任务上对CatFlow进行了评估，包括一个抽象图生成任务和两个分子生成任务。在所有情况下，CatFlow的性能都超过或达到了当前最先进的水平。
1. 背景介绍
近年来，生成式建模领域取得了显著进展，尤其是在图像生成方面，基于扩散（diffusion）的方法已证明非常有效。然而，扩散模型固有的随机采样过程，使其在推理时需要大量步骤。
作为一种更灵活的替代方案，连续归一化流（Continuous Normalizing Flows, CNFs）能够学习比扩散模型更通用的概率路径，但其训练成本高昂，因为每一步训练都需要求解一个常微分方程（ODE）。
最近，流匹配（Flow Matching, FM）作为一种高效、无模拟（simulation-free）的CNF训练方法被提出。它通过一个基于每个样本的插值来推导出一个更简单的目标函数，用于学习生成所需概率路径的边际向量场。
本文在此基础上，提出了流匹配的一个新公式，称之为变分流匹配（Variational Flow Matching, VFM）。在VFM中，作者将向量场明确地参数化为一个关于轨迹终点分布的期望，这个终点分布由一个变分分布来近似。VFM的目标就是最小化这个变分分布与真实的后验概率路径之间的KL散度。
这个变分视角有几个重要的理论和实践意义：
- 统一性：作者证明，当变分近似是高斯分布且条件向量场是线性的时，VFM可以恢复出原始的流匹配目标。
- 可分解性：在常见的线性条件下，高维的变分推断问题可以分解为一系列独立的低维问题，这使得我们可以使用一个完全分解的变分近似而不会损失通用性。
- 对分类数据的扩展：这个框架自然地导出了一个用于分类数据的新方法，作者称之为CatFlow。在CatFlow中，VFM的目标函数简化为了一个标准的交叉熵损失，使其易于实现且计算高效。
作者将CatFlow应用于图生成任务，并在多个基准上取得了最先进的结果，证明了VFM框架的有效性和实用性。
2. Variational Flow Matching
2.1 从流匹配到变分推断
在流匹配中，驱动概率流的边际向量场$u_t(x)$可以表示为所有条件向量场$u_t(x|x_1)$在一个后验概率路径$p_t(x_1|x)$上的期望：
\[u_t(x) = \int u_t(x | x_1) p_t(x_1 | x) dx_1 = \mathbb{E}_{p_t(x_1|x)} [u_t(x | x_1)]\]其中，$p_t(x_1|x)$表示在时间$t$经过点$x$的所有轨迹，其终点$x_1$的概率分布。这个公式直观地告诉我们：一个点$x$的速度，是所有从它出发并最终到达不同终点$x_1$的“续行路径”的速度的加权平均。
基于这个观察，作者提出了VFM的核心思想：不再直接用神经网络去近似$u_t(x)$，而是用一个变分分布$q_\theta^t(x_1|x)$去近似那个难以处理的后验概率路径$p_t(x_1|x)$。然后，将近似的向量场$v_\theta^t(x)$定义为在这个变分分布下的期望：
\[v_\theta^t(x) := \int u_t(x | x_1) q_\theta^t(x_1 | x) dx_1\]显然，当变分分布$q_\theta^t$与真实后验$p_t$完全匹配时，$v_\theta^t$也就等于真实的向量场$u_t$。因此，学习向量场的问题就被转化为了一个变分推断问题。其目标是最小化$q_\theta^t(x_1|x)$和$p_t(x_1|x)$之间的KL散度。经过推导，这等价于最大化以下对数似然，即最小化VFM损失：
\[\mathcal{L}_{VFM}(\theta) = - \mathbb{E}_{t, x, x_1} [\log q_\theta^t(x_1 | x)]\]其中，$t, x, x_1$是从联合分布$p_t(x, x_1)$中采样的。
2.2 均场变分流匹配 (Mean-Field VFM)
直接对高维后验$p_t(x_1|x)$进行变分推断仍然很困难。幸运的是，在FM中常用的线性条件向量场（如直线插值$u_t(x|x_1) = \frac{x_1 - x}{1-t}$）的假设下，问题可以被极大地简化。
如果条件向量场$u_t(x|x_1)$对$x_1$是线性的，那么边际向量场$u_t(x)$只依赖于后验分布$p_t(x_1|x)$的均值，而与其更高阶的矩（如协方差）无关。
这个定理意味着，我们不需要去学习后验分布$p_t(x_1|x)$中复杂的维度间相关性。我们只需要保证我们的变分近似$q_\theta^t(x_1|x)$能够匹配真实后验的边际均值即可。因此，我们可以安全地使用一个完全分解的变分分布，而不会损失任何通用性：
\[q_\theta^t(x_1 | x) := \prod_{d=1}^D q_\theta^t(x_1^d | x)\]这被称为均场变分流匹配（Mean-Field VFM, MF-VFM），其损失函数也相应地分解为各个维度的对数似然之和：
\[\mathcal{L}_{MF-VFM}(\theta) = - \mathbb{E}_{t, x, x_1} \left[ \sum_{d=1}^D \log q_\theta^t(x_1^d | x) \right]\]2.3 CatFlow：用于分类数据的MF-VFM
将MF-VFM框架应用于图生成等涉及分类数据的场景。对于一个具有$K_d$个类别的离散变量，其状态可以表示为一个one-hot向量，或者等价地，一个概率单纯形（probability simplex）上的顶点。
为每个维度$d$选择一个分类分布作为其变分近似：$q_\theta^t(x_1^d|x) = \text{Cat}(x_1^d | \mu_t^d(x))$，其中$\mu_t^d(x)$是由神经网络预测的、属于各个类别的概率向量。将这个分类分布代入MF-VFM的损失函数，我们得到了一个非常熟悉的形式——交叉熵损失 (cross-entropy loss)：
\[\mathcal{L}_{\text{CatFlow}}(\theta) = - \mathbb{E}_{t, x, x_1} \left[ \sum_{d=1}^D \sum_{k=1}^{K_d} \mathbb{I}[x_1^d=k] \log \mu_{t}^{dk}(x) \right]\]尽管损失函数是交叉熵，但其对应的向量场计算仍然遵循VFM的定义，即条件向量场在变分分布下的期望。对于直线插值，$d$维度的向量场为：
\[v_{\theta, t}^d(x) = \frac{\mu_t^d(x) - x^d}{1-t}\]这在几何上对应于学习一个到概率单纯形内部某点（由$\mu_t^d(x)$定义）的映射，然后沿着直线流向该点。
与标准的流匹配相比，CatFlow的优势在于：
- 归纳偏置：它预测的是单纯形内的点，并参数化指向该点的速度，这引入了一个有益的归纳偏置，确保所有生成路径都与合理的轨迹对齐。
- 梯度优化：使用交叉熵损失代替均方误差，改善了训练过程中的梯度行为。
2.4 VFM与其他生成模型的联系
- 与流匹配的关系：作者证明，当变分近似选择为高斯分布，并且其协方差与条件向量场的线性部分相匹配时，VFM目标函数可以精确地退化为标准的流匹配损失。这表明VFM是FM的一个更通用的框架。
- 与分数匹配的关系：作者进一步证明，VFM不仅可以用来学习向量场，还可以用来近似得分函数$\nabla_x \log p_t(x)$。这意味着VFM框架可以同时学习确定性动态（ODE）和随机动态（SDE）。此外，一个重加权的VFM目标还为模型的对数似然提供了一个变分下界。
3. 实验分析
作者在一系列图生成任务上评估了CatFlow的性能。
3.1 抽象图生成
作者在四个抽象图生成数据集上进行了测试，包括Ego-small, Community-small, Enzymes和Grid。评估指标是生成图与测试图在度分布、聚类系数分布和轨道计数分布上的最大均值差异（MMD）。
在Ego-small和Community-small两个任务上，CatFlow的性能都达到了最先进水平，与GDSS等基于扩散的模型相当，并显著优于GraphVAE和GNF等早期方法。
3.2 分子生成
作者在两个流行的分子生成基准QM9和ZINC250k上评估了CatFlow。评估指标包括有效性（Validity）、独特性（Uniqueness）和Fréchet ChemNet Distance (FCD)。
CatFlow在QM9和ZINC250k两个数据集上都取得了当前最先进的结果。在QM9上，CatFlow的有效性达到99.81%，独特性达到99.95%，FCD仅为0.441，显著优于所有基线方法，包括MoFlow, GDSS和DiGress。在ZINC250k上，CatFlow也取得了最高的有效性（99.21%）和独特性（100%），以及最低的FCD（13.211）。
图展示了CatFlow生成的分子样本，它们在化学结构上是合理且多样的。值得注意的是，CatFlow不仅性能优越，而且收敛速度比标准的流匹配更快，计算成本也不高于任何基线方法。
3.3 CatFlow消融研究
为了比较CatFlow和标准流匹配的性能差异，作者进行了一系列消融实验，考察了模型在不同参数数量和不同训练数据量下的泛化能力。
图a, b, c分别展示了在100%、20%和5%的训练数据下，CatFlow和标准流匹配的性能（以有效且独特的分子百分比作为分数）。结果清晰地显示，CatFlow不仅在全数据、大模型的情况下优于标准FM，而且在数据量和模型大小减小时，其性能下降得更慢，表现出更强的鲁棒性和泛化能力。
