Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action

TL;DR

This paper introduces Variational RUOT, a novel framework that enforces optimality conditions in unbalanced optimal transport, enabling more stable, faster convergence and solutions with lower action in high-dimensional dynamics recovery.

Contribution

We propose Variational RUOT, which explicitly incorporates optimality conditions into the model, simplifying the solution to a scalar field and improving convergence and stability.

Findings

Var-RUOT finds solutions with lower action.

Var-RUOT exhibits faster convergence.

Var-RUOT improves training stability.

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\"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.