Wasserstein Residuals: Learning Gradient Flows from Population Dynamics
This work provides a more flexible and robust method for reconstructing population dynamics, benefiting researchers in physical and data sciences who need to infer continuous-time dynamics from discrete observations.
The authors propose a residual-based approach to learn Wasserstein gradient flows from population dynamics, introducing a simulation-free particle method (stitching) that achieves state-of-the-art performance on trajectory inference benchmarks.
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 the continuity equations 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 to a new particle-based method, stitching, that is simulation-free and robust to large gaps between observations. We demonstrate that the stitching method achieves state-of-the-art performance across trajectory inference benchmarks. For code see github.com/BasisResearch/wasserstein-residuals.