Gromov-Monge Flow Matching for Equivariant Graph Generation
Moritz Piening, Christian Wald
cs.LG, math.OC, stat.ML
2026-08-27
训练阶段用Gromov-Wasserstein对齐源图与目标图的节点，再做流匹配，等变网络保持不变。连续SBM五步FGW-NNA从0.796降到0.568，ZINC五步分子有效性从56%升到73%。
图生成模型早就接受一件事实：换一次节点编号，还是同一张图。网络因此做成 permutation-equivariant，输入怎么换标签，输出跟着换。DiGress、GDSS、CatFlow 这条线都靠这件事成立。
Flow matching 还有另一层选择，等变网络管不到。训练要把一张噪声图接到一张数据图，中间走直线插值。噪声图的 1 号节点对数据图的哪一个?随机配对时，两边的边对不上号，中途会出现半有半无的边。论文用 12 节点圆环把这件事画清楚：源和目标都是半径 1 的正十二边形，只差一个竖直平移，边的规律完全一样。随机对齐时 t=1/2 的图塌成一团密网;把节点对齐之后，圆环全程保持完整。外层再选「哪两张图配对」，路径才会变直。内层管对象长什么样，外层管路怎么走。
这件事的几何名字是图商空间上的 Wasserstein 距离，和 Gromov-Monge 距离是一回事：找一个节点置换，让两边全部成对关系尽量对上。精确求解是 NP-hard 的二次分配。
理论做了两步。定理 3.1:商空间上的耦合可以抬到欧氏代表，二次代价不增加;最优耦合抬上去之后，直线插值投影回商空间是常速测地线。定理 3.2:把耦合沿群对角平均，在等变网络类上目标函数不变，并且存在等变的最优速度场。分类端点预测同样成立，不要求源、目标本身已经是群不变的。
工程上不求精确 Gromov-Monge，走两层近似，边、节点代价权重都取 1/2:
图编码成 N×N×C 张量，边特征铺在全体条目，节点特征只写在对角。变节点数数据按 N 组 batch，推理时从训练集经验分布抽 N。网络是标准等变 graph transformer（XEy，约 2.8M 参数）。连续任务回归速度，分子任务做坐标级分类端点预测再还原速度。改的只有训练耦合。
对照里还有一条无视置换的 MinibatchOT:目标图随机重标，再按欧氏距离在整个 minibatch 上配对。它只做外层、不做内层。
连续 SBM(N=10，无条件混合社区数 K=1 到 5)，5 步 Euler，3 个训练种子：
| 耦合 | FGW-NNA(越近 0.5 越好) | Degree MMD | Clustering MMD | Graphlet MMD |
| Random | 0.796 | 0.114 | 0.179 | 0.062 |
| MinibatchOT | 0.687 | 0.031 | 0.090 | 0.006 |
| FLB | 0.725 | 0.054 | 0.109 | 0.030 |
| GW | 0.577 | 0.027 | 0.049 | 0.014 |
| GW+out | 0.568 | 0.018 | 0.034 | 0.005 |
125 步时差距收窄，符合「步数够了都能走到」的预期。内层 GW 贡献最大，外层再加一点，FLB 吃到部分收益。MinibatchOT 多数指标赶上 FLB，明显落后 GW。条件版(把 K 喂给网络)排序一样：五步 FGW-NNA 从 Random 的 0.746 降到 GW+out 的 0.561。
分子短预算，同一套 2.8M 网络训 100 个 epoch，1 万样本：
| 设置 | QM9 五步有效性 / FCD | ZINC 五步有效性 / FCD |
| Random | 0.8849 / 1.731 | 0.5641 / 18.448 |
| MinibatchOT | 0.9357 / 1.747 | 0.6106 / 17.190 |
| GW | 0.9356 / 1.278 | 0.6431 / 15.024 |
| GW+out | 0.9421 / 1.232 | 0.7280 / 15.377 |
QM9 上 MinibatchOT 的有效性已经贴着内层 GW，FCD 却几乎没动，甚至略差于 Random。ZINC 上外层把有效性再抬一截，FCD 反而略差于纯内层 GW。步数加到 125，差距继续收窄。
放大后的 GW-CatFlow 关掉外层、加上 RRWP、self-conditioning 和两阶段 dropout，500 步：QM9 有效性 99.34%、FCD 0.115，和 DeFoG 的 99.30% / 0.120 同一档;ZINC 有效性 99.01%、FCD 0.966，表里最低(VBFN 是 1.307)。这张表不能单归因于对齐。
能直接拿去用的部分很窄，也很清楚。已经在跑等变 graph flow matching 的话，训练时加一层节点对齐，采样器不用改。收益集中在少步数：想 5 步出图，这件事值得做;已经 500 步，边际小得多。
和 GGFlow、Flowette 的差别在分工。那些工作主要用 OT 决定哪两张图配对。这篇把内层节点重标单独拆出来，消融显示内层才是大头。
训练代价不免费。附录里单线程 CPU 上，GW 内层占一次梯度步的 26%（SBM，B=16）到 55%–58%（QM9 / ZINC）;再加外层到 63%–88%。FLB 只要 1%–3%，但 tightness 只有 GW 代价的 0.1%–3.1%，Pearson 相关 0.15–0.23。当便宜启发用可以，别把它当成近似求解器。推理阶段对齐完全不出现。
作者写明了：精确 Gromov-Monge 不可解，GW 是非凸松弛，投出来的置换不必是最优硬对齐。Frank-Wolfe 只跑 10 步，外层配对还减到 5 步。对齐只在训练出现，墙钟会被 GW 拖住。
几处没钉死。全预算分子表混了更大模型(QM9 约 6.0M，ZINC 约 10.8M)和若干训练技巧，不能隔离对齐效应。SBM 只有 10 个节点。外层只在同 N、最多 8 张的子 batch 里做，对 ZINC 那种 N≤38、节点数分布不均的数据，外层能覆盖的配对空间很小。FLB 几乎不跟踪 GW 数值，却仍有中等收益，说明一个更朴素的结构启发式(按度排序之类)可能就够，论文没有把这类基线单独拆开。