Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schrödinger Bridge

Zhenyi Zhang, Zihan Wang, Yuhao Sun, Tiejun Li, Peijie Zhou

TL;DR

The paper introduces the Unbalanced Mean Field Schrödinger Bridge (UMFSB) to model interacting, unnormalized particle dynamics inferred from snapshot data and proposes CytoBridge, a neural solver that learns cell transitions, growth, and cell–cell interactions directly from data. By applying Fisher information regularization, SDE constraints are transformed into computationally tractable ODE constraints and simulated with weighted particles, aided by random batch methods. CytoBridge jointly learns a transition velocity, growth rate, density score, and interaction potential, and enforces physical constraints via a physics-informed loss that combines energy, mass reconstruction, and Fokker–Planck terms. The method is validated on synthetic gene networks and diverse scRNA-seq and spatiotemporal datasets, where it outperforms existing trajectory inference approaches in distribution and mass matching and uncovers interpretable growth and interaction patterns with biological relevance.

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ödinger 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ödinger 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.

Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schrödinger Bridge

TL;DR

The paper introduces the Unbalanced Mean Field Schrödinger Bridge (UMFSB) to model interacting, unnormalized particle dynamics inferred from snapshot data and proposes CytoBridge, a neural solver that learns cell transitions, growth, and cell–cell interactions directly from data. By applying Fisher information regularization, SDE constraints are transformed into computationally tractable ODE constraints and simulated with weighted particles, aided by random batch methods. CytoBridge jointly learns a transition velocity, growth rate, density score, and interaction potential, and enforces physical constraints via a physics-informed loss that combines energy, mass reconstruction, and Fokker–Planck terms. The method is validated on synthetic gene networks and diverse scRNA-seq and spatiotemporal datasets, where it outperforms existing trajectory inference approaches in distribution and mass matching and uncovers interpretable growth and interaction patterns with biological relevance.

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ödinger 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ödinger 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.
Paper Structure (50 sections, 2 theorems, 51 equations, 10 figures, 21 tables, 1 algorithm)

This paper contains 50 sections, 2 theorems, 51 equations, 10 figures, 21 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Various Dynamical OT Extensions and Deep Learning Solvers
  4. Modeling Cellular Interactions in Trajectory Inference
  5. Mean-Field Control Problem
  6. Preliminaries and Backgrounds
  7. Regularized Unbalanced Optimal Transport
  8. Mean Field Schrödinger Bridge Problem
  9. Unbalanced Mean Field Schrödinger Bridge
  10. Learning Cell Dynamics and Interactions through Neural Networks
  11. Simulating ODEs: Random Batch Methods
  12. Reformulating the Loss with Weighted Particle Simulation
  13. Energy Loss
  14. Reconstruction Loss
  15. Fokker-Planck Constraint
  16. ...and 35 more sections

Key Result

Theorem 4.1

The unbalanced mean field Schrödinger Bridge def:UISB is equivalent to where $\Psi (\cdot): \mathbb{R} \rightarrow \mathbb{R}$ is the growth cost function, and $\alpha$ is the weight of the growth cost. The infinium is taken over all $(\mathbf{b}, g, {\rho},\Phi)$ subject to $\rho(\mathbf{x},0)=\nu_0, \rho(\mathbf{x},1)=\nu_1$, and where $\Phi$ is the interaction potential and it is satisified

Figures (10)

  • Figure 1: Overview of CytoBridge.
  • Figure 2: (a) Illustration of the synthetic gene regulatory dynamics. (b) The ground truth cellular dynamics project on $(X_1 , X_2)$. The red lines indicate the ground truth trajectories of cells in (b), or inferred trajectories of cells in (c) to (e). (c) The dynamics learned by balanced Schrödinger bridge SF2M sflowmatch. (d) The dynamics learned by DeepRUOT. (e) The dynamics learned by CytoBridge. (f) The learned interaction potential. (g) The growth rates inferred by CytoBridge. (h) The constructed landscape at $t=4$. The z-axis represents the density of cells.
  • Figure 3: Application in mouse blood hematopoiesis data ($\sigma=0.1$), visualized in UMAP space. (a) The overall velocity learned by CytoBridge. (b) The growth rates learned by CytoBridge. (c) The score function learned by CytoBridge at $t=2$. (d) The correlation of velocity and interacting forces.
  • Figure 4: Results of synthetic gene regulatory data with Lennard-Jones Interaction Potential. (a) The dynamics learned by CytoBridge. (b) The learned interaction potential.
  • Figure 5: Comparasions of interacting forces learned from (a) dynamics with Lennard-Jones interaction potential, (b) dynamics without interaction potential.
  • ...and 5 more figures

Theorems & Definitions (8)

  • Definition 3.1: Regularized Unbalanced Optimal Transport
  • Definition 3.2: Mean Field Schrödinger Bridge Problem
  • Definition 4.1: Unbalanced Mean Field Schrödinger Bridge
  • Theorem 4.1
  • Remark 4.1
  • Proposition 5.1
  • proof
  • Remark D.1