mlsanityData to 14 Oct 2026

Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme

arXiv 2506.01502 ↗First posted 2 Jun 2025cs.LG cs.AI stat.MLcited by 9

Not on the map. Why?

Abstract (arXiv)

Learning population dynamics involves recovering the underlying process that governs particle evolution, given evolutionary snapshots of samples at discrete time points. Recent methods frame this as an energy minimization problem in probability space and leverage the celebrated JKO scheme for efficient time discretization. In this work, we introduce textttiJKOnet, an approach that combines the JKO framework with inverse optimization techniques to learn population dynamics. Our method relies on a conventional textitend-to-end adversarial training procedure and does not require restrictive architectural choices, e.g., input-convex neural networks. We establish theoretical guarantees for our methodology and demonstrate improved performance over prior JKO-based methods. The code of textttiJKOnet is available at https://github.com/MuXauJl11110/iJKOnet.

Key facts. Cited by 9 papers in the map (data through 14 Oct 2026). Builds on 120 of its 223 references that are in the map. First posted on arXiv 2 Jun 2025.

Builds on

120 of 223 references.

References · 120

  1. 2015Deep Residual Learning for Image Recognitioncited by 52,096
  2. 2014Adam: A Method for Stochastic Optimizationcited by 37,240
  3. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  4. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 9,122
  5. 2022Flow Matching for Generative Modelingcited by 5,702
  6. 2022Generative Adversarial Networkscited by 5,627
  7. 2018Neural Ordinary Differential Equationscited by 3,779
  8. 2022Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowcited by 3,700
  9. 2016Gaussian Error Linear Units (GELUs)cited by 3,502
  10. 2019Optuna: A Next-generation Hyperparameter Optimization Frameworkcited by 2,501
  11. 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,316
  12. 2017Improved Training of Wasserstein GANscited by 2,235
38 more
  1. 2018Spectral Normalization for Generative Adversarial Networkscited by 1,636
  2. 2022Building Normalizing Flows with Stochastic Interpolantscited by 995
  3. 2023Improving and generalizing flow-based generative models with minibatch optimal transportcited by 879
  4. 2021Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modelingcited by 602
  5. 2016Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithmcited by 507
  6. 2016Input Convex Neural Networkscited by 414
  7. 2020Scalable Gradients for Stochastic Differential Equationscited by 271
  8. 2023Flow Matching on General Geometriescited by 265
  9. 2020How to train your neural ODE: the world of Jacobian and kinetic regularizationcited by 228
  10. 2021Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theorycited by 219
  11. 2017Training GANs with Optimismcited by 217
  12. 2018Statistical Aspects of Wasserstein Distancescited by 203
  13. 2016Structured Inference Networks for Nonlinear State Space Modelscited by 182
  14. 2020TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamicscited by 160
  15. 2021Solving Schr"odinger Bridges via Maximum Likelihoodcited by 141
  16. 2022Neural Optimal Transportcited by 126
  17. 2022Optimal Transport Tools (OTT): A JAX Toolbox for all things Wassersteincited by 101
  18. 2023Guided Flows for Generative Modeling and Decision Makingcited by 101
  19. 2022Action Matching: Learning Stochastic Dynamics from Samplescited by 98
  20. 2019Wasserstein-2 Generative Networkscited by 97
  21. 2021Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmarkcited by 96
  22. 2021Proximal Optimal Transport Modeling of Population Dynamicscited by 93
  23. 2020Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimizationcited by 90
  24. 2022Variational inference via Wasserstein gradient flowscited by 85
  25. 2021Generative Modeling with Optimal Transport Mapscited by 82
  26. 2024ClimODE: Climate and Weather Forecasting with Physics-informed Neural ODEscited by 82
  27. 2023Diffusion Bridge Mixture Transports, Schr"odinger Bridge Problems and Generative Modelingcited by 80
  28. 2024Optimal Flow Matching: Learning Straight Trajectories in Just One Stepcited by 78
  29. 2021Large-Scale Wasserstein Gradient Flowscited by 77
  30. 2023Non-Denoising Forward-Time Diffusionscited by 77
  31. 2024Metric Flow Matching for Smooth Interpolations on the Data Manifoldcited by 77
  32. 2017Continuously Differentiable Exponential Linear Unitscited by 73
  33. 2022Manifold Interpolating Optimal-Transport Flows for Trajectory Inferencecited by 64
  34. 2022Supervised Training of Conditional Monge Mapscited by 61
  35. 2021Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networkscited by 53
  36. 2021Rates of Estimation of Optimal Transport Maps using Plug-in Estimators via Barycentric Projectionscited by 53
  37. 2021Variational Wasserstein gradient flowcited by 53
  38. 2021Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimizationcited by 51

50 most cited of 120 All 120 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 9.

9 papers.

  1. 2026A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshotscited by 5
  2. 2026Leveraging Gauge Freedom for Learning Non-Gradient Population Dynamics of Stochastic Systemscited by 1
  3. 2026Two-Parameter Flows for Learning Population Dynamics of Physical Systemscited by 1
  4. 2026Wasserstein Residuals: Learning Gradient Flows from Population Dynamicscited by 1
  5. 2026Multimarginal flow matching with optimal transport potentialscited by 0
  6. 2026Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimizationcited by 0
  7. 2026Light Entropic Optimal Transport on Riemannian Manifoldscited by 0
  8. 2026Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residualscited by 0
  9. 2026Discrete Action Matching: Learning Stochastic Dynamics from Samples via State Graphscited by 0