# Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme

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

## 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 (most cited first, 30 of 120)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [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)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [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)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theory](https://mlsanity.com/paper/2110.11291/) (2021)
- [Training GANs with Optimism](https://mlsanity.com/paper/1711.00141/) (2017)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [Structured Inference Networks for Nonlinear State Space Models](https://mlsanity.com/paper/1609.09869/) (2016)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Solving Schr"odinger Bridges via Maximum Likelihood](https://mlsanity.com/paper/2106.02081/) (2021)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Guided Flows for Generative Modeling and Decision Making](https://mlsanity.com/paper/2311.13443/) (2023)

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

- [A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots](https://mlsanity.com/paper/2605.08550/) (2026)
- [Leveraging Gauge Freedom for Learning Non-Gradient Population Dynamics of Stochastic Systems](https://mlsanity.com/paper/2605.25107/) (2026)
- [Two-Parameter Flows for Learning Population Dynamics of Physical Systems](https://mlsanity.com/paper/2605.26285/) (2026)
- [Wasserstein Residuals: Learning Gradient Flows from Population Dynamics](https://mlsanity.com/paper/2607.04738/) (2026)
- [Multimarginal flow matching with optimal transport potentials](https://mlsanity.com/paper/2606.05327/) (2026)
- [Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization](https://mlsanity.com/paper/2606.27767/) (2026)
- [Light Entropic Optimal Transport on Riemannian Manifolds](https://mlsanity.com/paper/2610.03085/) (2026)
- [Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals](https://mlsanity.com/paper/2610.03679/) (2026)
- [Discrete Action Matching: Learning Stochastic Dynamics from Samples via State Graphs](https://mlsanity.com/paper/2610.05071/) (2026)

## 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.