# Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

arXiv: https://arxiv.org/abs/2607.04738 · HTML: https://mlsanity.com/paper/2607.04738/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 57 references that are in the map.
- First posted on arXiv 6 Jul 2026.

## Builds on (most cited first, 30 of 33)

- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theory](https://mlsanity.com/paper/2110.11291/) (2021)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Diffusion Schr"odinger Bridge Matching](https://mlsanity.com/paper/2303.16852/) (2023)
- [Solving Schr"odinger Bridges via Maximum Likelihood](https://mlsanity.com/paper/2106.02081/) (2021)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Efficient and Accurate Gradients for Neural SDEs](https://mlsanity.com/paper/2105.13493/) (2021)
- [Neural Conservation Laws: A Divergence-Free Perspective](https://mlsanity.com/paper/2210.01741/) (2022)
- [A Computational Framework for Solving Wasserstein Lagrangian Flows](https://mlsanity.com/paper/2310.10649/) (2023)
- [Deep Momentum Multi-Marginal Schr"odinger Bridge](https://mlsanity.com/paper/2303.01751/) (2023)
- [Multi-marginal Schr"odinger Bridges with Iterative Reference Refinement](https://mlsanity.com/paper/2408.06277/) (2024)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling](https://mlsanity.com/paper/2505.13413/) (2025)
- [Curly Flow Matching for Learning Non-gradient Field Dynamics](https://mlsanity.com/paper/2510.26645/) (2025)
- [Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots](https://mlsanity.com/paper/2410.22729/) (2024)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Gradient-flow SDEs have unique transient population dynamics](https://mlsanity.com/paper/2505.21770/) (2025)

## Built on by (most cited first, 1 of 1)

- [Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals](https://mlsanity.com/paper/2610.03679/) (2026)

## Abstract (arXiv)

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.