A Memory-Efficient Hierarchical Algorithm for Large-scale Optimal Transport Problems

Abstract

In this paper we propose a memory-efficient hierarchical algorithm for solving large-scale optimal transport (OT) problems with squared Euclidean cost. The core of our proposed approach is the combination of multiscale hierarchical representation of the OT problem and a GPU-implemented Primal-Dual Hybrid Gradient (PDHG) method. Moreover, an active pruning technique is applied to further reduce computational complexity. Theoretically, we establish a scale-independent iteration-complexity upper bound for the refinement phase, which is consistent with our numerical observations. Numerically, experiments on image dataset DOTmark and point cloud dataset ModelNet10 demonstrate that the proposed algorithm effectively addresses the memory and scalability bottlenecks. Compared to state-of-the-art baselines, our method demonstrates significant advantages: for images with $n=1024^2$ pixels, it achieves an $8.9\times$ speedup and $70.5$\% reduction in memory usage under comparable accuracy; for 3D point clouds at scale $n=2^{18}$, it achieves a $1.84\times$ speedup and an $83.2$\% reduction in memory usage with $24.9$\% lower transport cost.