A Bidirectional DeepParticle Method for Efficiently Solving Low-dimensional Transport Map Problems

Abstract

This paper aims to efficiently compute transport maps between probability distributions arising from particle representation of bio-physical problems. We develop a bidirectional DeepParticle (BDP) method to learn and generate solutions under varying physical parameters. Solutions are approximated as empirical measures of particles that adaptively align with the high-gradient regions. The core idea of the BDP method is to learn both forward and reverse mappings (between the uniform and a non-trivial target distribution) by minimizing the discrete 2-Wasserstein distance (W2) and optimizing the transition map therein by a minibatch technique. We present numerical results demonstrating the effectiveness of the BDP method for learning and generating solutions to Keller-Segel chemotaxis systems in the presence of laminar flows and Kolmogorov flows with chaotic streamlines in three space dimensions. The BDP outperforms two recent representative single-step flow matching and diffusion models (rectified flow and shortcut diffusion models) in the generative AI literature. However when the target distribution is high-dimensional (4 and above), e.g. a mixture of two Gaussians, the single-step diffusion models scale better in dimensions than BDP in terms of W2-accuracy.