Title: A Memory-Efficient Hierarchical Algorithm for Large-scale Optimal Transport Problems
Primary area: optimization
Keywords: optimal transport, linear programming, multiscale framework, first-order methods

Abstract: In this paper we propose a memory-efficient hierarchical algorithm for solving large-scale optimal transport (OT) problems where the supports are in $\mathbb{R}^2$ with squared Euclidean cost. 
The key 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 which can further reduce computational complexity. Theoretically, we establish a scale-independent iteration-complexity upper bound for the refinement phase, which is consistent with our numerical results. Numerically, extensive experiments on image dataset DOTmark demonstrate that the proposed algorithm effectively addresses the memory and scalability bottlenecks of existing solvers. 
For images with resolution $1024\times 1024$, our method achieves an $8.9\times$ speedup and reduces GPU memory usage by $70.5$% compared to state-of-the-art baselines, under comparable accuracy.