Gromov-Monge Flow Matching for Equivariant Graph Generation
Moritz Piening, Christian Wald
cs.LG, math.OC, stat.ML
2026-08-27
Train-time Gromov-Wasserstein node alignment, same equivariant net: SBM 5-step FGW-NNA falls from 0.796 to 0.568, ZINC 5-step validity rises from 56% to 73%.
Graph generators already treat node order as arbitrary. Permute the labels and the graph is unchanged, so the network is built permutation-equivariant: relabel the input and the output relabels with it. That is the standard story behind DiGress, GDSS, and CatFlow.
Flow matching still has to pick a labelled representative for each pair. Training draws a noise graph and a data graph and interpolates them in a straight line. Which node on the noise graph should meet which node on the data graph? A random matching misaligns the edges, so the midpoint fills with half-present bonds. A 12-node ring makes the failure visible. Source and target are the same regular 12-gon of radius 1, offset only by a vertical translation, so the edge law is identical on both sides. With random inner matching the ring collapses into a dense mesh at t=1/2. Align the nodes and the ring stays intact all the way; an extra outer assignment then straightens the paths. Inner alignment sets the shape of the object being moved. Outer assignment sets the geometry of the path.
The matching cost is the Wasserstein metric on the graph quotient, which coincides with the Gromov-Monge distance: search over node permutations for the one that best matches all pairwise relations. Exact search is NP-hard quadratic assignment.
Two theorems justify the recipe. Theorem 3.1: any quotient coupling lifts to Euclidean representatives at the same quadratic cost; if the coupling is optimal, the linear interpolant projects to a constant-speed geodesic on the quotient. Theorem 3.2: averaging a coupling along the diagonal group action leaves the flow-matching objective unchanged on equivariant fields and yields an equivariant minimizer. The same holds for coordinate-wise categorical endpoint prediction, even when the source and target laws are not themselves group-invariant.
Exact Gromov-Monge is skipped. Two nested approximations are used, with edge and node weights both set to 1/2:
Graphs are N×N×C tensors: edge features on every entry, node features only on the diagonal. Variable-size datasets are batched by N; at sampling time N is drawn from the training empirical distribution. The backbone is a standard equivariant graph transformer (XEy, about 2.8M parameters). Continuous tasks regress velocity; molecular tasks predict a coordinate-wise categorical endpoint and recover velocity from it. Only the training coupling changes.
A permutation-blind reference, MinibatchOT, randomly relabels each target and then pairs graphs by Euclidean distance over the full minibatch. It does outer matching and skips inner alignment.
Unconditional SBM on N=10 nodes, community count K mixed from 1 to 5, 5 Euler steps, three training seeds:
| Coupling | FGW-NNA (closer to 0.5 better) | 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 |
By 125 steps the gap shrinks, as expected once every interpolant has enough time. Inner GW does most of the work; the outer assignment adds a smaller bump; FLB captures part of the gain. MinibatchOT matches or beats FLB on most descriptors and stays clearly behind GW. The class-conditional variant, which feeds K to the network, ranks the same way: five-step FGW-NNA falls from 0.746 (Random) to 0.561 (GW+out).
Molecular short-budget sweep, shared 2.8M network, 100 epochs, 10,000 samples:
| Setting | QM9 5-step validity / FCD | ZINC 5-step validity / 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 |
On QM9, MinibatchOT already matches inner GW on validity, yet FCD barely moves and is slightly worse than Random. On ZINC the outer assignment lifts validity further, while inner-only GW keeps a slightly better FCD. Extra Euler steps again compress the differences.
A scaled GW-CatFlow run drops the outer assignment and adds RRWP, self-conditioning, and two-stage dropout. At 500 steps: QM9 validity 99.34% and FCD 0.115, in the same band as DeFoG (99.30% / 0.120); ZINC validity 99.01% and FCD 0.966, the lowest FCD in the table (VBFN sits at 1.307). That table does not isolate the coupling.
The usable piece is narrow and concrete. If an equivariant graph flow matcher is already in the stack, add a node-alignment step at train time and leave the sampler alone. The gain lives in the few-step regime. Five Euler steps is where this coupling pays; five hundred steps leave much less on the table.
GGFlow and Flowette mostly use transport to decide which graphs to pair. This paper peels off the inner relabelling as its own operation, and the ablations put most of the gain there.
Training is not free. On single-thread CPU, inner GW is 26% of a gradient step on SBM (B=16) and 55%–58% on QM9 / ZINC; adding the outer assignment pushes that to 63%–88%. FLB costs 1%–3% of a step, but its tightness is only 0.1%–3.1% of the GW cost, with Pearson correlation 0.15–0.23. Treat it as a cheap heuristic, not as an approximate solver. Alignment is absent at inference.
The authors are clear: exact Gromov-Monge is intractable, GW is a non-convex relaxation, and the projected permutation need not be the best hard alignment. Frank-Wolfe is capped at 10 iterations, 5 for the outer cost. Alignment never appears at sampling, but it does stretch wall-clock training.
Several claims stay under-determined. The long-budget molecular table mixes a larger backbone (about 6.0M on QM9, 10.8M on ZINC) with extra training tricks, so it cannot isolate alignment. The SBM uses only 10 nodes. Outer matching runs inside same-N sub-batches of at most eight graphs; on ZINC, where N is at most 38 and the size distribution is uneven, that pairing pool is small. FLB barely tracks GW numerically and still helps, which suggests a still cheaper structural heuristic (degree sorting, for instance) might be enough. That baseline is not isolated.