Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

• URL: http://arxiv.org/abs/2607.04738v1
• Date: Mon, 06 Jul 2026 07:20:30 GMT
• Title: Wasserstein Residuals: Learning Gradient Flows from Population Dynamics
• Abstract summary: Reconstructing population dynamics is a central problem in the physical and data sciences.<n>The dominant algorithmic approach relies on the Jordan--Kinderlehrer--Otto (JKO) scheme.<n>We take a residual approach, enforcing the continuity equations via a non-negative loss function whose minimum is the WGF.<n>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.
• Score: 18.263762660153674
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Generative Path-Finding Method for Wasserstein Gradient Flow [4.372661069805067]
  We propose a generative path finding framework for Wasserstein gradient paths.<n>GenWGP learns a generative flow that transports mass from an initial density to an unknown equilibrium distribution.<n>Experiments on Fokker Planck and aggregation type problems show that GenWGP matches or exceeds high fidelity reference solutions.
  arXiv  Detail & Related papers  (2026-04-13T14:22:20Z)
• Learning Discrete Diffusion of Graphs via Free-Energy Gradient Flows [24.6419207456082]
  A parallel theoretical framework based on gradient flows has remained elusive due to intrinsic challenges in translating the $W$ distance directly into these settings.<n>In this work, we propose the first computational approach addressing these challenges, leveraging an appropriate metric $W_K$ on the simplex of probability distributions.<n>We introduce a novel methodology for learning diffusion dynamics over discrete spaces, which recovers the underlying functional directly by leveraging first-order optimality conditions.
  arXiv  Detail & Related papers  (2026-04-13T11:15:14Z)
• Curly Flow Matching for Learning Non-gradient Field Dynamics [49.480209466896035]
  We introduce Curly Flow Matching (Curly-FM), a novel approach to learning non-gradient field dynamics.<n>Curly-FM is capable of learning non-gradient field dynamics by designing and solving a Schr"odinger bridge problem.<n>Curly-FM can learn trajectories that better match both the reference process and population marginals.
  arXiv  Detail & Related papers  (2025-10-30T16:11:39Z)
• On the Wasserstein Convergence and Straightness of Rectified Flow [54.580605276017096]
  Rectified Flow (RF) is a generative model that aims to learn straight flow trajectories from noise to data.<n>We provide a theoretical analysis of the Wasserstein distance between the sampling distribution of RF and the target distribution.<n>We present general conditions guaranteeing uniqueness and straightness of 1-RF, which is in line with previous empirical findings.
  arXiv  Detail & Related papers  (2024-10-19T02:36:11Z)
• Score-based Diffusion Models in Function Space [137.70916238028306]
  Diffusion models have recently emerged as a powerful framework for generative modeling.<n>This work introduces a mathematically rigorous framework called Denoising Diffusion Operators (DDOs) for training diffusion models in function space.<n>We show that the corresponding discretized algorithm generates accurate samples at a fixed cost independent of the data resolution.
  arXiv  Detail & Related papers  (2023-02-14T23:50:53Z)
• Gradient flows and randomised thresholding: sparse inversion and classification [0.0]
  Sparse inversion and classification problems are ubiquitous in modern data science and imaging. In classification, we consider, e.g., the sum of a data fidelity term and a non-smooth Ginzburg--Landau energy. Standard (sub)gradient descent methods have shown to be inefficient when approaching such problems.
  arXiv  Detail & Related papers  (2022-03-22T09:21:14Z)
• Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations [57.15855198512551]
  We propose a novel score-based generative model for graphs with a continuous-time framework. We show that our method is able to generate molecules that lie close to the training distribution yet do not violate the chemical valency rule.
  arXiv  Detail & Related papers  (2022-02-05T08:21:04Z)
• Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation [110.09855163856326]
  This paper is about the problem of learning a policy for generating an object from a sequence of actions. We propose GFlowNet, based on a view of the generative process as a flow network. We prove that any global minimum of the proposed objectives yields a policy which samples from the desired distribution.
  arXiv  Detail & Related papers  (2021-06-08T14:21:10Z)
• Large-Scale Wasserstein Gradient Flows [84.73670288608025]
  We introduce a scalable scheme to approximate Wasserstein gradient flows. Our approach relies on input neural networks (ICNNs) to discretize the JKO steps. As a result, we can sample from the measure at each step of the gradient diffusion and compute its density.
  arXiv  Detail & Related papers  (2021-06-01T19:21:48Z)
• The Wasserstein Proximal Gradient Algorithm [23.143814848127295]
  Wasserstein gradient flows are continuous time dynamics that define curves of steepest descent to minimize an objective function over the space of probability measures. We propose a Forward Backward (FB) discretization scheme that can tackle the case where the objective function is the sum of a smooth and a nonsmooth geodesically convex terms.
  arXiv  Detail & Related papers  (2020-02-07T22:19:32Z)
• A Near-Optimal Gradient Flow for Learning Neural Energy-Based Models [93.24030378630175]
  We propose a novel numerical scheme to optimize the gradient flows for learning energy-based models (EBMs) We derive a second-order Wasserstein gradient flow of the global relative entropy from Fokker-Planck equation. Compared with existing schemes, Wasserstein gradient flow is a smoother and near-optimal numerical scheme to approximate real data densities.
  arXiv  Detail & Related papers  (2019-10-31T02:26:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.