🤖 machine learning

A Riemannian Approach to Low-Rank Optimal Transport

This paper proposes a unified Riemannian geometric framework for low-rank optimal transport that models factored couplings as smooth submanifolds equipped with the Fisher-Rao metric, enabling efficient, regularization-free first- and second-order solvers with linear complexity and superior convergence across balanced, unbalanced, and various optimal transport variants.

Original authors: Pratik Jawanpuria, Bamdev Mishra

Published 2026-06-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Pratik Jawanpuria, Bamdev Mishra

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to move a massive pile of sand from one pile (the source) to another pile (the target). In the world of mathematics and machine learning, this is called Optimal Transport. The goal is to figure out the most efficient way to move every grain of sand so that the total "effort" (or cost) is as low as possible.

For a long time, doing this for huge piles of sand was incredibly slow and expensive, like trying to map out a route for every single grain individually.

The Problem: The "Low-Rank" Shortcut

To speed things up, researchers came up with a clever shortcut called Low-Rank Optimal Transport. Instead of moving sand directly from every grain in the source to every grain in the target, they imagine a small group of central hubs (like major train stations).

All sand from the source goes to these hubs first.
Then, the hubs redistribute the sand to the targets.

This drastically reduces the number of connections you need to calculate. However, the paper points out a major flaw in how current computers solve this: they use a clumsy, trial-and-error method (called "mirror descent") that is slow, requires a lot of manual tweaking (like adjusting the sensitivity of a radio dial), and often gets stuck in local loops.

The Solution: A New Geometric Map

The authors of this paper propose a completely new way to navigate this problem using Riemannian Geometry.

Think of the possible solutions as a landscape.

Old Way: Imagine walking through a dense, foggy forest where the ground is uneven. You take small, cautious steps, constantly checking if you're going the right way, but you don't know the shape of the hills or valleys. You might get stuck in a small dip thinking it's the bottom of the valley.
New Way: The authors realize that the "forest" is actually a smooth, curved surface (a manifold). They equip this surface with a special map (the Fisher-Rao metric) that understands the true shape of the terrain.

Because they understand the shape of the land, they can use powerful tools:

First-Order Solvers: Like a hiker who knows the slope of the hill and walks straight down the steepest path.
Second-Order Solvers: Like a hiker who also knows the curvature of the hill. They can predict where the path will bend and take a giant, confident leap toward the bottom, rather than taking tiny, hesitant steps.

The Magic Trick: "Unbalanced" Transport

The paper makes a special breakthrough for a scenario called Unbalanced Transport. In real life, sometimes the source pile of sand is bigger than the target, or vice versa. You can't just move everything; you have to decide what to discard or what to create.

The Old Way: To handle this, computers had to run a complex, repetitive inner loop (like a robot checking its work 100 times before taking a single step). This was slow.
The New Way: The authors discovered that on their new geometric map, the rules for "unbalanced" sand are so simple that the computer can calculate the answer instantly with a single formula. No loops, no waiting. It's like realizing that instead of walking around a lake, you can just build a bridge across it in one step.

The Results: Faster and Smarter

The authors tested their new "geometric hikers" against the old "forest walkers" on massive datasets (up to 50,000 points).

Speed: Their method was often orders of magnitude faster. While the old methods took minutes or hours, the new method finished in seconds.
Accuracy: They reached better solutions (lower costs) without needing to manually tune any settings.
Confidence: They even built a "certificate" (a mathematical test) that tells you, "Yes, this is the absolute best solution possible," or "You are close, but here is exactly how to improve it."

Summary

In short, this paper takes a difficult, slow, and finicky math problem (moving data distributions efficiently) and reimagines it as a smooth journey on a curved surface. By using the right map and tools, they eliminated the need for slow, repetitive checking and manual tuning, allowing computers to solve these problems much faster and more accurately than ever before.

Technical Summary: A Riemannian Approach to Low-Rank Optimal Transport

Problem Statement

Classical Optimal Transport (OT) provides a rigorous metric for comparing probability distributions but suffers from a quadratic computational bottleneck ($O(mn)$) due to the need to compute dense coupling matrices. While entropic regularization enables the use of the Sinkhorn algorithm, it retains this scaling and introduces bias. To mitigate this, Low-Rank Optimal Transport (LOT) constrains the transport plan to a low-dimensional structure, factoring the coupling matrix $\Gamma$ into sub-couplings $(U, V)$ sharing a latent marginal $g$ .

Despite the structural efficiency of low-rank formulations, existing solvers (e.g., LOT, LR-GW, UB-LOT) rely heavily on first-order mirror-descent updates coupled with iterative inner loops (Dykstra or Bregman projections). These approaches face several limitations:

Hyperparameter Sensitivity: They require careful tuning of step-sizes and entropic regularization parameters.
Ignored Curvature: They fail to exploit the local curvature of the optimization landscape, leading to slower convergence.
Computational Overhead: The inner iterative loops for projection and retraction increase per-iteration complexity, particularly in unbalanced settings.

Methodology

The authors propose a unified Riemannian geometric framework that models the set of rank- $r$ positive factored couplings as smooth embedded submanifolds within the positive orthant. By equipping these manifolds with the Fisher-Rao product metric, the authors derive tractable formulations for all necessary geometric ingredients, enabling robust first- and second-order optimization.

1. Manifold Formulation

The framework defines two distinct manifolds based on marginal constraints:

Balanced Manifold ( $\mathcal{M}_{bal}$ ): Enforces strict reconstruction of source and target marginals ( $U\mathbf{1}=a, V\mathbf{1}=b$ ) and latent consistency ( $U^\top\mathbf{1}=V^\top\mathbf{1}$ ).
Unbalanced Manifold ( $\mathcal{M}_{ub}$ ): Relaxes outer marginal constraints, requiring only latent consistency ( $U^\top\mathbf{1}=V^\top\mathbf{1}$ ), with marginal deviations penalized via KL divergence in the objective.

2. Geometric Ingredients

Under the Fisher-Rao metric, the authors derive explicit expressions for:

Orthogonal Projectors:
- For $\mathcal{M}_{bal}$ , projection involves solving a KKT system (reduced to a Schur complement) via Conjugate Gradients (CG).
- For $\mathcal{M}_{ub}$ , the projection simplifies to a closed-form expression, eliminating the need for iterative solvers.
Retractions:
- The retraction maps tangent vectors back to the manifold via a Fisher-consistent exponential mirror step followed by a KL projection.
- For $\mathcal{M}_{bal}$ , this requires cyclic Bregman projections (iterative).
- For $\mathcal{M}_{ub}$ , the KL projection collapses into a closed-form column scaling, removing inner loops entirely.
Riemannian Gradient and Hessian:
- The framework provides a cost-agnostic formulation where the Riemannian gradient and Hessian-vector products depend only on the Euclidean derivatives of the specific cost function (Linear OT, Gromov-Wasserstein, etc.) and the manifold geometry.
- This enables the use of second-order solvers (Trust-Region, Conjugate Gradient) which utilize curvature information.

3. Rank-Sufficiency Certificate

The paper introduces a certificate for global optimality verification. By computing a "reduced gradient" (dual slack matrix) at a rank- $r$ critical point, the method can verify if the solution is globally optimal for the full-rank problem. If the rank is insufficient, the certificate identifies a direction for rank increase.

Key Contributions

Unified Riemannian Formulation: Formalizes balanced and unbalanced low-rank OT as smooth manifolds under the Fisher-Rao metric, providing a single geometric framework for linear OT, Gromov-Wasserstein (GW), and their unbalanced variants.
Closed-Form Unbalanced Operations: Demonstrates that for unbalanced OT, tangent space projectors and retractions reduce to closed-form expressions. This eliminates the inner Dykstra/Bregman loops required by state-of-the-art methods (e.g., UB-LOT), reducing per-iteration complexity to $O(rN)$.
Second-Order Solvers: Enables the use of curvature-aware solvers (CG, Trust-Region) that achieve faster convergence than first-order mirror descent, without requiring entropic regularization or step-size tuning.
Global Optimality Verification: Provides a computable certificate to verify if a low-rank solution is globally optimal or to guide rank-increase strategies.

Experimental Results

Extensive experiments across balanced and unbalanced settings (Linear OT, GW, Fused GW) with dataset sizes up to $N=50,000$ demonstrate:

Convergence Speed: Riemannian solvers converge in an order-of-magnitude fewer outer iterations compared to mirror-descent baselines.
Wall-Clock Time: The elimination of inner loops in the unbalanced setting leads to significant speedups. For example, in unbalanced OT with outliers, Riemannian solvers converge in seconds, whereas UB-LOT times out or diverges.
Performance: The methods achieve superior or comparable objective values without the entropic bias inherent in regularized baselines.
Second-Order Advantage: On non-convex problems (GW, Fused GW), Trust-Region methods leverage Hessian information to escape shallow local minima, outperforming first-order methods.

Significance

The paper claims that by shifting from a first-order mirror-descent paradigm to a Riemannian geometric perspective, low-rank OT solvers can overcome the limitations of ill-conditioning and hyperparameter sensitivity. The primary significance lies in the structural simplification of unbalanced operations, which transforms a computationally expensive iterative process into a closed-form evaluation. This allows for regularization-free, tuning-free optimization that scales linearly with dataset size while providing rigorous second-order convergence guarantees and global optimality verification. The framework is presented as a general, cost-agnostic approach applicable to a wide spectrum of optimal transport problems.