Hierarchical Refinement: Optimal Transport to Infinity and Beyond

Hierarchical Refinement: Optimal Transport to Infinity and Beyond

Speakers: Peter Halmos, Julian Gold, Xinhao Liu, Benjamin Raphael

Conference: ICML 2025

YouTube: https://slideslive.com/39043994

Overview

Optimal Transport (OT) is a powerful mathematical framework for comparing probability distributions by finding the least-cost mapping or coupling between them. It has garnered significant attention in machine learning for its ability to model least-action principles in generative models, facilitate domain-domain translation, and enable tasks like shape registration and barycenter computation in computer vision and computational biology. However, classical OT algorithms, such as the network simplex and Hungarian algorithm, suffer from prohibitive cubic time complexity and quadratic space complexity, rendering them impractical for the massive datasets prevalent in modern AI. While methods like Sinkhorn regularization and mini-batch OT offer some scalability improvements, they often introduce approximations, batch bias, or sacrifice the global, bijective nature of the optimal map.

This talk, presented by Peter Halmos and co-authored with Julian Gold, Xinhao Liu, and Benjamin Raphael, introduces Hierarchical Refinement (HR), a novel multi-scale optimal transport algorithm designed to overcome these long-standing scalability limitations. HR addresses the critical open question of achieving global, full-rank, and bijective Monge maps with linear space and log-linear time scaling. By leveraging a key local-to-global optimality property derived from low-rank optimal transport, HR can efficiently compute these maps for datasets comprising millions of points, opening new avenues for applying optimal transport in large-scale deep learning pipelines and other data-intensive applications.

The significance of Hierarchical Refinement lies in its ability to bridge the gap between the theoretical elegance and practical applicability of optimal transport. By providing a method that scales gracefully to modern dataset sizes—from single-cell transcriptomics to high-dimensional image embeddings—HR enables practitioners to harness the full power of OT without compromising on map quality or incurring the computational overhead that previously made it infeasible. This breakthrough promises to unlock new capabilities in generative modeling, data alignment, and a host of other machine learning tasks that demand robust and scalable distribution comparisons.

Background

▶ Watch: Introduction to Optimal Transport and problem statement (0:00)

Optimal Transport (OT) originated from Gaspard Monge's 1781 problem of moving a pile of earth to another location with minimal effort. Mathematically, OT seeks to find the most efficient way to transform one probability distribution, A (source), into another, B (target). There are two primary formulations:

1. Kantorovich Problem: This formulation seeks a least-cost coupling P (a joint distribution) between A and B, minimizing the expected cost across all pairs of points (X, Y). The minimum value of this cost is known as the Wasserstein distance. It is framed as a minimization over P in the space of all couplings between A and B.
2. Monge Problem: This formulation seeks a least-cost mapping T that pushes the source distribution A forward onto the target distribution B. When such a Monge map exists, it coincides with the optimal coupling P, and the problems are equivalent. The Monge map provides a bijective assignment between points of the two distributions.

The appeal of OT in machine learning stems from its ability to encode a "least-action principle." Applications are diverse and growing, including:

• Generative Models: Training models that adhere to a least-action principle, making them more stable and effective.
• Domain-Domain Translation: Finding least-cost mappings between different data distributions, e.g., transforming faces into cartoons.
• Shape and Point Cloud Registration: Aligning geometric data by minimizing the cost of transforming one shape into another.
• Barycenter Computation: Finding the "average" of multiple distributions.
• Computational Biology: Applications in areas like single-cell transcriptomics for comparing cell populations.

However, the computational cost of classical OT algorithms has been a significant barrier.

• Network Simplex and the Hungarian Algorithm solve the Monge problem as a bijective assignment problem, incurring a cubic time complexity (O(N^3)) and quadratic space complexity (O(N^2)) for N points. These methods are practical only for small datasets, typically up to a few thousand points.
• The Sinkhorn algorithm was a major breakthrough, regularizing the Kantorovich problem with entropy. It reduces the time complexity to quadratic (O(N^2)) and scales to tens of thousands of points. While faster, it provides an approximation and still struggles with larger N.

The advent of "modern datasets" has further exacerbated this scalability challenge. Applications in single-cell brain imaging, image classification (e.g., ImageNet), and text-image datasets now involve hundreds of thousands to tens of millions, or even billions, of points. Existing methods simply cannot keep pace:

• Mini-batch Optimal Transport (mini-batch OT) scales linearly in time and space, but it introduces severe batch bias because alignments are computed independently on small, sampled batches. This results in local, rather than global, alignments.
• Low-rank Optimal Transport (low-rank OT) also scales linearly in time and space and provides global alignments. However, by definition, it is intrinsically low-resolution and constrained by rank, forfeiting bijective correspondences.
• Multi-scale Optimal Transport (MOPT) methods, such as those by Gerber and Maggioni, construct multi-scale partitions to find optimal alignments. While promising, these often depend on a predefined gridding of the space and can suffer from a "curse of dimensionality" if not carefully constructed, leading to solutions that vary with the input grid.

The core problem addressed by Hierarchical Refinement is the lack of an algorithm that can compute global, full-rank, and bijective optimal transport Monge maps with linear space and log-linear time scaling for datasets of modern scale.

Key Findings

▶ Watch: Limitations of classical OT for modern large datasets (2:00)

The talk presents Hierarchical Refinement (HR) as a groundbreaking solution to the scalability challenges of Optimal Transport, offering several key findings:

• Scalability Breakthrough: HR computes global, bijective Monge maps in linear space complexity (O(N)) and log-linear time complexity (O(N log N)) for millions of points, assuming the rank of the cost matrix (D) is small relative to N. This significantly surpasses the capabilities of existing methods, which typically become intractable beyond 1-2 million points.
• Multi-Scale, Data-Intrinsic Partitioning: Unlike previous multi-scale OT algorithms that rely on predefined grids, HR constructs a multi-scale partition of the datasets in a manner that is intrinsic to the data itself. This avoids the "curse of dimensionality" associated with fixed gridding and ensures that the alignment quality is robust and data-driven.
• Leveraging Low-Rank OT Optimality: HR proves and leverages a key local-to-global optimality property of a specific low-rank optimal transport formulation (introduced by Meyer-Scetbon in 2021) as a recursive subroutine. This property is crucial for maintaining optimality across different scales of the hierarchy.
• Guaranteed Bijective Mapping: The hierarchical refinement process is designed to recover a bijective T mapping between point pairs at the finest scale of the hierarchy. This ensures that the generated map is a true one-to-one assignment, which is essential for many applications like generative modeling and flow matching.
• Implicit Block Coupling: HR implicitly defines a block coupling that, through a simple inductive argument, can be guaranteed to find an optimal solution for N points in the PDE Monge map. This parallels the structure of hierarchical matrices used in PDE solvers.
• Robustness to Cost Matrix Type: The approach is not limited to costs induced by kernels in an inner product space. It can approximate costs induced by other distance metrics (e.g., Euclidean distance squared) that satisfy appropriate metric properties, with the rank D growing with the ambient dimension.

In summary, Hierarchical Refinement delivers a computationally efficient, high-quality, and robust method for large-scale optimal transport, making it feasible to apply the full power of the Monge map to previously intractable problem sizes.

Technical Deep Dive

▶ Watch: Existing linear scaling OT methods and their drawbacks (2:40)

Hierarchical Refinement (HR) is fundamentally a multi-scale optimal transport algorithm that specifically targets the Monge problem to find global, bijective mappings. Its core innovation lies in its recursive, divide-and-conquer strategy, which leverages a specialized low-rank optimal transport (low-rank OT) optimization problem as a subroutine to build data-intrinsic partitions.

The foundation of HR's recursive partitioning is a specific low-rank OT formulation introduced in 2021 by Meyer-Scetbon. This formulation considers an optimization problem that replaces the traditional marginal constraint G with a constraint on G corresponding to hard clustering functions Q-star and R-star. These functions define a bi-partition of the two datasets, X (source, $\mu$) and Y (target, $\nu$), being aligned. For a rank R=2, it is shown that this problem has optimal solutions where Q-star and R-star effectively co-cluster points X with their image T-star(X) under the optimal mapping. This co-clustering property implies a cyclic monotonicity property, which is a key characteristic of Monge maps.

The HR algorithm operates recursively:

1. Partitioning: At each level of the hierarchy, the current set of points from both distributions is divided into sub-partitions. This division is achieved by solving the specialized low-rank optimization problem for Q-star and R-star. These clustering functions intrinsically identify how points from the source distribution should be grouped with their corresponding images in the target distribution.
2. Local-to-Global Optimality Invariant: A crucial aspect of HR is the maintenance of a local-to-global optimality invariant. This means that optimal alignments found within sub-partitions contribute to the overall global optimal alignment. The proof of this invariant ensures that independently recursing on each sub-partition does not compromise the global optimality of the final map.
3. Recursion: The algorithm then recurses independently on each sub-partition. This process continues until the finest scale is reached, at which point a direct bijective T mapping is recovered between the individual point pairs within the smallest clusters.
4. Implicit Block Coupling: The hierarchical structure generated by this process implicitly defines a block coupling. This coupling, through a simple inductive argument, is proven to yield an optimal solution for the N points in the PDE Monge map. This parallels the efficiency achieved by hierarchical matrices in solving Partial Differential Equations (PDEs).

Computational Complexity:

• Space Complexity: HR achieves linear space complexity (O(N)) for N points, primarily for storing the partitions.
• Time Complexity: The time complexity is stated as **O(N D log N)**, where D is the rank of the cost matrix C.
• The log N factor arises from the hierarchical, multi-scale nature of the algorithm (similar to a balanced binary tree traversal).
• The D factor comes from the low-rank solver subroutine.
• A critical condition for achieving the desired log-linear time scaling is that D must be small relative to N. For instance, if D represents the ambient dimension (e.g., 3, 256, 2056) and N represents the sample complexity (e.g., 1 million, 100 million, 1 billion), then D is effectively a constant in the asymptotic sense. If D were to equal N, the complexity would revert to O(N^2), recovering quadratic complexity.

Cost Matrices and Generalizability:

HR is not restricted to cost matrices induced by kernels in an inner product space (e.g., the squared Euclidean distance). For other distances that satisfy appropriate metric properties, HR can still approximate these costs. In such cases, the effective rank D might grow with the ambient dimension of the data, but the linear-time algorithms for low-rank approximation still allow for efficient processing.

Comparison with Multi-scale Optimal Transport (MOPT):

The talk explicitly contrasts HR with previous MOPT methods. While MOPT can be full-rank and scalable, its solution often varies with the input grid, and it can suffer from a "curse of dimensionality" due to its reliance on predefined spatial partitions. In contrast, HR automatically and intrinsically builds its partitions based on the dataset itself, which means the actual alignment quality does not suffer from arbitrary partitioning choices or the curse of dimensionality in the same way. HR maintains a full-rank, global, and bijective solution computed with respect to all points in the dataset.

This technical framework allows Hierarchical Refinement to achieve unprecedented scalability for optimal transport, pushing its applicability into domains previously considered computationally infeasible.

Experimental Setup & Results

▶ Watch: Hierarchical Refinement's recursive algorithm and properties (3:40)

The Hierarchical Refinement (HR) algorithm was benchmarked on several challenging, large-scale datasets to demonstrate its performance and scalability, particularly in comparison to existing optimal transport methods.

Datasets Used for Benchmarking:

• Single-cell Transcriptomics: This dataset involved hundreds of thousands of points (cells), each with an ambient dimension of 60, representing gene expression profiles. This is a common and growing area in computational biology where large-scale cell comparisons are crucial.
• ImageNet Embeddings: A massive dataset comprising millions of points, representing high-dimensional embeddings (dimension 2000) of images from the ImageNet collection. This showcases HR's ability to handle high-dimensional data at scale, which is critical for modern computer vision tasks.
• Merfish Brain Imaging Data: Another biologically relevant dataset, likely involving spatial transcriptomics data, which presents complex alignment challenges.
• Synthetic Datasets: A number of synthetic datasets were also used, presumably to control for specific properties and systematically evaluate the algorithm's performance under various conditions.

Baselines and Comparative Performance:

• The primary baseline mentioned is Sinkhorn, a widely used entropy-regularized optimal transport algorithm.
• The talk emphasizes that while Sinkhorn solutions are comparable to HR in terms of cost, HR significantly outperforms Sinkhorn and "most methods" in terms of scalability. Specifically, the presentation states that HR scales beyond 1-2 million points, a threshold where most other optimal transport methods typically fail to run or become impractically slow.

Headline Results and Metrics:

• Cost Comparability: The solutions found by HR were found to be comparable to Sinkhorn in cost. This is a crucial finding, indicating that HR achieves high-quality optimal transport maps without sacrificing accuracy for scalability.
• Linear Space and Log-Linear Time Scaling: The most significant result is that HR achieves linear space complexity (O(N)) and log-linear time complexity (O(N log N)) with respect to the input size N. This scaling was observed across the massive and high-dimensional datasets tested, assuming that the rank D of the cost matrix is asymptotically constant relative to N.
• Beyond Million-Point Scale: HR successfully processed datasets in the order of hundreds of thousands to millions of points, including those with high ambient dimensions (e.g., ImageNet embeddings with dimension 2000), where traditional and even many modern OT methods struggle or simply cannot execute within practical timeframes.
• Global, Bijective, Full-Rank Maps: Unlike mini-batch or low-rank approaches, HR consistently delivers global, bijective, and full-rank Monge maps, maintaining the fidelity and interpretability of the optimal assignment.

In summary, the experimental results robustly demonstrate that Hierarchical Refinement provides a highly scalable and accurate solution for computing optimal transport maps on datasets that are orders of magnitude larger than what was previously feasible. This empirical validation underscores HR's potential to unlock new applications in data-intensive fields.

Practical Implications

▶ Watch: Guaranteed bijective mapping and complexity analysis (4:10)

The development of Hierarchical Refinement carries profound practical implications for practitioners across various domains, particularly within machine learning infrastructure teams, model builders, and deployers.

1. Revolutionizing Deep Learning Pipelines with OT:

One of the most significant implications is the paradigm shift it offers for integrating Optimal Transport into deep learning workflows. Current methods for scaling OT in deep learning, such as mini-batch optimal transport, suffer from severe limitations:

• Batch Bias: Mini-batch OT computes local alignments on independent batches, leading to a biased overall solution that does not represent the true global optimal map.
• Rerunning OT: It requires rerunning the OT computation for every training iteration (T iterations), which is computationally expensive and slow, especially for large models and datasets.
• Entropic Bias: If Sinkhorn is used, it introduces an entropic regularization bias.

Hierarchical Refinement offers a powerful alternative: precomputing global OT maps offline, once and for all. These precomputed optimal global solution pairs can then be treated as if they were standard supervised training data points. This approach completely eliminates batch bias and the need for repetitive OT computations during training. For instance, in applications like flow matching for generative models, having a precomputed, globally optimal Monge map can significantly enhance training stability and performance.

2. Unlocking Previously Infeasible Applications:

The linear space and log-linear time scaling of HR opens the door to applying optimal transport to datasets and problems that were previously computationally infeasible. This includes:

• Large-scale Generative Models: Training models that adhere to the least-action principle of OT on vast image, text, or multimodal datasets.
• High-Dimensional Data Alignment: Aligning complex, high-dimensional data, such as single-cell multi-omics data (e.g., transcriptomics, epigenomics) or large-scale brain imaging data (e.g., Merfish), where millions of data points need precise correspondence.
• Global Domain Adaptation: Performing global, rather than local, domain translation or adaptation across extremely large and diverse datasets.
• Enhanced Data Preprocessing: Using OT for sophisticated data preprocessing steps (e.g., outlier detection, data cleaning, feature matching) on scales previously unimaginable.

3. Tradeoffs and Limitations:

While HR represents a significant leap, it's important to consider its practical tradeoffs and limitations:

• Dependence on Cost Matrix Rank (D): The log-linear time complexity is contingent on the rank D of the cost matrix being small relative to the number of points N. If D is large (e.g., approaching N), the complexity can degrade towards quadratic, which would diminish its advantage for certain problem types. Practitioners need to assess the intrinsic dimensionality or effective rank of their cost matrices.
• Complexity of Implementation: Implementing Hierarchical Refinement, with its recursive structure and reliance on specialized low-rank OT subroutines, is likely more complex than using off-the-shelf Sinkhorn or mini-batch OT implementations.
• Memory Footprint: While linear in N, the constant factor for memory could still be substantial for extremely large N, depending on the specific data types and auxiliary structures.

In essence, Hierarchical Refinement provides a robust and scalable toolkit for practitioners to leverage the full power of optimal transport in modern, data-intensive machine learning environments. It reduces computational bottlenecks and improves the quality of OT-based solutions, enabling the development of more accurate and principled AI systems.

Key Takeaways

• Scalability Solved: Hierarchical Refinement (HR) offers a breakthrough in scaling Optimal Transport (OT), enabling the computation of global, bijective Monge maps for datasets with millions of points.
• Log-Linear Performance: HR achieves linear space complexity and log-linear time complexity (O(N log N)), significantly outperforming classical and existing approximate OT methods for large-scale data.
• Data-Intrinsic Partitioning: The algorithm uses a recursive, divide-and-conquer approach, building multi-scale partitions intrinsically from the data itself by leveraging a key local-to-global optimality property of low-rank OT.
• Enhanced Deep Learning Pipelines: HR allows for the precomputation of global OT maps offline, eliminating batch bias and computational overhead during deep learning training, thus improving generative models and flow matching applications.
• Full-Rank, Bijective Maps: Unlike many scalable OT approximations, HR guarantees the computation of full-rank, bijective Monge maps, preserving the high fidelity and interpretability of the optimal assignments.
• Broad Applicability: This method opens up previously infeasible applications for OT across machine learning, computational biology, and computer vision, especially for high-dimensional, large-scale data.

About the Speaker(s)

The talk "Hierarchical Refinement: Optimal Transport to Infinity and Beyond" was presented by Peter Halmos. His co-authors on this work include Julian Gold, Xinhao Liu, and Benjamin Raphael. While specific titles and affiliations were not detailed in the transcript, their collective work focuses on advancing the theoretical and practical aspects of optimal transport, particularly its scalability and application in machine learning and related fields.