Kirill Neklyudov,Rob Brekelmans,Alexander Tong,Lazar Atanackovic,Qiang Liu,Alireza Makhzani
Abstract:The dynamical formulation of the optimal transport can be extended through various choices of the underlying geometry (kinetic energy), and the regularization of density paths (potential energy). These combinations yield different variational problems (Lagrangians), encompassing many variations of the optimal transport problem such as the Schrödinger bridge, unbalanced optimal transport, and optimal transport with physical constraints, among others. In general, the optimal density path is unknown, and solving these variational problems can be computationally challenging. We propose a novel deep learning based framework approaching all of these problems from a unified perspective. Leveraging the dual formulation of the Lagrangians, our method does not require simulating or backpropagating through the trajectories of the learned dynamics, and does not need access to optimal couplings. We showcase the versatility of the proposed framework by outperforming previous approaches for the single-cell trajectory inference, where incorporating prior knowledge into the dynamics is crucial for correct predictions.
What problem does this paper attempt to address?
The problem this paper attempts to address is how to develop a general computational framework for solving the Wasserstein Lagrangian flow problem. Specifically, the paper focuses on how to optimize various forms of Lagrangian functions through deep learning methods. These functions can describe different optimal transport problems (such as standard optimal transport, Schrödinger bridge, optimal transport under physical constraints, etc.). These problems have important applications in multiple fields, especially in single-cell trajectory inference, where it is necessary to recover the dynamic evolution process of the system from sample data at different time points.
### Main Contributions of the Paper:
1. **Unified Framework**: A general framework based on deep learning is proposed, which can handle various Lagrangian flow problems from a unified perspective.
2. **Efficient Solution**: By using the dual formulation, it avoids simulating or backpropagating through learned dynamic trajectories and does not require access to the optimal coupling.
3. **Flexibility**: The framework can flexibly incorporate different types of prior information, such as physical potentials or nonlinear functions, to improve the accuracy of trajectory inference.
4. **Practical Application**: The framework's application in single-cell RNA sequencing data is demonstrated, particularly in single-cell trajectory inference tasks, where it outperforms existing methods.
### Specific Problems:
- **Trajectory Inference**: Recovering the dynamic evolution process of the system from sample data at different time points.
- **Optimal Transport**: Solving different types of optimal transport problems, including standard optimal transport, Schrödinger bridge, optimal transport under physical constraints, etc.
- **Flexibility and Accuracy**: How to flexibly incorporate prior knowledge in the inference process to improve the accuracy of the inference.
### Application Background:
- **Single-Cell RNA Sequencing**: Provides a high-resolution snapshot of cell populations, allowing high-throughput observation of thousands of genes. However, due to the measurement process ultimately leading to cell death, only the time changes of cells during treatment, differentiation, or development can be observed.
- **Biomedical Research**: Understanding the dynamic evolution process of cells is crucial for predicting future states and understanding biological processes.
### Method Overview:
- **Lagrangian Function**: Defines a Lagrangian function that includes kinetic and potential energy, where kinetic energy describes the velocity dependence of the density path, and potential energy describes the position dependence.
- **Dual Formulation**: Uses the dual form of the Lagrangian function to transform the problem into an optimization problem, avoiding complex simulations and backpropagation.
- **Deep Learning Model**: Uses neural networks to parameterize cotangent vectors and distribution paths, optimizing through gradient descent and gradient ascent.
### Experimental Results:
- Experiments were conducted on multiple single-cell RNA sequencing datasets, and the results show that the framework performs excellently in trajectory inference tasks, especially when combined with physical constraints, significantly improving performance.
In summary, this paper proposes a general and efficient computational framework that solves various Lagrangian flow problems and demonstrates its superior performance in single-cell trajectory inference.