Hierarchical method speeds up optimal transport across large datasets
Dual-guided Hierarchical Edge Localization for Large-scale Optimal Transport Across Dimensions
Machine Learning
Summary
Optimal transport helps compare and match data, but doing this exactly can be slow and use a lot of memory when data is big. The authors propose HELLO, a method that breaks the problem into smaller parts and uses smart guesses to quickly focus on important connections. This approach uses less memory and runs faster, even on datasets with millions of points and thousands of features. HELLO can also handle various cost types and works as a fast tool for many transport-related problems.
What this means in practice
- •For machine learning engineers: Efficiently compute large-scale optimal transport maps to align high-dimensional datasets for tasks like domain adaptation and data integration.
- •For computer vision developers: Run fast and memory-efficient transport computations for image matching and large-scale feature alignment in high-resolution images.
Authors
Wenzhou Xia, Qiaoqiao Ding, Jingwei Liang, Xiaoqun Zhang
Abstract
Optimal transport (OT) compares distributions and aligns datasets in machine learning, yet unregularized discrete OT requires a linear program with quadratically many transport variables. We propose HELLO, a hierarchical solver that casts large-scale discrete OT as edge localization and uses dual potentials to guide both coarse-to-fine initialization and within-level refinement. Initialization propagates coarse dual potentials across a recursive subsampling hierarchy to assign candidate edges. Refinement then iteratively inserts the largest dual violators in each row and column until the relative KKT residual meets a prescribed tolerance, while budgeted pruning ensures linear memory complexity. For exact-arithmetic refinement, we prove finite termination at a global optimum under a symbolic lexicographic rule. At the million-point scale, HELLO attains lower transport objectives with order-of-magnitude runtime improvements over strong baselines across feature dimensions from single digits to thousands. It further scales to 1.28 million samples per marginal in 8192 dimensions on a single H100, using 41.6 GiB peak GPU memory while satisfying a full relative KKT residual below $10^{-6}$. Beyond standard discrete OT, the framework supports general pairwise costs and serves as a scalable balanced-OT oracle for semi-discrete OT, Gromov--Wasserstein, unbalanced OT, and OT-based Flow Matching.