Gromov-Wasserstein at Scale, Beyond Squared Norms
Abstract
A fundamental challenge in data science is to match disparate point sets with each other. While optimal transport efficiently minimizes point displacements under a bijectivity constraint, it is inherently sensitive to rotations. Conversely, minimizing distortions via the Gromov-Wasserstein (GW) framework addresses this limitation but introduces a non-convex, computationally demanding optimization problem. In this work, we identify a broad class of distortion penalties that reduce to a simple alignment problem within a lifted feature space. Leveraging this insight, we introduce an iterative GW solver with a linear memory footprint and quadratic (rather than cubic) time complexity. Our method is differentiable, comes with strong theoretical guarantees, and scales to hundreds of thousands of points in minutes. This efficiency unlocks a wide range of geometric applications and enables the exploration of the GW energy landscape, whose local minima encode the symmetries of the matching problem.
Lay Summary
Data in Machine Learning and imaging is frequently represented as point clouds. Because these sets of points often originate from distinct data modalities or different acquisition devices, direct comparison is usually impossible. Yet, many applications require transferring knowledge across these disparate data types by identifying shared patterns. It is hence crucial to establish meaningful point-to-point correspondences to process this data effectively, driving the need for fast and accurate matching techniques between point sets embedded in distinct feature spaces. Among existing algorithms, the Gromov-Wasserstein framework stands out as a highly promising solution due to its polyvalence and its strong mathematical guarantees. However, it is also very computationally demanding. Until now, the only methods capable of scaling to large, real-world datasets rely on heuristic approximations that introduce important artifacts into the resulting matches. In our work, we identify a wide range of scenarios in which the Gromov-Wasserstein problem can be significantly simplified. Specifically, we reformulate it as a sequence of optimal transport problems, a simpler and extensively researched class of optimization tasks. This reformulation allows us to apply decades of efficient numerical techniques developed for optimal transport directly to the Gromov-Wasserstein problem, dramatically accelerating its computation. Furthermore, by linking Gromov-Wasserstein to optimal transport, we provide new theoretical insights and numerical guarantees. These contributions make it possible to find accurate matchings between massive datasets across different spaces in just seconds or minutes, enabling the processing of heterogeneous data on inputs more than 10 times larger than existing baselines.