Gradient Flows and Riemannian Structure in the Gromov-Wasserstein Geometry

TL;DR

This work develops gradient-flow theory for the inner product Gromov–Wasserstein (IGW) geometry on $\mathcal{P}_2(\mathbb{R}^d)$, identifying a global mobility-based transformation that maps the Wasserstein gradient to an IGW gradient. It introduces a JKO-style minimizing-movement scheme with orthogonal-alignment steps, proves convergence to an IGW generalized minimizing movement (GMM) curve satisfying the continuity equation with velocity $v_t= -\mathcal{L}_{\mathbf{\Sigma}_t,\rho_t}^{-1}[\nabla\delta\mathsf{F}(\rho_t)]$, and derives a Benamou–Brenier–like formulation for IGW. The analysis establishes the IGW Riemannian structure, connects IGW to Wasserstein geometry via the mobility operator, and provides numerical demonstrations of global-structure-preserving interpolations and flow-matching tasks. This framework enables structure-preserving transport and interpolation beyond standard OT, with potential implications for tasks where global geometry and orientations should be preserved. The work also outlines open questions about the IGW tangent space, convergence rates, and extensions to the full GW framework.

Abstract

The Wasserstein space of probability measures is known for its intricate Riemannian structure, which underpins the Wasserstein geometry and enables gradient flow algorithms. However, the Wasserstein geometry may not be suitable for certain tasks or data modalities. Motivated by scenarios where the global structure of the data needs to be preserved, this work initiates the study of gradient flows and Riemannian structure in the Gromov-Wasserstein (GW) geometry, which is particularly suited for such purposes. We focus on the inner product GW (IGW) distance between distributions on $\mathbb{R}^d$. Given a functional $\mathsf{F}:\mathcal{P}_2(\mathbb{R}^d)\to\mathbb{R}$ to optimize, we present an implicit IGW minimizing movement scheme that generates a sequence of distributions $\{ρ_i\}_{i=0}^n$, which are close in IGW and aligned in the 2-Wasserstein sense. Taking the time step to zero, we prove that the discrete solution converges to an IGW generalized minimizing movement (GMM) $(ρ_t)_t$ that follows the continuity equation with a velocity field $v_t\in L^2(ρ_t;\mathbb{R}^d)$, specified by a global transformation of the Wasserstein gradient of $\mathsf{F}$. The transformation is given by a mobility operator that modifies the Wasserstein gradient to encode not only local information, but also global structure. Our gradient flow analysis leads us to identify the Riemannian structure that gives rise to the intrinsic IGW geometry, using which we establish a Benamou-Brenier-like formula for IGW. We conclude with a formal derivation, akin to the Otto calculus, of the IGW gradient as the inverse mobility acting on the Wasserstein gradient. Numerical experiments validating our theory and demonstrating the global nature of IGW interpolations are provided.

Figures (11)

• Figure 1: GW versus Wasserstein interpolation between rotated cat shapes: the Wasserstein interpolation (second line) breaks the structure of the shape to minimize the transportation cost, while the GW interpolation (first line) respects the structure and produces the desired effect. See \ref{['sec:experiment']} for details on this experiment.
• Figure 2: Generalized IGW geodesic between between $\mu_1$ and $\mu_2$ w.r.t. $\mu_0$.
• Figure 3: Illustration of steps along IGW proximal point method from \ref{['eq:minimizing_movement']}: each purple ring represents an orbit of $\tilde{\rho}_i$ along the orthogonal group $\mathrm{O}(d)$ in $\mathbb{R}^d$, i.e., $\mathrm{O}(d)_\sharp\tilde{\rho}_i\coloneqq \{\mathbf{U}_\sharp\tilde{\rho}_i:\,\mathbf{U}\in\mathrm{O}(d)\}$; after each update step of the algorithm, the obtained distribution $\tilde{\rho}_{i+1}$ (shown in red, arbitrarily chosen from $\mathrm{O}(d)_\sharp\rho_i$) is transformed using any $\mathbf{O}\in\mathcal{O}_{\rho_i,\tilde{\rho}_{i+1}}$ to obtain $\rho_{i+1}=\mathbf{O}_\sharp\tilde{\rho}_{i+1}$ (shown in teal; see \ref{['def:PSD_transform']}). The transformation aligns each two consecutive steps and guarantees the the cross-covariance matrix between any two consecutive steps is PSD. This results in a regular and controlled path (marked in dark teal), whose convergence and limiting characterization we establish in \ref{['thm:main']}. The dashed orange line shows the path that would have been obtained without the transformations, which notably lacks stability.
• Figure 4: Illustration of the action of the mobility operator $\mathcal{L}_{\mathbf{\Sigma}_\mu,\mu}$ and its inverse. Fixing $\mu$ as uniform distribution over the grid, and we initiate the vector field as pointing inward except for the periphery. The first shows the repeated application of $\mathcal{L}_{\mathbf{\Sigma}_\mu,\mu}$ to the initial vector field, while the second shows the inverse $\mathcal{L}_{\mathbf{\Sigma}_\mu,\mu}^{-1}$.
• Figure 5: Visualization of cross-partition error function: the IGW distance between the sequences is at most the cross-partition error function, plus an $O(\sqrt{\tau}+\sqrt{\eta})$ term.

Theorems & Definitions (57)

• Theorem 2.1
• Lemma 2.1: IGW duality
• Lemma 2.2: Gromov-Monge map
• Definition 2.1: First variation
• Definition 2.2: Fréchet subdifferential
• Proposition 3.1: IGW pseudometric
• Definition 3.1: Cross-covariance PSD transform
• Lemma 3.1: Cross-covariance PSD tranform
• Lemma 3.2: IGW and Wasserstein comparison
• Proposition 3.2: Lipschitz IGW and Wasserstein curves