通过非梯度向量流的流图学习

# 通过非梯度向量流学习流映射 来源：https://arxiv.org/html/2607.26398 Mark Goldstein¹, Anshuk Uppal², Raghav Singhal³, Aahlad Puli³, & Rajesh Ranganath³ ¹计算数学中心，Flatiron Institute。 ²应用数学与计算机科学系，丹麦技术大学。 ³Courant研究所，纽约大学。 ###### 摘要 扩散模型和基于流的模型受益于简单的回归损失，但推理会产生显著开销，因为采样需要积分。一致性模型通过直接学习ODE轨迹上的流映射来解决这个问题，开创了单步和多步方法之间的设计空间。然而，现有方法面临计算挑战，例如需要模型逆或通过迭代模型调用进行反向传播，并且并不总能证明所需的ODE流映射是损失的一个解。我们提出了SGFlow，一种学习流映射的方法，它绕过了显式的可逆性约束和通过模型迭代进行昂贵微分的需求。SGFlow训练一个模型从零开始计算ODE解和隐含速度，遵循非保守动力学，其驻点位于所需的流映射处。在CIFAR图像基准上，没有单一方法能在所有步数下获得最佳FID：SGFlow在10个采样步数下获得最佳FID，并且在其他步数下与flow matching、Meanflow和Lagrangian map matching保持竞争力，同时是唯一一个为其基于stopgrad的动力学提供驻点保证的方法。 ## 1 引言 扩散模型和流模型（Sohl-Dickstein等人，2015 (https://arxiv.org/html/2607.26398#bib.bib22)；Ho等人，2020 (https://arxiv.org/html/2607.26398#bib.bib29)；Song等人，2020 (https://arxiv.org/html/2607.26398#bib.bib23)；Kingma等人，2021 (https://arxiv.org/html/2607.26398#bib.bib10)；Albergo and Vanden-Eijnden,2022 (https://arxiv.org/html/2607.26398#bib.bib20)；Singhal等人，2023 (https://arxiv.org/html/2607.26398#bib.bib30)；Pandey and Mandt,2023 (https://arxiv.org/html/2607.26398#bib.bib32)；Bartosh等人，2024 (https://arxiv.org/html/2607.26398#bib.bib33)；Singhal等人，2024 (https://arxiv.org/html/2607.26398#bib.bib31)；Albergo等人，2023 (https://arxiv.org/html/2607.26398#bib.bib21)；Lipman等人，2022 (https://arxiv.org/html/2607.26398#bib.bib46)；Liu等人，2022 (https://arxiv.org/html/2607.26398#bib.bib47)）在蛋白质（Abramson等人，2024 (https://arxiv.org/html/2607.26398#bib.bib19)）和图像（Peebles and Xie,2023 (https://arxiv.org/html/2607.26398#bib.bib27); Esser等人，2024 (https://arxiv.org/html/2607.26398#bib.bib13)）等领域改进了生成。从这些模型中采样通常需要对常微分方程或随机微分方程进行数值积分。数值积分需要多次前向传播神经网络，导致采样延迟和成本增加。为了通过改变训练来改善这种生成成本，最近用于一致性建模和映射匹配的方法（Song等人，2023 (https://arxiv.org/html/2607.26398#bib.bib9)；Song and Dhariwal,2023 (https://arxiv.org/html/2607.26398#bib.bib8)；Kim等人，2023 (https://arxiv.org/html/2607.26398#bib.bib5)；Lu and Song,2024 (https://arxiv.org/html/2607.26398#bib.bib11)；Boffi等人，2024 (https://arxiv.org/html/2607.26398#bib.bib44)；2025 (https://arxiv.org/html/2607.26398#bib.bib38)）旨在学习从噪声到概率流ODE轨迹上中间或最终数据点的直接映射，从而避免昂贵的积分。然而，这些方法各有其复杂性。例如，流映射匹配需要模型可逆性，而一致性模型要么需要一步映射，要么引入使目标ODE轨迹偏离的额外步骤。 我们提出了SGFlow（代表StopGrad Flow），一种基于流和映射匹配方法的方法，并且： - **具有真实流映射作为唯一驻点** - **不限制所使用的神经网络类别（例如，可逆函数）** - **不需要涉及可逆性或对抗优化的辅助损失** - **不需要通过模型嵌套调用进行优化** - **允许沿ODE轨迹以任意步数进行生成** 现有的流映射学习方法在挑战上可分为几类；所有挑战都与流映射由某些导数性质表征以及损失最小化平方误差以使这些性质成立这一思想有关。流映射匹配和相关方法依赖于可逆映射与常微分方程（ODE）之间的基本关系。这种关系通常需要在训练期间显式计算前向映射（正在训练的模型）及其逆，从而复杂化训练，或者需要通过嵌套模型调用进行昂贵的反向传播。Boffi等人（2025 (https://arxiv.org/html/2607.26398#bib.bib38)）针对Boffi等人（2024 (https://arxiv.org/html/2607.26398#bib.bib44)）的流映射匹配损失提出了stopgrad放置策略，绕过了这种昂贵的嵌套微分，但并未证明stopgrad能保留真实流映射处的驻点。Meanflow（Geng等人，2025 (https://arxiv.org/html/2607.26398#bib.bib36)）没有显式强制执行模型逆恒等式，并完全避免了通过前向模式导数的反向传播，在低步数下实现了良好的图像生成性能；Meanflow也未被证明在真实流映射处有驻点。 SGFlow通过利用仅涉及雅可比向量积（JVP）而没有逆函数的替代恒等式，避免了跟踪模型及其逆的复杂性。这个恒等式使我们能够纯粹根据前向映射来制定目标，而无需显式访问其逆。由于ODE的解自然产生可逆映射，SGFlow目标隐式地鼓励可逆性，而无需显式强制执行。因此，在最优情况下，SGFlow生成一个连续可微函数，该函数精确地积分速度场，直接生成所需的数据分布。我们在表1 (https://arxiv.org/html/2607.26398#S1.T1)和6 (https://arxiv.org/html/2607.26398#S6)节中总结了近期方法之间的权衡。 实验上，对于使用相同通用架构的基本训练设置，我们询问当减少采样步数时，flow matching、Meanflow、SGFlow和Lagrangian map matching在中等维度（CIFAR-10）上的无条件指标（FID）如何比较。 表1：与先前工作的比较。我们根据以下标准对流映射学习（或一致性建模）技术以及我们提出的SGFlow方法进行分类：（1）训练后调整采样步数的能力，（2）是否遵循PF-ODE（Song等人，2020 (https://arxiv.org/html/2607.26398#bib.bib23)），（3）是否允许无模拟训练，（4）其目标是否使用回归，（5）训练是否无需模型逆，（6）是否证明真实流映射是最优或驻点，以及（7）训练是否避免通过嵌套模型调用的微分。详见第6节 (https://arxiv.org/html/2607.26398#S6)。 ## 2 背景 随机插值法（Lipman等人，2022 (https://arxiv.org/html/2607.26398#bib.bib46)；Albergo等人，2023 (https://arxiv.org/html/2607.26398#bib.bib21)），以及更广泛地说，大多数扩散和流方法（此后简称为流），将生成建模视为从简单基础密度到目标密度的运输。插值法按如下方式处理该问题。对于 \(t \in [0,1]\)： 1. 1. 选择 \((\alpha_t, \sigma_t)\)，其中 \(\alpha_0 = \sigma_1 = 1\) 且 \(\alpha_1 = \sigma_0 = 0\)。通常，\(\alpha_t = 1 - t\) 且 \(\sigma_t = t\)。 2. 2. 定义 \(X_t = \alpha_t X_0 + \sigma_t X_1\)，其中基础密度 \(X_0 \sim q_0\) 和数据 \(X_1 \sim q_1\)（或反之亦然）。 3. 3. 学习沿密度轨迹生成新样本。 对于函数 \(f\)，令 \(\dot{f}_t := \frac{d}{dt} f_t\)。因此 \(\dot{X}_t := \dot{\alpha}_t X_0 + \dot{\sigma}_t X_1\)。于是，\(X_t\) 具有密度 \(q_t\)，满足： \[ \partial_t q_t(x) = -\nabla_x \cdot (q_t(x) v(t,x)), \quad \quad v(t,x) := \mathbb{E}[\dot{X}_t ~|~ X_t = x], \tag{1} \] 其中 \(v\) 称为速度。方程1 (https://arxiv.org/html/2607.26398#S2.E1) 中的 PDE 在上述工作中推导得出。 为了完成第三步，首先观察到，当且仅当密度是概率流 ODE \(dX = v dt\) 从 \(X_0 \sim q_0\) 正向积分或从 \(X_1 \sim q_1\) 反向积分的解的密度时，该密度才满足方程1 (https://arxiv.org/html/2607.26398#S2.E1)（Albergo and Vanden-Eijnden, 2024 (https://arxiv.org/html/2607.26398#bib.bib39)）。然后，首先使用以下（无模拟）损失来近似 \(v\)： \[ \mathcal{L}_v(v_\theta) = \mathbb{E}\Big[ \| v_\theta(t, X_t) - (\dot{\alpha}_t X_0 + \dot{\sigma}_t X_1) \|^2 \Big]_{X_t = \alpha_t X_0 + \sigma_t X_1}, \tag{2} \] 其最小化器为 \(v_\theta = v\)，然后求解 \(dx = v_\theta dt\)。 #### 关于一致性方法的背景。 从流中采样需要积分，每一步都需要评估一个建模分数、速度或类似项的神经网络 \(v_\theta\)。直接知道 \(v\) 的积分原则上可以加速采样。一致性和映射匹配方法的目标是学习沿最优 \(v\) 隐含的轨迹进行映射。我们在此回顾一个例子，其他例子见第6节 (https://arxiv.org/html/2607.26398#S6)。Song等人（2023 (https://arxiv.org/html/2607.26398#bib.bib9)）；Song and Dhariwal（2023 (https://arxiv.org/html/2607.26398#bib.bib8)）试图学习一个映射 \(\hat{g}\)，该映射将插值样本 \(X_t \sim q_t\) 带到 \(\widehat{X}_0\)，即从 \(X_t\) 开始的 \(dx = v\,dt\) 在 \(t=0\) 的解（注意 \(\widehat{X}_0\) 通常与用于抽取 \(X_t\) 的端点 \(X_0\) 不同）。该损失度量两个相邻点处建模输出之间的距离。令 \(\text{SG}[\hat{g}]\) 表示 stopgrad。则： \[ \text{Consistency}(\hat{g}) := \mathbb{E}_{q(X_t)} [ \text{dist}(\hat{g}(t, X_t), \text{SG}[\hat{g}](t-\Delta t, \widehat{X}_{t-\Delta t})) ]. \tag{3} \] 目标 \(\widehat{X}_{t-\Delta t}\) 应来自从 \(X_t\) 开始积分真实速度一个小的步长 \(\Delta t\)，但由于 \(v\) 未知，通常使用预训练的 \(v_\theta\) 或与 \(\hat{g}\) 联合导出的 \(v_\theta\) 来近似——这增加了训练成本并引入了近似误差。允许多步采样需要重新加噪步骤，使轨迹偏离概率流 ODE，因此得到的更新不再对应于积分 PF-ODE。Kim等人（2023 (https://arxiv.org/html/2607.26398#bib.bib5)）观察到，这种多步方法“表现出随着 NFE 增加而退化的样本质量，缺乏计算预算（NFE）与样本保真度之间的清晰权衡”。后续工作引入了各种训练和推理时的修改，以弥合单步采样和多步采样之间的差距（Song等人，2023 (https://arxiv.org/html/2607.26398#bib.bib9)；Lu and Song,2024 (https://arxiv.org/html/2607.26398#bib.bib11)；Kim等人，2023 (https://arxiv.org/html/2607.26398#bib.bib5)；Boffi等人，2024 (https://arxiv.org/html/2607.26398#bib.bib44)；Sabour等人，2025 (https://arxiv.org/html/2607.26398#bib.bib37)；Geng等人，2025 (https://arxiv.org/html/2607.26398#bib.bib36)；Zhou等人，2025 (https://arxiv.org/html/2607.26398#bib.bib4)）；见第6节 (https://arxiv.org/html/2607.26398#S6)。 ## 3 方法 我们提出了 SGFlow，一种学习求解概率流 ODE 的方法，无需对抗训练、无需训练期间使用模型逆、无需表示显式导数矩阵，也无需来自预训练模型的昂贵模拟。SGFlow 训练一个模型从零开始计算 ODE 解和隐含速度，遵循非保守动力学。 考虑一个双时间映射 \(f\)，对于 \(t \leq u\)，通过求解概率流 ODE \(dx = v dt\) 将 \(X_t\) 带到 \(X_u\)。这样一个积分 \(v\) 的 \(f\) 可以定义如下： \[ f(t,u,x) = x + \int_t^u v(s, X_s) ds = x + \int_t^u v(s, f(t,s,x)) ds \tag{4} \] 使用总（物质）导数对 \(t\) 微分 RHS 上的递归形式得到： \[ \partial_t f + (\partial_x f) v(t,x) = 0, \quad f(u,u,x) = x \tag{5} \] 这在真实流映射 \(f\) 处唯一求解。我们可以对参数化的 \(f_\theta\) 平方左侧并取期望。通过抽取数据 \(X_1\)、噪声 \(X_0\) 并计算 \(X_t = \alpha_t X_0 + \sigma_t X_1\) 来采样 \(X_t\)： \[ L := \mathbb{E}_{X_t} [ \| \partial_t f_\theta + (\partial_x f_\theta) v \|^2 ] = \mathbb{E}_{X_t} [ \| \partial_t f_\theta + (\partial_x f_\theta) \mathbb{E}[\dot{X}_t | X_t] \|^2 ]. \tag{6} \] 真实映射 \(f\) 是该损失的唯一最小化器。使用 \(v(t,x) = \mathbb{E}[\dot{X}_t | X_t = x]\)，我们可以展开： \[ L = \mathbb{E}_{X_t} [ \| \partial_t f_\theta + (\partial_x f_\theta) \dot{X}_t \|^2 - \| (\partial_x f_\theta) (\dot{X}_t - \mathbb{E}[\dot{X}_t | X_t]) \|^2 ] \tag{7} \] 对于基础模型 \(\tilde{f}_\theta\)，我们可以使用参数化： \[ f_\theta(t,u,x) := x + (u - t) \tilde{f}_\theta(t,u,x) \tag{8} \] 该参数化自动满足边界条件 \(f_\theta(u,u,x) = x\)。该参数化产生两个额外性质： - • 时间导数：\(\partial_t f_\theta(t,t,x) = -\tilde{f}_\theta(t,t,x)\) - • 雅可比：\(\partial_x f_\theta(t,t,x) = I\) 利用这些性质并在 \(t=u\) 处评估，我们看到方程7 (https://arxiv.org/html/2607.26398#S3.E7) 的最小化简化为 flow matching，其中 \(\tilde{f}_\theta(t,t,x)\) 被训练以匹配速度： \[ L\big|_{t=u} = \mathbb{E}_{X_t} [ \| \tilde{f}_\theta(t,t,X_t) - \dot{X}_t \|^2 ], \tag{9} \] 这揭示了对于真实 \(f\)，我们有 \[ -\partial_t f(t,t,\cdot) = \tilde{f}(t,t,\cdot) = v(t,x) = \mathbb{E}[\dot{X}_t | X_t = x] \tag{10} \] 这启发我们用 \(\text{stopgrad}[\tilde{f}_\theta(t,t,\cdot)]\) 替换方程7 (https://arxiv.org/html/2607.26398#S3.E7) 中的未知 \(v\)。