Learning Stochastic Dynamics from Snapshots through Regularized Unbalanced Optimal Transport

2024-10-01

arXiv.org

AI Summary

This study addresses the critical challenge of reconstructing dynamic systems from sparse temporal snapshots. We propose a data-driven, model-agnostic framework for inferring nonequilibrium stochastic dynamics—without presupposing growth or death models. Our method is the first to integrate regularized unbalanced optimal transport (RUOT) with deep neural networks, establishing a rigorous theoretical connection to the Schrödinger bridge problem to enable differentiable probabilistic flow modeling. Evaluated on synthetic gene regulatory networks, high-dimensional Gaussian mixtures, and single-cell RNA-seq data of hematopoietic development, the approach significantly improves dynamic reconstruction fidelity: it accurately identifies cellular proliferation, apoptosis, and migration patterns; eliminates spurious transitions; and robustly recovers the Waddington epigenetic landscape. By dispensing with prior mechanistic assumptions, our framework establishes a new paradigm for modeling developmental and disease progression directly from observational snapshot data.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Machine Learning: Probabilistic Circuits and Graphical Models

Application Category

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

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.

Problem

• Reconstructing stochastic dynamics from snapshots
• Inferring continuous dynamics using RUOT
• Modeling growth and death processes from data

Innovation

• Deep learning for RUOT
• Infer stochastic dynamics
• Model without prior knowledge