Supervised Gromov-Wasserstein Optimal Transport

TL;DR

The paper addresses the limitation of traditional Gromov-Wasserstein in enforcing application-driven geometry constraints and partial couplings. It introduces supervised Gromov-Wasserstein (sGW), which imposes an infinity-pattern on the cost tensor to govern distance preservation, and reduces the resulting 4th-order constraint to a tractable 2nd-order form via minimal vertex covers. A complete numerical pipeline combines greedy MVC, a KL Mirror-C descent-based sOT solver, and an entropic regularization framework, with convergence guarantees for the nonconvex optimization. The authors demonstrate sGW’s ability to automatically regulate preserved geometry, enabling stable, flexible data matching across synthetic point clouds, complex surfaces, and real scRNA-seq datasets, including downsampling-based scalability strategies.

Abstract

We introduce the supervised Gromov-Wasserstein (sGW) optimal transport, an extension of Gromov-Wasserstein by incorporating potential infinity patterns in the cost tensor. sGW enables the enforcement of application-induced constraints such as the preservation of pairwise distances by implementing the constraints as an infinity pattern. A numerical solver is proposed for the sGW problem and the effectiveness is demonstrated in various numerical experiments. The high-order constraints in sGW are transferred to constraints on the coupling matrix by solving a minimal vertex cover problem. The transformed problem is solved by the Mirror-C descent iteration coupled with the supervised optimal transport solver. In the numerical experiments, we first validate the proposed framework by applying it to matching synthetic datasets and investigating the impact of the model parameters. Additionally, we successfully apply sGW to real single-cell RNA sequencing data. Through comparisons with other Gromov-Wasserstein variants on real data, we demonstrate that sGW offers the novel utility of controlling distance preservation, leading to the automatic estimation of overlapping portions of datasets, which brings improved stability and flexibility in data-driven applications.

Figures (12)

• Figure 1: Statistically decreasing trend of the transported mass $s$ with respect to the size of MVC. Left: two point clouds generated over two disks, one of radius 1 (blue), the other of radius 2 (red). Middle: for fixed $\rho = 0.2$, 100 MVCs are generated and the corresponding transported masses are calculated by sOT solver. A clear decreasing trend is observed. Right: for various values of $\rho$, MVCs and the corresponding masses are calculated, from which a similar decreasing trend is presented.
• Figure 2: The impact of the step size $\eta$ on sGW matching on synthetic point-cloud data.
• Figure 3: The impact of the shreshold $\rho$ on thesGW matching on synthetic point-cloud data.
• Figure 4: The $\gamma$-effect on the sGW matching on synthetic point-cloud data.
• Figure 5: sGW matching for smaller datasets and the intricate graph structure in the MVC calculation.

Theorems & Definitions (7)

• Definition 2.1
• Definition 2.2
• Definition 3.1
• Definition 3.2
• Lemma 3.1
• proof
• Theorem 3.1