Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport

TL;DR

A new deep learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dynamics from observed snapshots based on the RUOT form, which accurately identifies growth and transition patterns, eliminates false transitions, and constructs the Waddington developmental landscape.

Abstract

Reconstructing dynamics using samples from sparsely time-resolved snapshots is an important problem in both natural sciences and machine learning. Here, we introduce a new deep learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dynamics from observed snapshots. Based on the RUOT form, our method models these dynamics without requiring prior knowledge of growth and death processes or additional information, allowing them to be learned directly from data. Theoretically, we explore the connections between the RUOT and Schrödinger bridge problem and discuss the key challenges and potential solutions. The effectiveness of our method is demonstrated with a synthetic gene regulatory network, high-dimensional Gaussian Mixture Model, and single-cell RNA-seq data from blood development. Compared with other methods, our approach accurately identifies growth and transition patterns, eliminates false transitions, and constructs the Waddington developmental landscape. Our code is available at: https://github.com/zhenyiizhang/DeepRUOT.

Figures (9)

• Figure 1: Overview of DeepRUOT.
• Figure 2: (a) Illustration of the synthetic gene regulatory dynamics. (b) The ground truth cellular dynamics project on $(X_1 , X_2)$. (c) The ground truth growth rates. (d) The dynamics learned by balanced Schrödinger bridge (SF2M sflowmatch, $\sigma=0.25$). (e) The dynamics learned by our DeepRUOT solver ($\sigma=0.25$). (f) The growth rates inferred by our DeepRUOT solver. (g) The Waddington developmental landscape learned at $t=1$. (h) The constructed landscape at $t=4$.
• Figure 3: Application in hematopoiesis scRNA-seq data ($\sigma=0.25$). (a) The stochastic dynamics learned by RUOT. (b) The growth rates learned by DeepRUOT. (c) The constructed Waddington developmental landscape at $t=0$. (d) The landscape at $t=2$.
• Figure 4: Results obtained by UDSB bunne_unsb on gene regulatory network. (a) The trajectory learned by UDSB, where black dots indicate particle death, red dots signify particle growth, orange dots represent the target distribution, dark blue dots denote the source distribution, and gradient blue dots illustrate particle trajectories. (b) Predicted changes in cell population at intermediate time points, with dots representing the actual mass.
• Figure 5: Results of DeepRUOT on Gaussian mixtures ($\sigma=0.1$, 10D). (a) The learned trajectory by DeepRUOT ($\sigma=0.1$). (b) The growth rate inferred by our model. (c) The Waddington developmental landscape at $t=0$ ($\sigma=0.1$). (d) The Waddington developmental landscape at $t=1$ ($\sigma=0.1$).

Theorems & Definitions (16)

• Theorem 3.1
• Definition 4.1: Regularized unbalanced optimal transport
• Theorem 4.1
• Remark 4.1
• Remark 4.2
• Remark 4.3
• Theorem A.1
• proof
• proof
• Theorem D.1