A translated OT problem with distributions of different sizes for GPU Acceleration

Abstract

Widely used as a tool for comparing probability distributions, the optimal transport (OT) theory is very important in many machine learning tasks. Sinkhorn's algorithm successfully reduces its compuational cost from a cubic complexity to a quadratic one. Nevertheless, popular approaches of distribution comparison with OT on feature sets of different sizes could not support GPU parallelization. In order to overcome this difficulty, we propose the basis optimal transport which provides a translated OT problem with distributions of fixed sizes. Futhermore, we propose a deep dictionary learning framework for translating a given OT problem into our proposed basis optimal transport problem to make it solvable with GPU-based Sinkhorn's algorithm. A great reduction of computational time cost is reported according to our numerical experiments for computing the Wasserstein distance on datasets with size-variable distributions.