Wasserstein residuals: Learning Gradient Flows from Population Dynamics

Reconstructing population dynamics is a central problem in the physical and data sciences. Often, the dynamics are modeled as a Wasserstein gradient flow (WGF): a curve of distributions driven by an energy functional. Though there are multiple mathematical characterizations of a WGF, the dominant algorithmic approach relies on the Jordan--Kinderlehrer--Otto (JKO) scheme. JKO-based methods are inflexible to time discretisation and require solving costly optimal transport problems. We take a residual approach, enforcing that the continuity equation conditions are satisfied via a non-negative loss function, whose minimum is the WGF. Combined with a data-fitting divergence, this gives a single global objective. This perspective unifies several existing methods, and leads us to discover a new particle-based method, \emph{stitching}, that is simulation-free and robust to large gaps in observation times of the dynamics. We demonstrate that the stitching method achieves state-of-the-art performance across trajectory inference benchmarks.