A Neural Network Framework for High-Dimensional Dynamic Unbalanced Optimal Transport

Abstract: In this paper, we introduce a neural network-based method to address the high-dimensional dynamic unbalanced optimal transport (UOT) problem. Dynamic UOT focuses on the optimal transportation between two densities with unequal total mass, however, it introduces additional complexities compared to the traditional dynamic optimal transport (OT) problem. To efficiently solve the dynamic UOT problem in high-dimensional space, we first relax the original problem by using the generalized Kullback-Leibler (GKL) divergence to constrain the terminal density. Next, we adopt the Lagrangian discretization to address the unbalanced continuity equation and apply the Monte Carlo method to approximate the high-dimensional spatial integrals. Moreover, a carefully designed neural network is introduced for modeling the velocity field and source function. Numerous experiments demonstrate that the proposed framework performs excellently in high-dimensional cases. Additionally, this method can be easily extended to more general applications, such as crowd motion problem.