MLOT: Generalizing the Bipartite Structure to a Multi-Layered Framework for Optimal Transport

Liangliang Shi, Rui Tao, Junchi Yan

Keywords: Optimal Transport, Entropic Regularization, Sinkhorn Algorithm, Data Augmentation

Abstract: Although optimal transport (OT) has achieved significant success and widespread application in various fields, its structure remains relatively simple, relying on bipartite graphs with only two layers of nodes for transportation. In this paper, we propose a multi-layer optimal transport (MLOT) method that extends the original two-layer structure to handle transportation problems across multiple hierarchical levels, making it more adaptable to the complex structures found in deep learning tasks. In this framework, the source distribution flows through intermediate layers before reaching the target distribution, where estimating the intermediate distributions becomes crucial for solving the MLOT. Under entropic regularization, we further propose the MLOT-Sinkhorn algorithm to solve the multi-layer OT problem, where intermediate distributions can be estimated through the transportation calculations between adjacent layers. This algorithm can be accelerated using GPUs and significantly outperforms general solvers such as Gurobi. We also present theoretical results for the entropic MLOT, demonstrating its efficiency advantages and convergence properties. Furthermore, we find that our MLOT is well-suited for machine learning tasks based on data augmentation. As a result, we apply the MLOT-Sinkhorn algorithm to tasks such as text-image retrieval and visual graph matching. Experimental results show that reformulating these problems within the MLOT framework leads to significant improvements in performance.

Primary Area: optimization

Submission Number: 16142