Research map / Graph learning / GNN architectures
Unsupervised Learning for Optimal Transport plan prediction between unbalanced graphs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains a cross-attention GNN, conditioned on FUGW hyperparameters, to predict unbalanced Gromov-Wasserstein transport plans between graphs in an unsupervised way by minimizing the FUGW loss directly.
Problem
Computing fused unbalanced Gromov-Wasserstein (FUGW) optimal transport plans between graphs requires solving non-convex problems with cubic-cost iterative solvers, which does not scale to large graphs or repeated hyperparameter selection.
Approach
Amortized, unsupervised learning of FUGW plans: a GNN with cross-graph attention, conditioned on (alpha, rho) via Fourier encoding and node-weight prediction for unbalanced mass, is trained by directly minimizing the expected FUGW loss over graph pairs and hyperparameters, so no solver-computed plans are needed as labels.
Key results
- ULOT predicts FUGW plans up to 100x faster than classical solvers on 1000-node IBC graphs, with losses correlated with IBPP at Pearson 0.99.
- ULOT is about 10x faster than entropic Sinkhorn while achieving smaller loss error.
- Using ULOT plans as warm start accelerates IBPP convergence.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Predict transport plans with ULOT instead of classical solvers | FUGW on IBC test graph pairs · IBC test graph pairs | runtime: up to 100 times faster | Helped |
| Predict transport plans with ULOT instead of Sinkhorn | FUGW on IBC test graph pairs · IBC test graph pairs | speed and loss error: 10 times faster for a smaller error | Helped |
| Use ULOT-predicted plans rather than IBPP plans | FUGW on IBC test graph pairs · IBC test graph pairs | FUGW loss correlation: 0.99 Pearson correlation | Helped |
| Optimize the FUGW parameters for label propagation | ULOT on SBM pairs with similar clusters · SBM label propagation, similar clusters (1,2,3) | accuracy: 0.87± 0.092 | Helped |
| Optimize the FUGW parameters for label propagation | ULOT on SBM pairs with different clusters · SBM label propagation, different clusters (1,2,3) to (1,2) | accuracy: 0.74\pm 0.11 | Helped |
| Use ULOT as a warm start for IBPP rather than starting IBPP without it | IBPP solver · IBC FUGW | convergence speed | Helped |
| Increase $\rho$ for graphs sharing cluster 2 | ULOT transport plan · SBM graphs with clusters (1,2) and (2,3) | transport mass assigned to the shared cluster | Helped |
| Set $\alpha=1$ rather than using node-feature information | classical FUGW solver · SBM graphs with clusters (1,2,3) | cluster matching | Hurt |
Limitations
- Small prediction error limits use where high precision is needed.
- Experiments limited to graphs with n <= 10000 due to GPU memory.
- Training still requires computing the cubic FUGW loss for each plan.
Experimental setup
Datasets: simulated stochastic block model graphs Individual Brain Charting (IBC) fMRI cortical surfaces
Benchmarks: simulated SBM graphs IBC fMRI brain surfaces
Metrics: FUGW loss loss error relative to best solver computation time Pearson correlation of losses label propagation accuracy
Baselines: IBPP MM L-BFGS-B Sinkhorn Meta OT
Builds on: Graph Matching Network Meta OT Fused unbalanced Gromov-Wasserstein Flow matching positional encoding
Closest papers by meaning
- Submodular framework for structured-sparse optimal transport (ICML 2024 · similarity 0.86)
- Universal Neural Optimal Transport (ICML 2025 · similarity 0.85)
- Towards Pre-trained Graph Condensation via Optimal Transport (NeurIPS 2025 · similarity 0.85)
- Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport (NeurIPS 2024 · similarity 0.84)
- Light Unbalanced Optimal Transport (NeurIPS 2024 · similarity 0.84)
- On Partial Optimal Transport: Revising the Infeasibility of Sinkhorn and Efficient Gradient Methods (AAAI 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.