# Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport

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

## Key facts

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

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (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)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [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)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theory](https://mlsanity.com/paper/2110.11291/) (2021)
- [Denoising Diffusion Bridge Models](https://mlsanity.com/paper/2309.16948/) (2023)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Deep Generative Learning via Schr"odinger Bridge](https://mlsanity.com/paper/2106.10410/) (2021)
- [On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpoint](https://mlsanity.com/paper/1412.4430/) (2014)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems](https://mlsanity.com/paper/1912.01825/) (2019)
- [Manifold Interpolating Optimal-Transport Flows for Trajectory Inference](https://mlsanity.com/paper/2206.14928/) (2022)
- [The Schr"odinger Bridge between Gaussian Measures has a Closed Form](https://mlsanity.com/paper/2202.05722/) (2022)
- [A Computational Framework for Solving Wasserstein Lagrangian Flows](https://mlsanity.com/paper/2310.10649/) (2023)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Trajectory Inference via Mean-field Langevin in Path Space](https://mlsanity.com/paper/2205.07146/) (2022)
- [Convergence of Continuous Normalizing Flows for Learning Probability Distributions](https://mlsanity.com/paper/2404.00551/) (2024)
- [Convergence Analysis of Flow Matching in Latent Space with Transformers](https://mlsanity.com/paper/2404.02538/) (2024)
- [Neural Optimal Transport with Lagrangian Costs](https://mlsanity.com/paper/2406.00288/) (2024)
- [Unnormalized Optimal Transport](https://mlsanity.com/paper/1902.03367/) (2019)
- [The most likely evolution of diffusing and vanishing particles: Schrodinger Bridges with unbalanced marginals](https://mlsanity.com/paper/2108.02879/) (2021)
- [Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior](https://mlsanity.com/paper/2406.07475/) (2024)
- [Learning High Dimensional Wasserstein Geodesics](https://mlsanity.com/paper/2102.02992/) (2021)

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

- [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)
- [Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schr"odinger Bridge](https://mlsanity.com/paper/2505.11197/) (2025)
- [Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action](https://mlsanity.com/paper/2505.11823/) (2025)
- [Momentum Multi-Marginal Schr"odinger Bridge Matching](https://mlsanity.com/paper/2506.10168/) (2025)
- [Integrating Dynamical Systems Modeling with Spatiotemporal scRNA-seq Data Analysis](https://mlsanity.com/paper/2503.11347/) (2025)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Modeling Microenvironment Trajectories on Spatial Transcriptomics with NicheFlow](https://mlsanity.com/paper/2511.00977/) (2025)
- [WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport](https://mlsanity.com/paper/2601.06810/) (2026)
- [Half-order Fine-Tuning for Diffusion Model: A Recursive Likelihood Ratio Optimizer](https://mlsanity.com/paper/2502.00639/) (2025)
- [Branched Schr"odinger Bridge Matching](https://mlsanity.com/paper/2506.09007/) (2025)
- [Inferring stochastic dynamics with growth from cross-sectional data](https://mlsanity.com/paper/2505.13197/) (2025)
- [CellStream: Dynamical Optimal Transport Informed Embeddings for Reconstructing Cellular Trajectories from Snapshots Data](https://mlsanity.com/paper/2511.13786/) (2025)
- [Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots](https://mlsanity.com/paper/2605.00545/) (2026)
- [Equilibrium flow: From Snapshots to Dynamics](https://mlsanity.com/paper/2509.17990/) (2025)
- [Learning Explicit Single-Cell Dynamics Using ODE Representations](https://mlsanity.com/paper/2510.02903/) (2025)
- [Entangled Schr"odinger Bridge Matching](https://mlsanity.com/paper/2511.07406/) (2025)
- [WFR-MFM: One-Step Inference for Dynamic Unbalanced Optimal Transport](https://mlsanity.com/paper/2601.20606/) (2026)
- [Leveraging Gauge Freedom for Learning Non-Gradient Population Dynamics of Stochastic Systems](https://mlsanity.com/paper/2605.25107/) (2026)
- [TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics](https://mlsanity.com/paper/2608.21070/) (2026)
- [Private Continuous-Time Synthetic Trajectory Generation via Mean-Field Langevin Dynamics](https://mlsanity.com/paper/2506.12203/) (2025)
- [MODE: Learning compositional representations of complex systems with Mixtures Of Dynamical Experts](https://mlsanity.com/paper/2510.09594/) (2025)
- [On the Unique Recovery of Transport Maps and Vector Fields from Finite Measure-Valued Data](https://mlsanity.com/paper/2604.07671/) (2026)
- [Multiscale Supervised Unbalanced Optimal Transport Flow Matching](https://mlsanity.com/paper/2605.16529/) (2026)
- [Modeling Temporal scRNA-seq Data with Latent Gaussian Process and Optimal Transport](https://mlsanity.com/paper/2605.20989/) (2026)
- [Multimarginal flow matching with optimal transport potentials](https://mlsanity.com/paper/2606.05327/) (2026)
- [Order Is Not Control: Driven-Dissipative Response Laws Across Artificial and Biological Systems](https://mlsanity.com/paper/2606.12923/) (2026)
- [Trajectory Inference of Human Aging from Cross-Sectional DNA Methylation Data](https://mlsanity.com/paper/2607.06583/) (2026)
- [Simulation-free Unbalanced Dynamic Optimal Transport with General Growth Penalty](https://mlsanity.com/paper/2609.04710/) (2026)
- [Dynamic Generalized Gromov-Wasserstein Optimal Transport](https://mlsanity.com/paper/2609.20008/) (2026)

## Abstract (arXiv)

Reconstructing dynamics using samples from sparsely time-resolved snapshots is an important problem in both natural sciences and machine learning. Here, we introduce a new deep learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dynamics from observed snapshots. Based on the RUOT form, our method models these dynamics without requiring prior knowledge of growth and death processes or additional information, allowing them to be learned directly from data. Theoretically, we explore the connections between the RUOT and Schr"odinger bridge problem and discuss the key challenges and potential solutions. The effectiveness of our method is demonstrated with a synthetic gene regulatory network, high-dimensional Gaussian Mixture Model, and single-cell RNA-seq data from blood development. Compared with other methods, our approach accurately identifies growth and transition patterns, eliminates false transitions, and constructs the Waddington developmental landscape. Our code is available at: https://github.com/zhenyiizhang/DeepRUOT.