Gromov-Monge Flow Matching for Equivariant Graph Generation

Moritz Piening; Christian Wald

Source-linked AI summary

Gromov-Monge Flow Matching for Equivariant Graph Generation

Moritz Piening, Christian Wald

arXiv:2608.26961v1Aug 27, 2026cs.LGmath.OCstat.ML

TL;DR

Graph flow matching must account for node-permutation symmetry in its source–target coupling, since the relevant geometry is the graph quotient space. The paper develops quotient transport theory and practical Gromov–Wasserstein-based alignments, finding better few-step sample quality while retaining competitive many-step molecular performance.

Problem
Graph generative models enforce permutation equivariance architecturally, but often do not choose aligned labelled representatives for source–target flow-matching couplings.
Method
The paper lifts quotient couplings to aligned representatives, symmetrizes them for equivariant minimization, and approximates graph alignments with inner Gromov–Wasserstein relaxations and optional outer assignments.
Results
Structure-aware couplings improve sample quality for few-step continuous and categorical graph generation, while scaled-up molecular models remain competitive with conventional many-step sampling.
Takeaways & Limitations
Only the training coupling changes, so the approach remains compatible with standard permutation-equivariant velocity and categorical endpoint models.
Takeaways & Limitations
The Gromov–Wasserstein representative is a relaxation and need not solve the original hard Gromov–Monge alignment problem.

Abstract

from arXiv · show

Graphs are invariant under node permutations, motivating the use of permutation-equivariant architectures in generative models. In flow matching, however, symmetry may also enter the source--target coupling: once graph pairs are compared up to node relabeling, the natural Wasserstein geometry is that of the graph quotient space. The Euclidean quotient metric of this space coincides with the Gromov--Monge distance, obtained by optimally relabeling the nodes. We develop this perspective theoretically, showing that quotient couplings can be lifted to aligned representatives without additional cost and that symmetrization yields equivariant flow-matching minimizers, including for categorical endpoint prediction. In practice, exact Gromov--Monge alignment is intractable, so we construct minibatch couplings using efficient Gromov--Wasserstein-type relaxations and lower bounds for the inner node alignment, optionally combined with an outer assignment between graphs. The resulting procedure changes only the training coupling and is compatible with standard permutation-equivariant architectures. Across continuous graph and categorical molecular generation, these structure-aware couplings substantially improve sample quality at small integration budgets, while our scaled-up molecular models remain competitive under conventional many-step sampling.

1 Introduction

The paper frames graph flow matching on permutation quotient spaces, where choosing aligned representatives induces Gromov–Monge transport. It develops theory and practical couplings that improve few-step graph generation while preserving standard equivariant models.

Graph node relabeling makes the source–target coupling part of the symmetry problem, not merely an architectural constraint.
Gromov–Monge alignment minimizes bridge length by optimally relabeling target nodes, providing the hard-assignment counterpart of Gromov–Wasserstein transport.
Optimal quotient couplings lift to Euclidean couplings with the same quadratic cost, and their interpolations project to constant-speed Wasserstein geodesics.
Diagonal symmetrization preserves the flow-matching objective on equivariant fields and yields equivariant minimizers, including categorical endpoint prediction.
The practical method uses inner Gromov–Wasserstein-type alignment and optional outer graph assignment to construct minibatch couplings without changing the model architecture.
Across continuous and categorical graph-generation benchmarks, structure-aware couplings improve few-step sample quality, while scaled-up molecular models remain competitive with many-step sampling.

2 Background on transport and flow matching

This section introduces Wasserstein transport and flow matching as a dynamic construction connecting a tractable source distribution to a sampled target. It also describes the vector-field and categorical endpoint formulations used later.

Flow matching connects an easy-to-sample source measure with a target measure available through samples by constructing a curve of probability measures.
Wasserstein distance equips probability measures with finite second moments with a metric based on optimal coupling costs.
Continuity equations and characteristic ODEs describe how a velocity field transports the source law along the evolving measure curve.
Flow matching turns this dynamic transport description into mean-squared-error regression for a conditional vector field.
For categorical targets, coordinatewise class distributions suffice to train endpoint predictors and recover the velocity without a conditional-independence assumption.

3 Flow matching on quotient spaces

On quotient spaces induced by compact isometric group actions, the paper shows that representative lifts preserve transport cost and that symmetrization makes flow matching compatible with equivariant fields and quotient flows.

The quotient metric identifies points up to group actions, and its induced Wasserstein distance governs transport between unordered representations.
Choosing graph representatives is precisely a hard Gromov–Monge node-alignment problem under simultaneous relabeling of both node indices.
Every quotient coupling can be lifted through aligned representatives without increasing cost; optimal quotient couplings lift to optimal representative couplings.
The lifted interpolation projects to a constant-speed quotient Wasserstein geodesic, and its flow-matching field satisfies the corresponding quotient transport relation.
Although representative selection can break symmetry, diagonal averaging restores invariance while preserving the quotient coupling.
For equivariant fields, symmetrization preserves the objective and yields an equivariant minimizer; well-posed equivariant ODEs induce flows depending only on quotient laws.

4 Constructing graph couplings via Gromov–Monge approximation

The method approximates hard Gromov–Monge node alignment with Gromov–Wasserstein relaxations or an inexpensive eccentricity-based lower bound, optionally combining these inner alignments with outer graph assignment.

Inner alignment: Exact quotient-cost evaluation is NP-hard because one permutation must consistently align all pairwise graph relations.The method therefore uses approximate inner alignment rather than solving the quadratic assignment exactly.
Inner alignment: Gromov–Wasserstein relaxation optimizes a soft transport plan, then recovers a node permutation with Frank–Wolfe iterations and the Hungarian algorithm.The recovered permutation relabels the target graph before flow-matching training.
Inner alignment: The first lower bound sorts nodes by eccentricity and pairs nodes at matching ranks, providing an inexpensive alignment for candidate graph pairs.Computing eccentricities costs O(N^2C), while sorting costs O(N log N).
Inner alignment: GW preserves detailed pairwise structure at higher cost, whereas FLB is cheap enough to evaluate across B^2 candidate pairs for outer minibatch assignment.This creates a computational trade-off between alignment detail and minibatch scalability.
Outer–inner coupling: Inner alignment selects representatives within paired graphs, while optional outer assignment selects which source and target graphs are connected.The two decisions are distinct, and inner alignment can shorten interpolation even under independent outer pairing.
Outer–inner coupling: After pairing and alignment, the method interpolates between each source graph and its aligned target and trains continuous or categorical flow-matching losses.A shared permutation-equivariant graph transformer handles batches with a fixed node count, while node count may vary across batches.

5 Numerical experiments

The experiments compare graph couplings that differ in node alignment and minibatch assignment while keeping the equivariant backbone fixed. Structure-aware couplings improve graph and molecular quality most at five Euler steps, with strong molecular results also at 500 steps.

Experimental setup: The experiments compare RANDOM, FLB, and 10-iteration GW inner alignments, optionally adding outer assignment, with a shared equivariant backbone.MINIBATCHOT provides a permutation-blind reference, and the models differ only in their training coupling.
Illustrative translated cycles: Inner GW preserves the ring through interpolation, while outer assignment straightens its trajectory.Random inner relabelling creates denser intermediate graphs with half-present edges; GW removes this edge ambiguity.
Continuous SBM: At five steps, GW with outer assignment reduces FGW–NNA from 0.796 to 0.568, degree MMD from 0.114 to 0.018, and clustering MMD from 0.179 to 0.034.The comparison is for the continuous SBM experiment; the gap narrows by 125 steps.
Categorical molecular generation: At five steps, inner GW raises QM9 validity from 0.8849 to 0.9356 and ZINC250k validity from 0.5641 to 0.6431.It also reduces FCD from 1.731 to 1.278 on QM9 and from 18.448 to 15.024 on ZINC250k.
Categorical molecular generation: Outer assignment further improves validity, while inner-only GW has marginally better ZINC250k FCD and performance differences narrow with more steps.Outer assignment is uniformly best on QM9 FCD and gives the highest ZINC250k validity.
Full-budget molecular generation: At 500 Euler steps, GW-CATFLOW reaches 99.34% validity and 0.115 FCD on QM9, and 99.01% validity with 0.966 FCD on ZINC250k.The authors describe these results as competitive rather than isolating alignment because baselines use published protocols and the full-budget model includes architectural changes.

6 Conclusion

The method uses Gromov–Wasserstein alignments to specialize equivariant flow matching for graphs. Its inner and optional outer alignments improve few-step generation, while the approximate GW solvers remain computationally expensive.

Method: Inner alignment reorders each target graph’s nodes, while optional outer assignment reorders the minibatch to choose source–target graph pairings.These are distinct coupling decisions.
Conclusion: The resulting couplings improve few-step generation over random assignment and permutation-blind minibatch OT.The procedure changes the training coupling and is absent at inference.
Limitation: GW solvers approximate the Gromov–Monge problem and are expensive, although their cost is incurred during training rather than inference.The paper points to accelerated approximations and semi-discrete formulations as future directions.

A Proof of representative lifts and quotient geodesics

The proof establishes representative lifting for quotient couplings and the resulting geodesic structure. It shows that optimal quotient transport can be realized by aligned representatives without increasing quadratic cost.

Quotient metric: The quotient Q is formed by identifying points in R^D under the group action, with distance defined by minimizing Euclidean distance over group transformations.Compact isometric group actions ensure the minimum is attained, and q is 1-Lipschitz.
Representative lifting: Any quotient coupling can be lifted to a representative-space coupling whose pushforward is the original quotient coupling and whose cost attains the corresponding minimum.The construction fixes a representative of one marginal and measurably selects aligned representatives for the other.
Optimality: An optimal quotient coupling admits a representative lift with quadratic cost equal to the quotient Wasserstein distance.The proof constructs the lift through an aligned representative map and verifies the marginal and cost identities.
Geodesic conclusion: The lifted interpolation is a constant-speed geodesic between the endpoint measures.Equality in the interval bounds and the triangle inequality forces the geodesic property.

B Proof of equivariant flow matching via symmetrization

Diagonal symmetrization converts arbitrary couplings into couplings compatible with equivariant flow matching while preserving the relevant objective. The resulting equivariant fields, flows, and categorical endpoint predictors remain well defined on graph quotients.

Equivariant minimizers: For a diagonally invariant coupling, the averaged velocity field changes only the representative of the L2 minimizer and is exactly equivariant.Haar averaging preserves the minimizer almost everywhere, then right invariance establishes equivariance everywhere under the stated convention.
Symmetrization: Diagonal symmetrization preserves the flow-matching objective over equivariant fields and produces an equivariant minimizer.The constrained problem is equivalent to ordinary flow matching for the symmetrized coupling.
Symmetrization: Averaging a representative coupling over the group yields a diagonally invariant coupling with the same quotient coupling.The symmetrized coupling represents the same quotient pushforward while coupling invariantized marginals.
Quotient flows: An equivariant velocity field with a unique ODE flow induces a well-defined quotient flow independent of the initial representative lift.The flow commutes with the group action, so its projection factors through the quotient.
Categorical prediction: Coordinatewise categorical endpoint probabilities admit equivariant representatives that minimize the constrained coordinatewise objective.This construction does not require the endpoint measures themselves to be invariant.
Categorical prediction: Finite categorical data are implemented with one-hot feature blocks and softmax heads rather than scalar class positions.Scalar encodings introduce arbitrary geometry between classes, motivating blockwise categorical treatment.

D Conditional continuous-SBM experiment

The conditional continuous-SBM experiment supplies the community count K to the model and evaluates matching K-specific graph sets. Its ranking reproduces the unconditional experiment, with GW alignment strongest at five steps and diminishing gains as the integration budget grows.

Conditional evaluation: GW alignment has the largest advantage at five Euler steps, and the performance gap narrows with the integration budget.The conditional experiment pools generated and real graphs with matching ground-truth K multisets.

E Full-budget molecular training protocol

The full-budget molecular protocol retains categorical endpoint prediction with GW inner alignment and independent outer pairing while scaling model capacity and training toward a published baseline. Its results establish competitiveness, not an isolated alignment effect.

Protocol: The full-budget protocol uses categorical endpoint prediction, GW inner alignment, independent outer pairing, larger models, and extended training budgets.It also adds structural graph statistics, relative random-walk probabilities, and self-conditioning.

F Reproducibility details

The reproducibility specifications define controlled coupling comparisons in which model and optimization settings are shared, while datasets, solvers, metrics, and evaluation protocols are explicitly fixed. The experiments cover continuous graph, cycle-graph, and molecular settings.

Shared settings: Controlled coupling experiments keep architecture, optimizer, EMA, gradient clipping, and learning-rate schedule identical, varying only the specified coupling conditions.The larger-model protocol is specified separately.
Datasets: Cycle-graph data use randomly relabelled 12-node rings with translated and independently rotated source and target configurations.The representation contains one adjacency channel and planar node positions on the diagonal.
Datasets: The continuous SBM uses 10-node graphs with K communities, Beta-distributed edges, and pooled or class-conditional K settings.The unconditional dataset pools K ∈ {1, . . . , 5}, while the conditional model receives K as a class label.
Datasets: Molecular graphs use one-hot bond channels and diagonal atom-type and formal-charge features for QM9 and ZINC250k.The categorical representation supports endpoint prediction over the listed molecular feature vocabularies.
Coupling solvers: The inner GW solver uses 10 Frank–Wolfe iterations with Hungarian linearized steps, while eccentricity sorting provides an O(N^2C + N log N) lower-bound alignment.Outer reordering applies Hungarian assignment within same-size sub-batches of at most eight graphs.
Metrics: Molecular validity counts sanitizable RDKit decodes without valence correction, while FCD compares ChemNet embeddings after fragment retention, canonicalization, and duplicate removal.Both metrics are defined on generated molecules processed through the stated evaluation pipeline.

G Solver cost and bound tightness

The solver comparison measures FLB and GW costs on same-size graph pairs and at minibatch scale. FLB is substantially cheaper but numerically loose, while GW alignment can materially contribute to gradient-step cost.

Per-pair comparison: Table 7 evaluates FLB against computed GW over 300 same-N real graph pairs per dataset, reporting runtime, tightness, and correlation.Runtime is measured in milliseconds per pair on a single CPU thread; tightness is mean(FLB)/mean(GW).
Per-pair comparison: FLB is cheaper than GW by roughly 21× at N = 10, 34× at N = 9, and 44× at N = 23, with the gap widening as node count increases.The comparison follows O(N^2C + N log N) scaling for FLB versus O(N^3 × iters) for GW.
Bound tightness: FLB is numerically loose and only weakly correlated with GW, so its empirical role is a cheap alignment heuristic rather than an accurate GW-cost approximation or screening rule.This qualifies how lower-bound efficiency should be interpreted in the solver pipeline.
Batch-level cost: At batch level, GW rows dominate alignment cost because every aligned pair requires its own Frank–Wolfe solve, although pair separability enables a 12-worker pool to reduce inner QM9 alignment to about 28 ms.The reported gradient-step shares are upper bounds because the alignments can be parallelized.