Parameterized Wasserstein Hamiltonian Flow

SIAM Journal on Numerical Analysis, Volume 63, Issue 1, Page 360-395, February 2025.
Abstract. In this work, we propose a numerical method to compute the Wasserstein Hamiltonian flow (WHF), which is a Hamiltonian system on the probability density manifold. Many well-known PDE systems can be reformulated as WHFs. We use the parameterized function as a push-forward map to characterize the solution of WHF, and convert the PDE to a finite-dimensional ODE system, which is a Hamiltonian system in the phase space of the parameter manifold. We establish theoretical error bounds for the continuous time approximation scheme in the Wasserstein metric. For the numerical implementation, neural networks are used as push-forward maps. We design an effective symplectic scheme to solve the derived Hamiltonian ODE system so that the method preserves some important quantities such as Hamiltonian. The computation is done by a fully deterministic symplectic integrator without any neural network training. Thus, our method does not involve direct optimization over network parameters and hence can avoid errors introduced by the stochastic gradient descent or similar methods, which are usually hard to quantify and measure in practice. The proposed algorithm is a sampling-based approach that scales well to higher dimensional problems. In addition, the method also provides an alternative connection between the Lagrangian and Eulerian perspectives of the original WHF through the parameterized ODE dynamics.