Special Session 32: Inverse Problems and Image Processing

Hierarchical Exact Solver for Large-scale Optimal Transport and its applciations
Xiaoqun Zhang
Shanghai Jiao Tong University
Peoples Rep of China
Co-Author(s):    Wenzhou Xia, Jingwei Liang, Qiaoqiao Ding
Abstract:
Optimal transport (OT) underlies many machine-learning and computer-vision tasks, yet solving large-scale OT problems under a tight memory budget remains challenging because of an inherent trilemma among precision, memory efficiency and dimensional scalability. To resolve this trade-off we propose a memory-efficient, dimension-scalable hierarchical exact solver for large-scale OT with low-rank cost. The proposed algorithm is a multi-resolution decomposition of the OT problem coupled with parallel-friendly linear-programming solvers. To guarantee memory efficiency we strictly bound the storage to linear complexity via active support pruning. Theoretically, we prove a scale-independent iteration-complexity upper bound for the refinement phase (consistent with the empirical observation) and show that, under standard assumptions, the algorithm converges to the global optimum. Numerical experiments demonstrate the scalability and exactness of the proposed algorithm across a wide range of regimes. The algorithm is capable to handle million-scale problems in very high dimensions (n=106, d=8192) on a single GPU, while still delivering exact solutions on synthetic high-dimensional data. The method thus alleviates the memory and scalability bottlenecks of existing solvers and resolves the precision-memory-dimensionality trilemma for large-scale OT.