Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

2026-07-06

AI Summary

This work addresses the reconstruction of population dynamics governed by Wasserstein gradient flows from sparse observational data. The authors propose a particle-based "stitching" method that bypasses the traditional Jordan–Kinderlehrer–Otto (JKO) time discretization, instead enforcing the continuity equation via a non-negative residual loss and integrating a data-fidelity divergence into a unified optimization objective. By eliminating the need for costly optimal transport computations and fixed time steps, the approach is simulation-agnostic and robust to irregular or widely spaced observation intervals. Evaluated on multiple trajectory inference benchmarks, the method achieves state-of-the-art performance, demonstrating particular superiority in regimes with highly sparse observations or large temporal gaps between measurements.

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.

Problem

• population dynamics
• Wasserstein gradient flow
• trajectory inference
• optimal transport
• continuity equation

Innovation

• Wasserstein gradient flow
• residual learning
• continuity equation
• particle-based method
• trajectory inference