Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport

TL;DR

This paper presents a deep learning method based on regularized unbalanced optimal transport to infer continuous stochastic dynamics from sparse snapshot data, applicable to biological and high-dimensional systems.

Contribution

It introduces a novel deep learning approach for RUOT that learns dynamics directly from data without prior knowledge, connecting it to Schrödinger bridge problems.

Findings

Successfully applied to synthetic gene networks and single-cell RNA-seq data.

Accurately identifies growth and transition patterns, reducing false positives.

Constructs the Waddington developmental landscape effectively.

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