Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schr\"odinger Bridge

TL;DR

This paper introduces CytoBridge, a deep learning framework based on the Unbalanced Mean-Field Schr"odinger Bridge, to model cellular dynamics and interactions from snapshot data, improving accuracy over existing methods.

Contribution

It formulates the UMFSB framework for unbalanced stochastic cell interactions and develops CytoBridge, a neural network-based algorithm to learn these dynamics directly from data.

Findings

Accurately models cell transitions, proliferation, and interactions.

Eliminates false cell state transitions in data.

Reconstructs developmental landscapes more precisely.

Abstract

Modeling the dynamics from sparsely time-resolved snapshot data is crucial for understanding complex cellular processes and behavior. Existing methods leverage optimal transport, Schr\"odinger bridge theory, or their variants to simultaneously infer stochastic, unbalanced dynamics from snapshot data. However, these approaches remain limited in their ability to account for cell-cell interactions. This integration is essential in real-world scenarios since intercellular communications are fundamental life processes and can influence cell state-transition dynamics. To address this challenge, we formulate the Unbalanced Mean-Field Schr\"odinger Bridge (UMFSB) framework to model unbalanced stochastic interaction dynamics from snapshot data. Inspired by this framework, we further propose CytoBridge, a deep learning algorithm designed to approximate the UMFSB problem. By explicitly modeling cellular transitions, proliferation, and interactions through neural networks, CytoBridge offers the flexibility to learn these processes directly from data. The effectiveness of our method has been extensively validated using both synthetic gene regulatory data and real scRNA-seq datasets. Compared to existing methods, CytoBridge identifies growth, transition, and interaction patterns, eliminates false transitions, and reconstructs the developmental landscape with greater accuracy. Code is available at: https://github.com/zhenyiizhang/CytoBridge-NeurIPS.