Supervised Gromov–Wasserstein Optimal Transport with Metric-Preserving Constraints

We introduce the supervised Gromov-Wasserstein (sGW) optimal transport, an extension of Gromov-Wasserstein that incorporates potential infinity entries in the cost tensor. These infinity entries enable sGW to enforce application-induced constraints on preserving pairwise distance to a certain extent. 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 apply sGW to aligning single-cell RNA sequencing data where the datasets are partially overlapping and only intra-dataset metrics are used. Through comparisons with other Gromov-Wasserstein variants, we demonstrate that sGW offers an additional utility of controlling distance preservation, leading to automatic estimation of overlapping portions of datasets, which brings improved stability and flexibility in data-driven applications. The codes for sGW and for reproducing the results are available on Github [https://github.com/zcang/supervisedGW].