Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action

2025-05-17

AI Summary

Reconstructing dynamics from sparse temporal snapshots in high-dimensional systems—such as single-cell evolutionary processes—remains challenging due to ill-posedness and instability. Method: We propose a variational non-equilibrium optimal transport framework grounded in the principle of least action. Unlike existing regularized unbalanced optimal transport (RUOT) and Schrödinger bridge methods—which lack explicit enforcement of optimality conditions and suffer from unstable convergence—we embed the RUOT optimality condition directly into both the parametric structure and the loss function, requiring only a single neural scalar field. We further introduce a biologically informed growth penalty term. The method integrates variational inference, the Wasserstein–Fisher–Rao metric, and neural scalar field modeling. Results: On both synthetic and real single-cell datasets, our approach yields solutions with lower action values, exhibits improved training stability, faster convergence, and significantly outperforms state-of-the-art RUOT and Schrödinger bridge methods.

Technology Category

• Search and Optimization: Evolutionary Computation
• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Evolutionary Learning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems

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.

Problem

• Recovering dynamics from high-dimensional system snapshots
• Enforcing optimality conditions in unbalanced optimal transport
• Improving convergence and stability in transport solutions

Innovation

• Variational RUOT integrates optimality conditions explicitly
• Learns scalar field for lower action solutions
• Improved growth penalty aligns with biological priors