动态广义Gromov-Wasserstein最优传输

# 动态广义Gromov-Wasserstein最优传输  
来源：https://arxiv.org/html/2609.20008  

Junda Ying††thanks:同等贡献  
Zhiwei Zeng11footnotemark:1  
*所属机构：北京大学定量生物学中心*  
Peijie Zhou††thanks:通讯作者: pjzhou@pku\.edu\.cn, zhangl@math\.pku\.edu\.cn  
*所属机构：北京大学定量生物学中心*  
*所属机构：北京大学机器学习研究中心*  
*所属机构：北京大数据分析与应用国家工程实验室*  
*所属机构：北京科学智能研究院*  
Lei Zhang22footnotemark:2  
*所属机构：北京国际数学研究中心，北京大学*  
*所属机构：北京大学定量生物学中心*  
*所属机构：北京大学机器学习研究中心*  
*所属机构：北京大学数学科学学院*  
*所属机构：北京大学人工智能研究院*  

###### 摘要  
Gromov–Wasserstein最优传输（GW-OT）通过引入结构感知传输代价扩展了经典最优传输。这对于空间转录组学尤为重要，因为动态重构除了匹配表达模式外，还应保持组织结构。尽管静态形式已广泛用于此类结构感知对齐，但重建连续轨迹的通用动态形式仍缺失。我们引入“旅行对动态对齐与轨迹估计”（TP-DATE），这是一个以无模拟方式推广GW-OT的理论和计算框架。我们通过路径作用量表述了一大类静态和动态二次型最优传输（QOT），并证明了静态-动态等价性。我们进一步开发了旅行对流匹配，该方法允许交互条件路径并将其相互作用边缘化为单一向量场。在合成和真实空间转录组学数据上，TP-DATE能更好地保持空间结构并提升连续三维动态重建性能。  

## 1 引言  
最优传输（OT）（Kantorovich, 1958 [20]）通过最小化传输代价为概率分布匹配提供了原则性方法，而其动态形式（Benamou and Brenier, 2000 [21]）将此端点匹配提升为测度的连续时间演化。这一视角使OT成为从群体快照重构动态的自然工具，包括无法纵向追踪单个细胞的单细胞系统。静态OT通过对时间点间的耦合推断来对齐快照，而动态OT则实现了观测快照间的连续插值与轨迹推断（Schiebinger et al., 2019 [31]；Tong et al., 2020 [34]；Tong et al., 2024a [19]；Sha et al., 2024 [16]；Zhang et al., 2025a [15]；Peng et al., 2026b [14]）。Schrödinger桥（Schrödinger, 1932 [25]；Léonard, 2014 [26]）和Wasserstein–Fisher–Rao（Chizat et al., 2018b [23]；Chizat et al., 2018a [22]；Liero et al., 2018 [24]）等扩展进一步将该框架推广到随机和非平衡动态。  

然而，当数据具有有意义的内部结构（如空间结构）时，仅有点态传输代价可能不足。Gromov–Wasserstein OT（GW-OT）比较分布内成对关系，因此特别适合结构感知匹配任务（Mémoli 2011 [35]）。融合GW-OT（FGW-OT）进一步结合了特征与结构信息（Vayer et al., 2020 [9]；Klein et al., 2025 [32]）。此类方法已广泛应用于图、异构域和空间组学等结构化数据。然而，大多数类GW的OT仍缺乏动态形式，因此只能对齐快照而无法重建连续时间动态。虽然当对齐不同模态的快照时，动态类GW形式可能无意义，但当快照代表同一结构化系统的不同状态时，它就变得至关重要。空间转录组学就是一个典型例子。不同时间采集的快照描述同一组织随时间演化，有意义的插值不仅要恢复基因表达变化，还需恢复空间组织的连续演化（Chen et al., 2022 [4]；Wei et al., 2022 [3]）。同一原理适用于连续组织切片，其中插值轴是组织深度而非时间，未观测的中间切片对应具有物理意义的状态（Zhang et al., 2026a [5]）。  

近期研究已从互补方向开始推广GW-OT。静态方面，Wang and Zhang (2025) [6] 定义了一族静态二次型最优传输（QOT），GW-OT作为特例。动态方面，内积GW-OT（IGW-OT）（Zhang et al., 2026b [7]）基于梯度流和黎曼几何为特定类GW-OT开发了动态形式并进行了研究。然而，通用动态QOT形式及其理论仍然缺失，高效无模拟的求解算法也尚未出现。  

为解决这些局限性，我们引入“旅行对动态对齐与轨迹估计”（TP-DATE），这是一个通用的动态QOT框架及配套的无模拟旅行对流匹配求解器。我们的贡献总结如下：  
- • 我们开发了动态QOT的数学理论，将标准OT、GW-OT和IGW-OT都纳入我们的QOT框架作为特例。  
- • 我们开发了无模拟旅行对流匹配框架，用于学习具有交互条件路径的动态并求解动态QOT问题。  
- • 我们提出了OT与GW-OT的动态融合，并展示了其在时空动态和三维结构重建等生物任务中的能力。  

## 2 相关工作  
**最优传输及其扩展**  
最优传输（Kantorovich, 1958 [20]）及其动态形式（Benamou and Brenier, 2000 [21]）已广泛应用于从快照数据重构动态。随机对应（Schrödinger, 1932 [25]；Léonard, 2014 [26]）和非平衡扩展（Chizat et al., 2018a [22]；Liero et al., 2018 [24]）进一步应用于建模更复杂的动态。另一重要的扩展方向是QOT，以GW-OT和FGW-OT等形式为代表（Mémoli, 2011 [35]；Vayer et al., 2020 [9]；Klein et al., 2025 [32]），可建模传输过程中的全局结构保持。最近，Wang and Zhang (2025) [6] 提出了静态QOT框架，统一了大类此类扩展；而Zhang et al. (2026b) [7] 开发了IGW-OT，作为特定QOT问题的动态形式。我们提出TP-DATE，这是一类通用QOT问题的统一动态形式，与现有形式兼容。  

**基于流匹配的最优传输求解器**  
流匹配（Lipman et al., 2023 [18]）是一种高效的无模拟生成建模框架，已用于求解OT（Tong et al., 2024a [19]）、Schrödinger桥（Tong et al., 2024b [28]）、WFR（Peng et al., 2026b [14]）以及多种基于OT的动态重构问题（Eyring et al., 2024 [11]；Rathod et al., 2026 [2]；Klein et al., 2024 [33]；Ying et al., 2026 [38]）。受条件路径技术和最优传输中旅行Dirac测度（Chizat et al., 2018a [22]）启发，我们从路径作用量视角表述动态QOT，并为此类问题开发了流匹配求解器。  

**基于OT的空间转录组学动态重构**  
近期已开发若干方法用于从空间转录组学推断时空动态。静态对齐方面，Klein et al. (2025) [32] 通过GW-OT和FGW-OT制定跨时间对齐，而Halmos et al. (2025) [1] 将时间和空间结构信息编码到静态传输目标中。动态方面，Rathod et al. (2026) [2] 将结构感知静态耦合整合到流匹配中。Peng et al. (2026a) [12] 直接将结构保持项引入动态传输目标，Zhang et al. (2025b) [13] 则通过基于模拟的NeuralODE（Chen et al., 2018 [17]）显式建模学习动态中的交互作用。这为TP-DATE——一个无模拟、直接建模结构感知传输动态的流匹配框架——留下了明确空白。  

## 3 预备知识  
**静态最优传输**  
设 \(\mathcal{P}(\mathcal{X})\) 为某个 \(X \subset \mathbb{R}^d\) 上所有概率密度的集合。\(\mu_0, \mu_1\) 是 \(\mathcal{P}(\mathcal{X})\) 中的概率密度。\(\Pi(\mu_0, \mu_1)\) 表示所有耦合的集合。  
\[
\Pi(\mu_0, \mu_1) = \left\{ \gamma \in \mathcal{P}(\mathcal{X}^2) \mid \int_{\mathcal{X}} \gamma(\bm{x}, \bm{y}) \mathrm{d}\bm{y} = \mu_0(\bm{x}), \int_{\mathcal{X}} \gamma(\bm{x}, \bm{y}) \mathrm{d}\bm{x} = \mu_1(\bm{y}) \right\}
\]  
静态最优传输（OT），也称为Kantorovich形式（Kantorovich, 1958 [20]），定义为  
\[
\text{OT}(\mu_0, \mu_1) = \inf_{\gamma \in \Pi(\mu_0, \mu_1)} \int_{\mathcal{X}^2} C(\bm{x}, \bm{y}) \gamma(\bm{x}, \bm{y}) \mathrm{d}\bm{x} \mathrm{d}\bm{y}
\]  
其中 \(C(\bm{x}, \bm{y})\) 表示将单位质量从 \(\bm{x}\) 运输到 \(\bm{y}\) 的代价，而 \(\gamma(\bm{x}, \bm{y}) \in \Pi(\mu_0, \mu_1)\) 称为耦合。从优化视角看，(1) 可视为关于耦合 \(\gamma\) 的线性规划。一个著名例子是当选择 \(C(\bm{x}, \bm{y}) = \frac{1}{2} \|\bm{x} - \bm{y}\|^2\)（即欧几里得距离平方）时，下确界称为2-Wasserstein距离平方（\(\mathcal{W}_2\)）。  

**静态二次型最优传输**  
近期，一个静态二次型最优传输（QOT）问题被定义并进行了数学研究（Wang and Zhang, 2025 [6]）。给定两个空间 \(\mathcal{X}, \mathcal{Y}\) 和一个代价函数 \(C: (\mathcal{X} \times \mathcal{Y})^2 \to \mathbb{R}\)，静态QOT定义为  
\[
\text{QOT}(\mu_0, \mu_1) = \inf_{\gamma \in \Pi(\mu_0, \mu_1)} \int_{\mathcal{X}^4} C(\bm{x}, \bm{y}, \bm{x}', \bm{y}') \gamma(\bm{x}, \bm{y}) \gamma(\bm{x}', \bm{y}') \mathrm{d}\bm{x} \mathrm{d}\bm{y} \mathrm{d}\bm{x}' \mathrm{d}\bm{y}'
\]  
若选择 \(C(\bm{x}, \bm{y}, \bm{x}', \bm{y}') = \|d_{\mathcal{X}}(\bm{x}, \bm{x}') - d_{\mathcal{Y}}(\bm{y}, \bm{y}')\|^2\)，其中 \((\mathcal{X}, d_{\mathcal{X}})\), \((\mathcal{Y}, d_{\mathcal{Y}})\) 是两个度量空间，则 (2) 恢复了Gromov-Wasserstein OT（GW-OT）。此类QOT已广泛应用于单细胞轨迹推断（Mémoli, 2011 [35]；Klein et al., 2025 [32]）。直观上，GW-OT在传输后保持局部结构。若选择 \(C(\bm{x}, \bm{y}, \bm{x}', \bm{y}') = f(\bm{x}, \bm{y}) + g(\bm{x}', \bm{y}')\)，则 (2) 退化为代价为 \(f+g\) 的标准静态OT。与静态OT不同，静态QOT是关于耦合 \(\gamma\) 的二次规划。  

**动态最优传输**  
对于静态OT，Benamou and Brenier (2000) [21] 为 \(\mathcal{W}_2\) 情形建立了动态形式，也称为BB形式。  
\[
\text{OT}(\mu_0, \mu_1) = \inf_{\rho, \bm{u}} \int_0^1 \int_{\mathcal{X}} \frac{1}{2} \|\bm{u}(\bm{x}, t)\|_2^2 \rho_t(\bm{x}) \mathrm{d}\bm{x} \mathrm{d}t
\]  
\[
\text{s.t.} \quad \partial_t \rho + \nabla_{\bm{x}} \cdot (\rho \bm{u}) = 0, \quad \rho_0 = \mu_0, \quad \rho_1 = \mu_1
\]  
他们证明了当 \(C(\bm{x}, \bm{y}) = \frac{1}{2} \|\bm{x} - \bm{y}\|^2\) 时，(3) 与 (1) 的等价性。直观上，它旨在找到连接 \(\mu_0, \mu_1\) 的连续概率流，同时最小化总动能。由于良好的流体动力学解释，它也已广泛应用于单细胞轨迹推断（Tong et al., 2020 [34]；Tong et al., 2024a [19]；Klein et al., 2024 [33]）。  

**旅行Dirac测度**  
为求解动态OT，可先考虑两个Dirac测度间的动态OT：  
\[
\text{OT}(\delta_{x_0}, \delta_{x_1}) = \inf_{x_t} \frac{1}{2} \int_0^1 \|\dot{x}_t\|_2^2 \mathrm{d}t \quad \text{s.t.} \quad x_0 = x_0, \, x_1 = x_1
\]