Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

arXiv:2607.04738 · stat.ML, cs.AI, cs.LG · Submitted 2026-07-06 · Read on arXiv

Markus Heinonen, Yair Shenfeld, Ricardo Baptista, Daniel Waxman, Dmitry Batenkov, Tim Cooijmans, Eli Bingham

stat.ML, cs.AI, cs.LG

Submitted: 2026-07-06

Code: https://github.com/BasisResearch/wasserstein-residuals

License: http://creativecommons.org/licenses/by/4.0/

The gist: Reconstructing population dynamics is a central problem in the physical and data sciences.

Terminology

Abstract

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.

Sources

Related papers