Hierarchical Refinement: Optimal Transport to Infinity and Beyond

2025-03-04

AI Summary

Optimal transport (OT) suffers from quadratic space complexity under large-scale data due to the Sinkhorn algorithm, hindering scalable computation of exact bijective Monge maps. Method: We propose the first linear-space, log-linear-time OT framework capable of computing exact bijective Monge mappings. Our approach leverages the inherent co-clustering structure embedded in low-rank OT factors, enabling a hierarchical multi-scale partitioning and iterative refinement scheme. At each level, it progressively approximates the globally optimal bijective coupling via low-rank approximations augmented with Monge structural priors. Contribution/Results: Our method achieves the first scalable, full-rank OT solution on datasets exceeding one million points. It guarantees strictly linear memory complexity while matching the matching accuracy of exact OT—thereby substantially surpassing current scalability limits.

Technology Category

• Machine Learning: Scalability of ML Systems
• Search and Optimization: Non-convex Optimization
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Security and Privacy: Large-scale security measurements

Abstract

Optimal transport (OT) has enjoyed great success in machine-learning as a principled way to align datasets via a least-cost correspondence. This success was driven in large part by the runtime efficiency of the Sinkhorn algorithm [Cuturi 2013], which computes a coupling between points from two datasets. However, Sinkhorn has quadratic space complexity in the number of points, limiting the scalability to larger datasets. Low-rank OT achieves linear-space complexity, but by definition, cannot compute a one-to-one correspondence between points. When the optimal transport problem is an assignment problem between datasets then the optimal mapping, known as the Monge map, is guaranteed to be a bijection. In this setting, we show that the factors of an optimal low-rank coupling co-cluster each point with its image under the Monge map. We leverage this invariant to derive an algorithm, Hierarchical Refinement (HiRef), that dynamically constructs a multiscale partition of a dataset using low-rank OT subproblems, culminating in a bijective coupling. Hierarchical Refinement uses linear space and has log-linear runtime, retaining the space advantage of low-rank OT while overcoming its limited resolution. We demonstrate the advantages of Hierarchical Refinement on several datasets, including ones containing over a million points, scaling full-rank OT to problems previously beyond Sinkhorn's reach.

Problem

• Overcomes Sinkhorn's quadratic space complexity
• Enables bijective coupling in large datasets
• Scales optimal transport to million-point datasets

Innovation

• Hierarchical Refinement for scalable optimal transport
• Low-rank OT with linear-space complexity
• Dynamic multiscale partition for bijective coupling