技术摘要：基于通量匹配的生成建模

问题陈述

当前最先进的生成建模方法，特别是在图像合成、蛋白质设计和机器人学等高维领域，严重依赖基于分数的模型（score-based models）。这些模型通常通过分数匹配（score matching）学习Stein 分数函数 $\nabla \log p_{\text{data}}(x)$。一旦学习完成，该分数即被用于驱动扩散过程（例如朗之万动力学），以从目标分布 $p_{\text{data}}$ 中采样。

然而，分数函数并非生成 $p_{\text{data}}$ 的唯一向量场。根据福克 - 普朗克方程（FPE），任何满足条件 $\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$ 的向量场 $f_\theta$，其稳态分布均为 $p_{\text{data}}$。这意味着存在一个无限的生成向量场族（非保守场），它们可以生成相同的目标分布，但诱导不同的动力学特性（例如不同的混合速率、循环模式或有向依赖关系）。

现有的分数匹配目标迫使模型收敛到该空间中的单一点（$\nabla \log p_{\text{data}}$），即使某种偏差保留了目标分布，也会对其施加惩罚。这限制了将有用的归纳偏置、结构先验或动力学特性（如更快的混合或有向依赖）直接编码到生成场中的能力。

方法论：通量匹配

作者提出了通量匹配（Flux Matching），这是一种新的范式，通过学习目标分布生成场族内的任意向量场来推广分数匹配，而不是强制模型逐点匹配分数。

核心目标

通量匹配最小化概率通量的散度，而不是最小化模型 $f_\theta$ 与分数 $\nabla \log p_{\text{data}}$ 之间的 Fisher 散度。一个场具有生成性的条件为：
$$\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$$
该条件比逐点相等更弱，允许存在一个包含无限多个有效向量场的“零空间”。

理论推导

为了优化该条件，作者定义了一个投影 Fisher 散度：
$$\tilde{J}(\theta) := \mathbb{E}_{x \sim p_{\text{data}}} \left[ \| \Pi_{\text{flux}} f_\theta(x) - \nabla \log p_{\text{data}}(x) \|^2 \right]$$
其中 $\Pi_{\text{flux}}$ 将向量场投影到具有相同通量散度的梯度场空间。虽然 $\tilde{J}(\theta)$ 在理论上是合理的，但在高维情况下是不可处理的。

作者通过以下步骤推导出了可处理的通量匹配损失（$L_{\text{flux}}$）：

1. 定义残差场 $u_\theta = f_\theta - \nabla \log p_{\text{data}}$ 和朗之万 - Stein 算子 $r_\theta = \nabla \cdot u_\theta + u_\theta \cdot \nabla \log p_{\text{data}}$。
2. 利用扩散模拟 $d x_t = \nabla \log p_{\text{data}}(x_t) dt + \sqrt{2} dW_t$ 在数据分布上传播 $r_\theta$。
3. 应用分部积分和停止梯度（stop-gradient）技巧，将目标转换为一种避免通过昂贵的路径导数进行反向传播的形式。

最终损失为：
$$L_{\text{flux}}(\theta) := -\mathbb{E}_{t \sim q, x_0 \sim p_{\text{data}}, x_t | x_0} \left[ \frac{1}{q(t)} u_\theta(x_0)^\top \text{sg}\left( \frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t) \right) \right]$$
其中 $\text{sg}$ 表示停止梯度算子。作者证明了 $\nabla_\theta \tilde{J}(\theta) = 2 \nabla_\theta L_{\text{flux}}(\theta)$。

噪声退火与实际实现

该方法扩展到噪声退火（noise-annealed）设置（类似于扩散模型），通过为噪声水平 $\sigma$ 学习连续的场。

• 分数近似：由于 $\nabla \log p_{\text{data}}$ 未知，作者在小批量上使用非参数核密度估计（KDE）近似来估计分数。
• 方差缩减：为了估计路径项 $\frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t)$，作者采用跨链估计器，利用小批量中的所有样本并以高斯转移概率加权，而不是依赖单条链。

选择特定场

通量匹配学习的是一个场族。为了选择具有所需特性的特定场，作者提出了两种策略：

1. 特定应用损失：在通量匹配目标中添加辅助损失（例如针对混合速度或结构约束的损失）。
2. 模型参数化：通过架构约束（例如因果注意力掩码）将特定结构直接强制到 $f_\theta$ 中。

主要贡献

1. 通量匹配范式：引入了一种生成建模框架，通过匹配概率通量的散度来学习分数之外的向量场。
2. 高效损失推导：推导出了可扩展的通量匹配损失，在保留 Fisher 散度几何结构的同时，允许非保守向量场。
3. 噪声退火扩展：将该方法扩展到噪声退火分布，使其能够应用于复杂的高维图像数据集。
4. ** demonstrated 灵活性**：实证证据表明，通量匹配实现了标准分数匹配无法完成的应用，包括：
   • 更快的混合：学习加速采样收敛的向量场。
   • 可解释性：拟合结构化的、具有生物学意义的模型（例如 RNA 速度），而无需黑盒神经近似。
   • 结构化动力学：将有向依赖（例如时间自回归）直接嵌入到生成场中。

结果

该论文在五个设置中评估了通量匹配：

1. 可控场（玩具数据）：在二维高斯混合模型上，通量匹配对保留目标分布但具有不同动力学（例如旋转、循环）的场赋予零损失，而分数匹配则对这些偏差施加惩罚。
2. 可解释场（RNA 速度）：应用于单细胞基因组学，通量匹配通过梯度下降拟合 scVelo 动力学模型参数。在五个数据集上，其在一致性和跨边界正确性方面优于标准的基于 EM 的 scVelo 拟合，证明收益来自优化过程而非模型表达能力的增加。
3. 无限制生成（图像）：在 CIFAR-10 和 CelebA 上，通量匹配作为独立目标。虽然它取得了具有竞争力的 FID 分数，但与去噪分数匹配（DSM）相比，其训练速度目前慢 3–4 倍，且内存使用量多 2–3 倍。然而，它能够扩展到高维分布。
4. 快速混合：通过优化混合代理，通量匹配学习到的场能够以比传统通量匹配或 DSM 少得多的采样步骤达到合理的 FID 分数。
5. 嵌入结构：在轨迹生成任务（非线性弹簧）中，通量匹配成功结合了因果注意力掩码以强制时间自回归。相比之下，将相同的掩码应用于 DSM 会损害性能，因为分数函数的对称雅可比矩阵与有向依赖相冲突。

意义与主张

作者声称，通量匹配通过将向量场本身视为一种设计选择而非固定目标，在生成建模中开辟了一个“新维度”。

• 超越采样：主要价值不一定是在标准图像生成中取代 DSM（DSM 在此领域已高度优化），而是实现具有有用特性的动力学。
• 灵活性：它允许从业者将归纳偏置、结构先验和机制约束直接强加到生成过程中。
• 应用：论文强调了这种灵活性所解锁的具体用例：通过不可逆动力学实现更快的采样、可解释的科学模型（如 RNA 速度），以及具有指定因果或时间结构的生成模型。

该论文得出结论，尽管通量匹配目前会产生更高的计算成本，但其学习非保守、结构化生成场的能力，为那些生成动力学与输出分布同样重要的应用提供了一种强大的工具。