Computing Wasserstein-$p$ Distance Between Images with Linear Cost
Yidong Chen, Chen Li, Zhonghua Lu · 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) · 2022
When the images are formulated as discrete measures, computing Wasserstein-p distance between them is challenging due to the complexity of solving the corresponding Kantorovich's problem. In this paper, we propose a novel algorithm to compute the Wasserstein-p distance between discrete measures by restricting the optimal transport (OT) problem on a subset. First, we define the restricted OT problem and prove the solution of the restricted problem converges to Kantorovich's OT solution. Second, we propose the SparseSinkhorn algorithm for the restricted problem and provide a multi-scale algorithm to estimate the subset. Finally, we implement the proposed algorithm on CUDA and illustrate the linear computational cost in terms of time and memory requirements. We compute Wasserstein-p distance, estimate the transport mapping, and transfer color between color images with size ranges from$64\times 64$to$1920\times 1200$. (Our code is available at https://github.com/ucascnic/CudaOT)