Quotient-Aligned Flow Matching
Implementation & benchmark of arXiv:2608.26961 — Gromov-Monge Flow Matching for Equivariant Graph Generation
Source paper: Gromov-Monge Flow Matching for Equivariant Graph Generation arXiv:2608.26961 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace the arbitrary labelled target graph in each flow-matching pair by the node permutation that minimizes its Euclidean distance to the source graph. Construct the conditional interpolation and velocity target using this aligned representative, so the model learns transport in graph quotient space rather than wasting capacity on node-order mismatches.
Formulas
Mathematical statement
A graph with edge-channel tensor E in R^{N x N x C} is acted on by a permutation sigma in the symmetric group S_N as (rho_N^graph(sigma)E)_{ijc} = E_{sigma^{-1}(i),sigma^{-1}(j),c}; the unordered graph is the orbit [E] = {rho_N^graph(sigma)E : sigma in S_N}. For source tensor X and target tensor Y, use the Gromov-Monge alignment sigma* = argmin_{sigma in S_N} ||X - rho_N^graph(sigma)Y||_F^2. The quotient distance is d_Q([X],[Y]) = min_{sigma in S_N} ||X - rho_N^graph(sigma)Y||_F. For linear flow matching at time t in [0,1], define the aligned target Y* = rho_N^graph(sigma*)Y, interpolation Z_t = (1-t)X + tY*, and conditional velocity u_t = Y* - X. The training objective is E_{X,Y,t}[||v_theta(Z_t,t) - u_t||_F^2]. Because applying a common permutation to X and Y only changes sigma* by conjugation and leaves the quotient distance unchanged, an equivariant v_theta can represent the resulting quotient-space flow.
Implementation notes
Integrate this at the pair-construction stage of a continuous graph flow-matching trainer; do not change the equivariant graph neural network or the ODE solver initially. Inputs are a source tensor X, a target tensor Y, and channel weights w_c. First compute a node-cost matrix using node degree, node features, and local edge-neighborhood statistics. Obtain an approximate permutation sigma_star using a graph-matching solver, Sinkhorn relaxation followed by Hungarian rounding, or a few FAQ/Frank-Wolfe iterations. Apply the permutation to both target indices: Ystar[i,j,c] = Y[sigma_star^{-1}(i), sigma_star^{-1}(j), c]. Sample t uniformly, form Z = (1-t)*X + t*Ystar, and train the existing model on target velocity U = Ystar-X. Pseudocode: sigma = match(X,Y); Ystar = permute_edges(Y,sigma); t=rand(); Z=(1-t)*X+t*Ystar; loss=weighted_mse(v_theta(Z,t),Ystar-X). Estimate the exact quotient distance only for diagnostics; the training alignment is approximate. Log pre- and post-alignment Frobenius distances and matching stability across repeated solver runs. The first experiment should use a small molecular graph dataset, an equivariant flow model, and Euler sampling with 4, 8, and 16 steps. Compare random or canonical ordering against aligned coupling at equal training FLOPs. Success is lower conditional-target variance, faster validation loss decrease, and higher validity or FID at the same integration budget.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built Quotient-Aligned Flow Matching with exact enumeration over all 5-node permutations, aligned interpolation targets, a small flow MLP, and Euler evaluation at 4/8/16 steps. The math checks passed: common-permutation quotient distance error was 0, alignment was non-worsening, and the interpolation derivative matched the target velocity up to 0.0021 finite-difference error. On synthetic graphs, alignment reduced mean squared pair distance from 9.891 to 0.0365 and velocity variance from 0.3949 to 0.00146; validation conditional MSE improved from 0.4423 to 0.00158 and endpoint MSE from about 0.448 to 0.00158. This is a clear signal for the proposed effect, but not evidence yet on realistic molecular graphs or with matching cost included.
- Agent confidence
- 8/10
- Baseline
- Raw/unpaired targets: train velocity variance 0.3949149; validation conditional MSE 0.4422936; Euler endpoint MSE at 4/8/16 steps = 0.4471358/0.4482336/0.4488158.
- Idea
- Exact quotient-aligned targets: velocity variance 0.0014580; validation conditional MSE 0.0015799; Euler endpoint MSE at 4/8/16 steps = 0.0015771/0.0015781/0.0015787. Mean aligned squared pair distance was 0.0365341 versus baseline 9.8910437.
Limitations: This used synthetic 5-node dense graphs, exact factorial permutation search, and a non-equivariant MLP rather than a molecular dataset and equivariant GNN. Matching/search runtime was not charged to training FLOPs, graph validity/FID was not measured, approximate matching stability was not tested, and results are from one fixed seed and one small model.
How to run: python3 experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- experiment.py 5.5 KB View Raw
- report.md 1.8 KB View
- results.json 1.3 KB View Raw