Wasserstein Residuals: Learning Gradient Flows from Population Dynamics
2026-07-06 • Artificial Intelligence
Artificial IntelligenceMachine Learning
AI summaryⓘ
The authors study ways to track how groups or populations change over time, using a mathematical approach called Wasserstein gradient flow (WGF). Traditional methods to do this are slow and inflexible because they rely on solving complex transport problems step-by-step. The authors propose a new method called stitching that uses a different way to measure errors and fits data globally, making it faster and able to handle bigger gaps in observations. Their method performs very well compared to others in tests. They also provide their code online for others to use.
Wasserstein gradient flowJordan-Kinderlehrer-Otto (JKO) schemeoptimal transportcontinuity equationresidual approachparticle-based methodtrajectory inferenceenergy functionaldata-fitting divergence
Authors
Markus Heinonen, Yair Shenfeld, Ricardo Baptista, Daniel Waxman, Dmitry Batenkov, Tim Cooijmans, Eli Bingham
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.