Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action

TL;DR

Var-RUOT addresses the challenge of inferring continuous dynamics from sparse, high-dimensional snapshots by enforcing the RUOT minimum-action principle through a variational approach. It parameterizes a single scalar potential $\lambda(\mathbf{x},t)$ and derives the velocity $\mathbf{u}$ and growth $g$ from first-order optimality conditions, enabling a reduced, end-to-end model with potentially lower action and faster convergence. The method combines a weighted-particle loss with reconstruction, HJB, and action terms, and examines how different growth penalties $\psi(g)$ encode biological priors. Experimental results on synthetic three-gene dynamics, EMT data, and hematopoiesis datasets show Var-RUOT achieving competitive distribution fit while minimizing path action and improving training stability, suggesting practical value for trajectory inference in single-cell and other high-dimensional systems. Limitations include local convergence of neural nets, potential deterioration under modified metrics, and the need for automated selection of $\psi(g)$, offering clear directions for future work and broader impact in mesoscopic inference tasks.

Abstract

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ödinger 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.

Figures (14)

• Figure 1: Overview of Variational RUOT (Var-RUOT).
• Figure 2: Comparison between DeepRUOT and Var-RUOT on dynamics reconstruction on three-gene simulation dataset. a) The trajectory and growth rate learned by DeepRUOT. b) The trajectory and growth rate learned by Var-RUOT on the same dataset. In the figure, the circles and '+' signs denote the generated data and the original data points, respectively.
• Figure 3: Comparison between DeepRUOT and Var-RUOT on dynamics reconstruction on EMT dataset. a) The trajectory and growth rate learned by DeepRUOT. b) The trajectory and growth rate learned by Var-RUOT on the same dataset. In the figure, the circles and '+' signs denote the generated data and the original data points, respectively.
• Figure 4: Comparison of Var-RUOT using different growth metric on mouse blood hematopoiesis dataset. a) The trajectory and growth at time points $t=0,1,2$ learned using the standard WFR metric. b) The trajectory and growth at time points $t=0,1,2$ learned using the modified metric. In the figure, the circles and '+' signs denote the generated data and the original data points, respectively.
• Figure 5: Diagram of the Gaussian mixture datasets. PCA is used to reduce the high dimensional data to 2 dimensions.

Theorems & Definitions (20)

• Definition 3.1: Regularized Unbalanced Optimal Transport (RUOT) Problem
• Remark 3.1
• Definition 4.1: Isotropic Time-Invariant (ITI) RUOT Problem
• Theorem 4.1: Necessary Conditions for Achieving the Optimal Solution in the ITI-RUOT Problem
• Remark 4.1
• Remark 4.2
• Theorem 4.2: The relationship between $\mathbf{u}$ and $g$ ; Biological prior
• Theorem 5.1
• Remark 5.1
• Remark 5.2