The stability and statistic of domain decomposition algorithm with mini-batch learning for optimal transport

Abstract

The Wasserstein distance in the optimal transport (OT) model is being widely used in machine learning applications due to its mathematical and geometric properties. Since the entropy regularized OT is a convex problem, and for its large-scale problems, many algorithms have been proposed. Existing methods mainly adopt iterative random projection techniques to approximate the Wasserstein distance in the OT model. In this paper, we employ the entropic domain decomposition algorithm with mini-batch learning to calculate the entropy regularized OT problem and investigate the stability of this algorithm. The algorithm performs on different cells of the domain instead of using the projection technique and it is based on a mini-batch loss function. The corresponding solution in each iteration is derived from the dual problem of entropy regularized OT on each decomposed cell. Once the solutions to the dual problem on all decomposed cells are obtained, the corresponding transport plan in each iteration can be obtained. We show that the transport plan computed by its dual variables is admissible, and we also give the sub-gradient approximation result. In addition, we exploit the statistical concentration result of the subsampled mini-batch Wasserstein distance on decomposed cells. Finally, the numerical experiments validate that the proposed algorithm can save computational time further.