基于最优传输势的多边缘流匹配

# 基于最优输运势的多边际流匹配 来源：https://arxiv.org/html/2606.05327 ###### 摘要 流匹配（Flow Matching，FM）已成为学习两个经验分布之间动态输运映射的强大框架。然而，中间观测边际分布（可用于约束端点之间的流）的场景研究较少。这种“多边际”机制对于许多科学领域中动态系统的时序演化建模至关重要，因为这些领域可以采样序贯分布。我们通过一种新颖方法解决该问题，该方法利用FM与动态最优输运（OT）之间的联系，通过在动态OT作用量中引入势项，柔和地将流引导向中间边际分布。通过将条件FM学习目标扩展以纳入这些势，我们推导出一种高效、无模拟的多边际FM算法，该算法在学习流的时空动态方面提供了相当大的灵活性。我们在多种单细胞RNA测序、海洋学和气象学数据集上展示了OT势FM（OTP-FM）的最先进性能和训练效率。我们的代码可在 https://github.com/Bexorg-Inc/OTP-FM 获取。 机器学习，ICML ## 1 引言 理解物理系统复杂的非线性动态在许多科学领域至关重要，包括发育生物学中的转录组状态转变、神经退行性疾病中的疾病进展以及气候建模。这些学科在构建系统演化过程中的静态快照集方面取得了显著进展，例如纵向单细胞RNA测序（scRNA-seq）测量；然而，从这些通常是独立的快照中推断出准确的每个样本轨迹，对于机制理解、治疗靶点识别和预测性推理仍然是一个关键挑战。 条件流匹配（CFM）[Lipman 等人，2023 (https://arxiv.org/html/2606.05327#bib.bib1)；Tong 等人，2024 (https://arxiv.org/html/2606.05327#bib.bib3)；Liu 等人，2023 (https://arxiv.org/html/2606.05327#bib.bib4)；Albergo 等人，2025 (https://arxiv.org/html/2606.05327#bib.bib2)] 已成为解决该问题的领先框架，通过简单的条件轨迹的高效、无模拟回归来学习连续时间输运映射，这些轨迹通常是配对源样本和目标样本之间的条件最优输运（OT）解。当中间边际分布可用时，自然的扩展是在连续边际分布之间逐段应用 CFM，将条件路径端到端拼接起来（图 1 (https://arxiv.org/html/2606.05327#S1.F1)）。然而，这会在每个边际边界处产生不具物理连续性的轨迹。最近的多边际方法，如 MMFM [Rohbeck 等人，2025 (https://arxiv.org/html/2606.05327#bib.bib54)] 和 3MSBM [Theodoropoulos 等人，2026 (https://arxiv.org/html/2606.05327#bib.bib55)]，试图平滑这些边界，但采用了我们认为不必描述物理系统的规定性启发式策略。我们提出了一种有理论基础的松弛：我们证明逐段 CFM 对应于对每个中间边际施加*硬约束*的条件动态 OT，并将这些约束松弛为动态 OT 作用量中的平滑、有限强度势能项，即 OT 势 FM（OTP-FM）——将逐段 CFM 作为极限情况恢复。我们的贡献如下： - •多边际动态 OT 的推广，具有精确条件解和对真实中间边际分布对齐的理论界限，由势强度和训练损失控制。 - •灵活、无模拟的训练算法，在势参数方面具有广泛的设计空间，有效允许通过在此空间上进行优化来让数据确定插值动态，而不是先前方法的规定性方法。 - •在多样化的科学数据集上实现最先进（SOTA）的性能和训练效率，并进行系统的消融研究以及将 OTP-FM 应用于新数据集的具体建议。 参见图注 图 1：（左）比较标准 CFM——忽略中间边际分布的直线轨迹；多边际 CFM——在连续边际分布之间逐段拼接 CFM 轨迹；规定性方法如 MMFM 和 3MSBM——使用固定插值策略平滑尖点；以及 OTP-FM，其具有可调强度 ww、时间宽度 ττ 和 λλ 形状的软势驱动动态，产生平滑*且*灵活的轨迹。（右）方法概述：轨迹 XtX_{t} 分解为基础 CFM 路径 XtbaseX_{t}^{\mathrm{base}} 加上边际驱动的修正项 Xk,tcorrX^{\mathrm{corr}}_{k,t}，这些修正项由课程参数 αα 逐渐缩放，收敛到 OTP-FM 解。 ## 2 背景与预备知识 ### 2.1 最优输运 原始的静态 Monge OT 问题求解源测度 μ0\mu_{0} 和目标测度 μ1\mu_{1} 之间的输运映射 ψ:Rd→Rd\psi\!:\mathbb{R}^{d}\to\mathbb{R}^{d}，使得推前操作 ψ#μ0=μ1\psi_{\#}\mu_{0}=\mu_{1}，并且相对于代价 c:Rd×Rd→Rc\colon\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R} 来说是*最优*的 (Villani, 2009 (https://arxiv.org/html/2606.05327#bib.bib7))：LOT:=minψ:ψ#μ0=μ1∫c(x,ψ(x))dμ0(x)\mathcal{L}_{\mathrm{OT}}:=\min_{\psi:\psi_{\#}\mu_{0}=\mu_{1}}\int\!c(x,\psi(x))\mathrm{d}\mu_{0}(x)。对于代价 cp(x,y)=‖x−y‖pc_{p}(x,y)=\|x-y\|^{p}，LOT1/p\mathcal{L}_{\mathrm{OT}}^{1/p} 是 Wasserstein pp-距离 Wp\mathcal{W}_{p}。特别相关的是动态 OT（DOT）公式 [Benamou and Brenier, 2000 (https://arxiv.org/html/2606.05327#bib.bib6)]，其中我们定义动态概率路径 ρt:[0,1]×Rd→R+\rho_{t}\colon[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{+}，插值 μ0\mu_{0} 到 μ1\mu_{1}。该路径由速度场 utu_{t} 通过连续性方程 ∂tρ=−∇⋅(ρu)\partial_{t}\rho=-\nabla\cdot(\rho u) 产生，样本轨迹满足 Ẋt=ut(Xt)\dot{X}_{t}=u_{t}(X_{t})，其中 X0∼ρ0X_{0}\sim\rho_{0}，相应的映射或流 ψt:[0,1]×Rd→Rd\psi_{t}\colon[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d} 将样本沿轨迹 Xt=ψt(X0)∼ρtX_{t}=\psi_{t}(X_{0})\sim\rho_{t} 输运。平方欧氏代价 c2/2\nicefrac{{c_{2}}}{{2}} 的目标是：111我们采用惯例因子 1/2\nicefrac{{1}}{{2}} 以保持与物理动能的类比。 LDOT:=minρt,ut∫01∫Rd12‖ut(x)‖2dρt(x)dt,\mathcal{L}_{\mathrm{DOT}}:=\min_{\rho_{t},u_{t}}\int_{0}^{1}\int_{\mathbb{R}^{d}}\frac{1}{2}\|u_{t}(x)\|^{2}\mathrm{d}\rho_{t}(x)\mathrm{d}t, (1) 满足连续性方程和边界条件 ρ0=μ0\rho_{0}=\mu_{0}，ρ1=μ1\rho_{1}=\mu_{1}。有趣的是，这可以解释为最小化具有拉格朗日量 L=T−VL=T-V 的流体的作用量 S[L]S[L]，其中动能 T=12ρ‖u‖2T=\frac{1}{2}\rho\|u\|^{2}，势能 V=0V=0： S[L]=∫01∫Rd[12ρt(x)∥ut(x)∥2−V(x,t)+φt(x)[∂tρt(x)+∇⋅(ρt(x)ut(x))]]dxdt,S[L]=\int_{0}^{1}\int_{\mathbb{R}^{d}}\Big[\frac{1}{2}\rho_{t}(x)\|u_{t}(x)\|^{2}-V(x,t)+\\ \varphi_{t}(x)[\partial_{t}\rho_{t}(x)+\nabla\cdot(\rho_{t}(x)u_{t}(x))]\Big]\mathrm{d}x\mathrm{d}t, (2) 其中 φ\varphi 是施加连续性方程的拉格朗日乘子。对于 V=0V=0，最小化者满足 Euler-Lagrange（E-L）方程： ut=∇φt,∂tφt(x)+‖∇φt(x)‖22=0,u_{t}=\nabla\varphi_{t},\qquad\partial_{t}\varphi_{t}(x)+\frac{\|\nabla\varphi_{t}(x)\|^{2}}{2}=0, (3) 并遵循直线轨迹 Xt=(1−t)X0+tψ(X0)X_{t}=(1-t)X_{0}+t\psi(X_{0})，Ẋt=ψ(X0)−X0\dot{X}_{t}=\psi(X_{0})-X_{0}，其中 ψ\psi 是代价 c2c_{2} 的静态 OT 映射。详细推导见附录 A.1 (https://arxiv.org/html/2606.05327#A1.SS1)。这些直线最小化者对应于 CFM 中常用的条件路径，我们接下来将描述，而在 OTP-FM 中我们探索非零 VV 的情况。 ### 2.2 条件流匹配 条件流匹配（CFM）同样旨在通过学习参数化速度 utθu^{\theta}_{t} 来找到测度之间的流。获得一个有效的*边际*训练目标来进行回归通常是难以处理的；然而，Lipman 等人 (2023 (https://arxiv.org/html/2606.05327#bib.bib1))；Tong 等人 (2024 (https://arxiv.org/html/2606.05327#bib.bib3)) 表明，我们可以通过将边际 ρt(x)=∫ρt(x|z)q(z)dz\rho_{t}(x)=\int\rho_{t}(x|z)q(z)\mathrm{d}z 构造为条件路径 ρt(x|z)\rho_{t}(x|z) 的混合来推导一个等价的、更简单的*条件*目标，其中关联的条件速度 ut(x|z)u_{t}(x|z) 以潜在变量 z∼q(z)z\sim q(z) 为条件，并满足 ρ0(x)=μ0(x)\rho_{0}(x)=\mu_{0}(x)，ρ1(x)=μ1(x)\rho_{1}(x)=\mu_{1}(x)。最常用且最简单的方法是选择 z=(x0,x1)∼π(x0,x1)z=(x_{0},x_{1})\sim\pi(x_{0},x_{1})，其中 π\pi 是端点的某个联合分布，而 ut(x|z)u_{t}(x|z) 是 ρ0(x|z)=δ(x−x0)\rho_{0}(x|z)=\delta(x-x_{0}) 和 ρ1(x|z)=δ(x−x1)\rho_{1}(x|z)=\delta(x-x_{1}) 之间的条件动态 OT 解（对于 V=0V=0），定义 CFM 目标： Xt(x|z)=(1−t)x0+tx1,\displaystyle X_{t}(x|z)=(1-t)x_{0}+tx_{1}, (4) ut(x|z)=Ẋt(x|z)=x1−x0,\displaystyle u_{t}(x|z)=\dot{X}_{t}(x|z)=x_{1}-x_{0}, (5) LCFM(θ):=Et,z,x∼ρt(x|z)∥uθ(t,x)−ut(x|z)∥2.\displaystyle\mathcal{L}_{\mathrm{CFM}}(\theta):=\mathbb{E}_{t,z,x\sim\rho_{t}(x|z)}\left\|u^{\theta}(t,x)-u_{t}(x|z)\right\|^{2}. (6) 令人惊讶的是，这个简单得多的目标与边际 FM 目标梯度等价，从而提供了一种高效且可扩展的学习流的训练算法。在 OTP-FM 中，我们将这个回归目标推广到非零势 VV，灵活地纳入中间边际约束。 ### 2.3 少步和一致性模型 正如我们将在第 4.1 节 (https://arxiv.org/html/2606.05327#S4.SS1) 中描述的，OTP-FM 可能需要在训练期间评估中间时间的样本位置 XtkX_{t_{k}}。为了避免为此进行昂贵的 ODE 模拟，我们训练一个*一致性模型*用于少步推理。广义上，一致性模型旨在学习（的变体）流映射 Ψt1,t2:[0,1]2×Rd→Rd\Psi_{t_{1},t_{2}}\colon[0,1]^{2}\times\mathbb{R}^{d}\to\mathbb{R}^{d}——这是 ψt\psi_{t} 的推广，可以在两个任意时间点 t1t_{1} 和 t2t_{2} 之间输运样本：Ψt1,t2(Xt1)=Xt2\Psi_{t_{1},t_{2}}(X_{t_{1}})=X_{t_{2}}，边界条件 Ψt,t(x)=x\Psi_{t,t}(x)=x [Boffi 等人，2025 (https://arxiv.org/html/2606.05327#bib.bib16)，2026 (https://arxiv.org/html/2606.05327#bib.bib17)]。虽然 OTP-FM 对用于学习 Ψt1,t2\Psi_{t_{1},t_{2}} 的特定训练过程不敏感，但在我们的实验中，我们主要采用*改进的 MeanFlow*（iMF）[Geng 等人，2026b (https://arxiv.org/html/2606.05327#bib.bib76)]，在撰写本文时，该模型在一步和两步推理中达到了 SOTA。即，我们根据 t1t_{1} 和 t2t_{2} 之间的*平均速度*来参数化 Ψt1,t2\Psi_{t_{1},t_{2}}，vt1,t2(Xt1)=1t2−t1∫t1t2ut(Xt)dtv_{t_{1},t_{2}}(X_{t_{1}})=\frac{1}{t_{2}-t_{1}}\int_{t_{1}}^{t_{2}}u_{t}(X_{t})\mathrm{d}t，由此导出回归目标： LiMF(θ):=Et1,t2,z,x∼ρt1(⋅|z)∥Vt1,t2θ(x)−ut1(x|z)∥2\displaystyle\mathcal{L}_{\mathrm{iMF}}(\theta):=\mathbb{E}_{t_{1},t_{2},z,x\sim\rho_{t_{1}}(\cdot|z)}\!\left\|V^{\theta}_{t_{1},t_{2}}(x)-u_{t_{1}}\!(x|z)\right\|^{2}\!\!\!\!\!\!\! (7) Vt1,t2θ(x)≡vθ(x)−(t2−t1)sg[vt1,t1θ(x)∂xvθ+∂t1vθ],\displaystyle V^{\theta}_{t_{1},t_{2}}(x)\equiv v^{\theta}\!(x)-(t_{2}-t_{1})\,\mathrm{sg}\!\left[v^{\theta}_{t_{1},t_{1}}\!(x)\,\partial_{x}v^{\theta}+\partial_{t_{1}}v^{\theta}\right], 其中 sg\mathrm{sg} 是停止梯度算子，ut(x|z)u_{t}(x|z) 是 CFM 中使用的*瞬时*条件速度目标，而 Vt1,t2θV^{\theta}_{t_{1},t_{2}} 是直接回归该目标的模型 vt1,t2θv^{\theta}_{t_{1},t_{2}} 的参数化。222我们的公式与 Geng 等人 (2026b (https://arxiv.org/html/2606.05327#bib.bib76)) 的原始公式大部分等价，除了训练目标修改为沿时间*向前*流动。细节见附录 C.2 (https://arxiv.org/html/2606.05327#A3.SS2)。关于一致性方法的进一步讨论以及我们 iMF 目标的推导见附录 C (https://arxiv.org/html/2606.05327#A3)，附录 G (https://arxiv.org/html/2606.05327#A7) 中给出了替代方法的结果，展示了 OTP-FM 对一致性模型选择的灵活性。 ## 3 动态 OTP 问题 参见图注 图 2：顶部：对于具有不同势、强度和 λk(t)\lambda_{k}(t) 的一维高斯边际分布的边际动态 OTP 问题的精确解。底部：相同边际分布和势的 OTP-FM 解，除了最右边的图，它展示了 D=W2∞\mathcal{D}=\mathcal{W}_{2}^{\infty}。我们现在将 CFM 推广到多边际轨迹推理问题：给定在时间 tk∈[0,1]t_{k}\in[0,1] 的经验边际分布 {μtk}k=0K+1\{\mu_{t_{k}}\}_{k=0}^{K+1}（两个端点和 KK 个中间点），我们旨在学习一个速度场 utθu^{\theta}_{t}，其流产生的密度 ρt\rho_{t} 与 μtk\mu_{t_{k}} 对齐，同时描述物理上合理的插值轨迹，通过与保留的边际分布的对齐程度来衡量（第 6 节 (https://arxiv.org/html/2606.05327#S6)）。我们首先观察到，标准的逐段 CFM 目标可以被重新解释为具有*硬*惩罚项或奇异势的条件动态 OT 解（第 3.1 节 (https://arxiv.org/html/2606.05327#S3.SS1)）。这激发了向所谓动态 OT + 势（OTP）问题的平滑松弛（第 3.2 节 (https://arxiv.org/html/2606.05327#S3.SS2)），其中中间边际约束以*软*势项的形式出现在 OT 作用量中。最后，我们推导出样本轨迹并讨论势的设计空间（第 3.3 节 (https://arxiv.org/html/2606.05327#S3.SS3) 和 3.4 节 (https://arxiv.org/html/2606.05327#S3.SS4)）。 ### 3.1 逐段 CFM 作为具有奇异势的 OT 标准的 CFM 多边际扩展在连续边际分布之间逐段拼接条件 OT 解（例如 Tong 等人 (2024 (https://arxiv.org/html/2606.05327#bib.bib3)，2023 (https://arxiv.org/html/2606.05327#bib.bib69))，见附录 A.2 (https://arxiv.org/html/2606.05327#A1.SS2)），其目标密度在每个中间时间精确满足 ρtk=μtk\rho_{t_{k}}=\mu_{t_{k}}。这相当于在条件动态 OT 问题（方程 1 (https://arxiv.org/h