# Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 74 references that are in the map.
- First posted on arXiv 17 May 2025.

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (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)
- [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)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [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)
- [Diffusion Schr"odinger Bridge Matching](https://mlsanity.com/paper/2303.16852/) (2023)
- [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)
- [Simulation-free Schr"odinger bridges via score and flow matching](https://mlsanity.com/paper/2307.03672/) (2023)
- [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)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation](https://mlsanity.com/paper/2311.15100/) (2023)
- [Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport](https://mlsanity.com/paper/2410.00844/) (2024)
- [Trajectory Inference via Mean-field Langevin in Path Space](https://mlsanity.com/paper/2205.07146/) (2022)
- [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)
- [Multi-Modal and Multi-Attribute Generation of Single Cells with CFGen](https://mlsanity.com/paper/2407.11734/) (2024)
- [Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schr"odinger Bridge](https://mlsanity.com/paper/2505.11197/) (2025)

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

- [Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling](https://mlsanity.com/paper/2505.13413/) (2025)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport](https://mlsanity.com/paper/2601.06810/) (2026)
- [A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots](https://mlsanity.com/paper/2605.08550/) (2026)
- [CellStream: Dynamical Optimal Transport Informed Embeddings for Reconstructing Cellular Trajectories from Snapshots Data](https://mlsanity.com/paper/2511.13786/) (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)
- [Multiscale Supervised Unbalanced Optimal Transport Flow Matching](https://mlsanity.com/paper/2605.16529/) (2026)
- [Simulation-free Unbalanced Dynamic Optimal Transport with General Growth Penalty](https://mlsanity.com/paper/2609.04710/) (2026)
- [UniMergeTrack: Hierarchical Model Merging for General RGB-X Visual Object Tracking](https://mlsanity.com/paper/2610.13053/) (2026)

## Abstract (arXiv)

Recovering the dynamics from a few snapshots of a high-dimensional system is a challenging task in statistical physics and machine learning, with important applications in computational biology. Many algorithms have been developed to tackle this problem, based on frameworks such as optimal transport and the Schr"odinger bridge. A notable recent framework is Regularized Unbalanced Optimal Transport (RUOT), which integrates both stochastic dynamics and unnormalized distributions. However, since many existing methods do not explicitly enforce optimality conditions, their solutions often struggle to satisfy the principle of least action and meet challenges to converge in a stable and reliable way. To address these issues, we propose Variational RUOT (Var-RUOT), a new framework to solve the RUOT problem. By incorporating the optimal necessary conditions for the RUOT problem into both the parameterization of the search space and the loss function design, Var-RUOT only needs to learn a scalar field to solve the RUOT problem and can search for solutions with lower action. We also examined the challenge of selecting a growth penalty function in the widely used Wasserstein-Fisher-Rao metric and proposed a solution that better aligns with biological priors in Var-RUOT. We validated the effectiveness of Var-RUOT on both simulated data and real single-cell datasets. Compared with existing algorithms, Var-RUOT can find solutions with lower action while exhibiting faster convergence and improved training stability. Our code is available at https://github.com/ZerooVector/VarRUOT.