# Proximal Optimal Transport Modeling of Population Dynamics

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

## Key facts

- Cited by 93 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 69 references that are in the map.
- First posted on arXiv 11 Jun 2021.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Lipschitz regularity of deep neural networks: analysis and efficient estimation](https://mlsanity.com/paper/1805.10965/) (2018)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Unrolled Generative Adversarial Networks](https://mlsanity.com/paper/1611.02163/) (2016)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Structured Inference Networks for Nonlinear State Space Models](https://mlsanity.com/paper/1609.09869/) (2016)
- [Multiple Object Tracking: A Literature Review](https://mlsanity.com/paper/1409.7618/) (2014)
- [On Wasserstein Two Sample Testing and Related Families of Nonparametric Tests](https://mlsanity.com/paper/1509.02237/) (2015)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Improving GANs Using Optimal Transport](https://mlsanity.com/paper/1803.05573/) (2018)
- [Towards a Neural Statistician](https://mlsanity.com/paper/1606.02185/) (2016)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Optimal Control Via Neural Networks: A Convex Approach](https://mlsanity.com/paper/1805.11835/) (2018)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Scalable Unbalanced Optimal Transport using Generative Adversarial Networks](https://mlsanity.com/paper/1810.11447/) (2018)
- [Regularity as Regularization: Smooth and Strongly Convex Brenier Potentials in Optimal Transport](https://mlsanity.com/paper/1905.10812/) (2019)
- [Stochastic Deep Networks](https://mlsanity.com/paper/1811.07429/) (2018)
- [Multi-marginal optimal transport and probabilistic graphical models](https://mlsanity.com/paper/2006.14113/) (2020)
- [Estimating ensemble flows on a hidden Markov chain](https://mlsanity.com/paper/1905.09119/) (2019)

## Built on by (most cited first, 30 of 93)

- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Diffusion Schr"odinger Bridge Matching](https://mlsanity.com/paper/2303.16852/) (2023)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Simulation-free Schr"odinger bridges via score and flow matching](https://mlsanity.com/paper/2307.03672/) (2023)
- [Metric Flow Matching for Smooth Interpolations on the Data Manifold](https://mlsanity.com/paper/2405.14780/) (2024)
- [Manifold Interpolating Optimal-Transport Flows for Trajectory Inference](https://mlsanity.com/paper/2206.14928/) (2022)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [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)
- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Aligned Diffusion Schr"odinger Bridges](https://mlsanity.com/paper/2302.11419/) (2023)
- [Generative Sliced MMD Flows with Riesz Kernels](https://mlsanity.com/paper/2305.11463/) (2023)
- [Neural Optimal Transport with General Cost Functionals](https://mlsanity.com/paper/2205.15403/) (2022)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [Kernel Neural Optimal Transport](https://mlsanity.com/paper/2205.15269/) (2022)
- [Trajectory Inference via Mean-field Langevin in Path Space](https://mlsanity.com/paper/2205.07146/) (2022)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Multi-marginal Schr"odinger Bridges with Iterative Reference Refinement](https://mlsanity.com/paper/2408.06277/) (2024)
- [Dynamic Conditional Optimal Transport through Simulation-Free Flows](https://mlsanity.com/paper/2404.04240/) (2024)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Efficient Gradient Flows in Sliced-Wasserstein Space](https://mlsanity.com/paper/2110.10972/) (2021)
- [Riemannian Metric Learning via Optimal Transport](https://mlsanity.com/paper/2205.09244/) (2022)
- [Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling](https://mlsanity.com/paper/2505.13413/) (2025)
- [Causal Modeling with Stationary Diffusions](https://mlsanity.com/paper/2310.17405/) (2023)
- [Light Schr"odinger Bridge](https://mlsanity.com/paper/2310.01174/) (2023)
- [Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings](https://mlsanity.com/paper/2209.15621/) (2022)
- [Taming Hyperparameter Tuning in Continuous Normalizing Flows Using the JKO Scheme](https://mlsanity.com/paper/2211.16757/) (2022)

## Abstract (arXiv)

We propose a new approach to model the collective dynamics of a population of particles evolving with time. As is often the case in challenging scientific applications, notably single-cell genomics, measuring features for these particles requires destroying them. As a result, the population can only be monitored with periodic snapshots, obtained by sampling a few particles that are sacrificed in exchange for measurements. Given only access to these snapshots, can we reconstruct likely individual trajectories for all other particles? We propose to model these trajectories as collective realizations of a causal Jordan-Kinderlehrer-Otto (JKO) flow of measures: The JKO scheme posits that the new configuration taken by a population at time t+1 is one that trades off an improvement, in the sense that it decreases an energy, while remaining close (in Wasserstein distance) to the previous configuration observed at t. In order to learn such an energy using only snapshots, we propose JKOnet, a neural architecture that computes (in end-to-end differentiable fashion) the JKO flow given a parametric energy and initial configuration of points. We demonstrate the good performance and robustness of the JKOnet fitting procedure, compared to a more direct forward method.