用于等变图生成的Gromov-Monge流匹配

# 面向等变图生成的 Gromov-Monge 流匹配
来源: https://arxiv.org/html/2608.26961

**作者信息：**
- **Moritz Piening**
    - 单位：柏林工业大学数学研究所
    - 单位：德国柏林
    - 邮箱：[[email protected]](mailto:)
- **Christian Wald**
    - 单位：卡米尔·约旦研究所，法国里昂国立应用科学学院
    - 单位：法国维勒班
    - 邮箱：[[email protected]](mailto:)

###### 摘要
图结构在节点置换下保持不变，这促使生成模型中采用置换等变架构。然而在流匹配中，对称性也可能体现在源-目标耦合方式中：一旦图对在节点重标定意义下进行比较，其自然的 Wasserstein 几何结构便对应于图商空间。该空间的欧几里得商度量恰好等同于通过最优节点重标定得到的 Gromov-Monge 距离。我们从理论角度发展了这一观点，证明商耦合可以在无额外成本下提升至对齐的代表元，且对称化处理能够产生等变流匹配极小化子，包括对范畴化端点预测的情形。实践中，精确的 Gromov-Monge 对齐难以计算，因此我们采用高效的 Gromov-Wasserstein 型松弛与内节点对齐的下界来构建小批量耦合，并可选地结合图间外部指派。所得流程仅改变训练耦合方式，可与标准置换等变架构兼容。在连续图生成与范畴化分子生成任务中，这些结构感知耦合在少量积分步数下显著提升了样本质量，而我们的扩展分子模型在常规多步采样中仍保持竞争力。

## 1 引言
流匹配通过无仿真回归与常微分方程（ODE）采样，学习一个将易采样源分布输运至数据分布的速度场（Albergo 和 Vanden-Eijnden, 2023；Lipman 等, 2023；Liu 等, 2023）。其条件路径由源-目标耦合决定。尽管所有此类耦合在起始和终止时间具有相同的边际分布，但基于输运信息的选择能产生更短、更直的轨迹，从而简化数值积分（Chemseddine 等, 2025；Tong 等, 2023；Pooladian 等, 2023）。当数据仅在对称性意义下定义时，正确的输运概念就不那么直观了（Klein 等, 2023；Köhler 等, 2020）。对于图而言，对称性体现为节点重标定。我们将一个 N 节点图的边通道和对角节点特征编码为 E∈ℝ^{N×N×C}。由于图没有预定义的顶点顺序，置换 σ∈𝔖_N 同时作用于两个节点索引：(ρ_N^{graph}(σ)E)_{ijc} ≔ E_{σ^{-1}(i)σ^{-1}(j)c}，其中 1≤i,j≤N。无序图即轨道 [E]∈Q≔ℝ^{N²C}/G_N^{graph}，其中 G_N^{graph}=ρ_N^{graph}(𝔖_N)⊂O(N²C)。图生成模型中采用的等变架构尊重此对称性（Eijkelboom 等, 2024；Jo 等, 2022；Vignac 等, 2023），但许多模型本身并不选择在流匹配过程中连接哪个带标签的代表元。从 E 到无序目标 [F] 的每个条件桥可能选择某个 σ·F。最小化其长度是一个 **Gromov-Monge** 节点对齐问题，即 Gromov-Wasserstein 输运的硬指派对应问题（Bauer 等, 2025；Mémoli, 2011；Mémoli 和 Needham, 2024）。我们将此代表元选择与一般商空间 ℝ^D/G 上的等变流匹配联系起来（Klein 等, 2023；Köhler 等, 2020；Song 等, 2023）。我们的理论表明，最优商耦合可以在无额外成本下提升至欧几里得耦合，且选择代表元之间的耦合经对称化处理后，会产生流匹配目标的等变极小化子，包括对范畴化端点预测的情形。对于图，我们用 Gromov-Wasserstein 求解器和廉价下界近似得到的 Gromov-Monge 对齐（Mémoli, 2011；Bauer 等, 2025；Piening 和 Beinert, 2025），并可选地接续外部小批量指派。对齐后的图对训练与标准流匹配相同的等变速度场或范畴化端点模型；仅耦合方式发生变化。

##### 与现有工作的关系
图生成器通常通过等变架构来强制节点重标定对称性，包括分子与离散扩散（Hoogeboom 等, 2022；Jo 等, 2022；Vignac 等, 2023）以及变分流匹配（Eijkelboom 等, 2024）。近期的图流方法还额外使用小批量最优传输优化源-目标图配对（Hou 等, 2026；Wijesinghe 等, 2026）。我们的方法通过在每个图轨道内显式选择对齐代表元来补充这些方法，同时将外部图指派视为可选项，详见附录 C（https://arxiv.org/html/2608.26961#A3）的更详细比较。最优传输耦合此前已被用于拉直欧几里得流匹配路径（Chemseddine 等, 2025；Pooladian 等, 2023；Tong 等, 2023）。当底层样本空间本身具有最优传输几何时，这自然引出嵌套的输运表述，正如近期用于点云生成模型（Haviv 等, 2025；Piening 等, 2026；Piening 和 Beinert, 2026）。对于图，类似构造还需额外考虑节点对应关系，从而产生 Gromov-Monge 基础成本。Gromov-Wasserstein 方法提供了一种自然的松弛，因为它们比较关系结构，并长期用于图与结构化数据对齐（Beier 等, 2025；Chowdhury 等, 2021；Peyré 等, 2016；Vayer 等, 2020）。我们利用所得传输方案在条件流匹配桥内为每个目标图选择一个节点置换，而非直接生成图。我们的商空间分析建立在等变生成流的基础上（Klein 等, 2023；Köhler 等, 2020；Song 等, 2023）。超越通过线性指派实现的粒子置换（Klein 等, 2023；Haviv 等, 2025；Hui 等, 2025；Piening 等, 2026），图重标定同时作用于节点索引，从而产生二次型 Gromov-Monge 问题。我们将此图特定对齐与商最优传输和代表元提升联系起来（Mémoli 和 Needham, 2024）。我们的贡献包括：
- • 我们证明最优商耦合可以提升为具有相同二次成本的欧几里得耦合，其线性插值投影到恒速 Wasserstein 测地线上。
- • 我们证明对角对称化保持等变场上的流匹配目标，并产生等变极小化子，包括对范畴化端点预测的情形。
- • 一种实用的外部-内部 Gromov-Wasserstein 对齐方法为连续与范畴化图流构建小批量耦合。
- • 在连续与范畴化图生成基准测试中，当学习到的流使用少量欧拉积分步数进行采样时，结构感知耦合能提升样本质量，而扩展的分子模型在常规多步采样中仍保持竞争力。

## 2 输运与流匹配背景
在流匹配生成模型中，构造连接易采样源分布与仅通过样本可得的目标分布的概率测度曲线（Lipman 等, 2023；Wald 和 Steidl, 2025）。此构造不限于欧几里得空间上的概率测度（Chen 和 Lipman, 2024）。在简要回顾 Wasserstein 几何（Ambrosio 等, 2005；Santambrogio, 2015；Wald 和 Steidl, 2025）后，我们将看到如何在正交群作用诱导的商空间上表述相同构造。

### 2.1 Wasserstein 几何与动态输运
设 (X, d_X) 为波兰度量空间，𝒫_2(X) 为 X 上具有有限二阶矩的 Borel 概率测度集，即对某 x₀∈X 满足 ∫_X d_X(x, x₀)² dμ(x) < ∞。集合 𝒫_2(X) 在 *Wasserstein 距离* W_{2,X}（Santambrogio, 2015；Villani, 2009）下成为完备度量空间，其对于任意 μ, ν∈𝒫_2(X) 定义为：

W_{2,X}²(μ, ν) ≔ min_{π∈c_X(μ, ν)} ∫_{X×X} d_X(x, x')² dπ(x, x').  (1)

*耦合*或*传输计划*由 c_X(μ, ν) = {π∈𝒫_2(X×X) : proj⁰_♯π = μ, proj¹_♯π = ν} 给出，其中 projⁱ: X×X → X, (x₀, x₁) ↦ xᵢ。我们用 c_X^{opt}(μ, ν) 表示*最优*耦合，即式 (1) 取得最小值。始终令 I=[0,1]。对 X=ℝ^D，令 (μ_t)_{t∈I} 狭义连续，v: I×ℝ^D → ℝ^D 为 Borel 场。对 (μ_t, v_t) 若满足连续性方程 ∂_t μ_t + div(v_t μ_t) = 0 (2)（分布意义下）。若进一步 ∫_I ∫_{ℝ^D} ‖v_t(x)‖² dμ_t(x) dt < ∞，则 (μ_t)_{t∈I} 关于 W_{2,ℝ^D} 绝对连续。反之，关于 W_{2,ℝ^D} 绝对连续且具有平方可积速度的曲线，承认满足式 (2) 及所述可积条件的 Borel 场 v（Ambrosio 等, 2005, 第 8 章）。以下定义和界在附录 A（https://arxiv.org/html/2608.26961#A1）中回顾。若 v 足够正则，特征 ODE γ̇(t,x) = v(t, γ(t,x)), γ(0,x)=x, (3) 将 μ₀ 输运至 μ_t，即 μ_t = γ_{t,♯}μ₀（Ambrosio 等, 2005, 命题 8.1.8）。

### 2.2 欧几里得流匹配
流匹配将上述动态描述转化为均方误差（MSE）速度场回归（Lipman 等, 2023；Wald 和 Steidl, 2025，参见例如）。对 μ, ν∈𝒫_2(ℝ^D) 和 π∈c_{ℝ^D}(μ, ν)，定义 μ_t ≔ proj^t_♯π，其中 proj^t(x, x') ≔ (1-t)x + tx'。则 (μ_t, v_t^π) 满足式 (2)，其中 v^π 最小化 J_π(v) ≔ ∫_I ∫_{ℝ^D×ℝ^D} ‖v_t(proj^t(x, x')) - (x' - x)‖² dπ(x, x') dt (4)（联合 Borel 场 v）。等价地，对 (X, X') ∼ π，有 v_t^π(z) = E_π[X' - X | proj^t(X, X') = z] (5)（对几乎处处的 (t, z)）。若 π∈c_{ℝ^D}^{opt}(μ, ν)，则动能满足 ∫_I ‖v_t^π(z)‖_{L^2(μ_t)}^2 dt = W_{2,ℝ^D}²(μ, ν)。

##### 范畴化端点预测
设 μ, ν∈𝒫_2(ℝ^D) 和 π∈c_{ℝ^D}(μ, ν)。假设目标是坐标范畴化的。记 [m]≔{1,...,m}（正整数 m），并固定正整数 M。对每个 d∈[D]，固定互异值 a_{d,1},...,a_{d,M}∈ℝ 使得 ν(∏_{d=1}^D {a_{d,1},...,a_{d,M}}) = 1。对 (X, X') ∼ π，...