Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling

Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, Jian Sun

TL;DR

VGFM tackles the challenge of inferring time-evolving single-cell distributions from snapshot data where cell number changes violate mass conservation. It jointly learns a velocity field $v_\theta(\mathbf{x},t)$ and a growth function $g_\omega(\mathbf{x},t)$ by tying a two-period dynamic view of semi-relaxed OT to a practical joint flow-matching objective, enabling stable high-dimensional learning. The method combines a flow-matching loss with a distribution-fitting loss based on Wasserstein distance, and demonstrates superior performance on synthetic and real scRNA-seq datasets in reconstructing dynamics and mass changes, surpassing unbalanced baselines. VGFM's ability to capture branch-specific growth and state transitions offers a scalable tool for trajectory inference and growth-rate estimation in complex cellular systems, with potential impact on developmental biology and disease modeling. Limitations include reliance on a partially simulation-free framework and potential entanglement of biological growth with observational effects; future work proposes incorporating priors and pursuing fully simulation-free alternatives.

Abstract

Learning the underlying dynamics of single cells from snapshot data has gained increasing attention in scientific and machine learning research. The destructive measurement technique and cell proliferation/death result in unpaired and unbalanced data between snapshots, making the learning of the underlying dynamics challenging. In this paper, we propose joint Velocity-Growth Flow Matching (VGFM), a novel paradigm that jointly learns state transition and mass growth of single-cell populations via flow matching. VGFM builds an ideal single-cell dynamics containing velocity of state and growth of mass, driven by a presented two-period dynamic understanding of the static semi-relaxed optimal transport, a mathematical tool that seeks the coupling between unpaired and unbalanced data. To enable practical usage, we approximate the ideal dynamics using neural networks, forming our joint velocity and growth matching framework. A distribution fitting loss is also employed in VGFM to further improve the fitting performance for snapshot data. Extensive experimental results on both synthetic and real datasets demonstrate that VGFM can capture the underlying biological dynamics accounting for mass and state variations over time, outperforming existing approaches for single-cell dynamics modeling.

Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling

TL;DR

VGFM tackles the challenge of inferring time-evolving single-cell distributions from snapshot data where cell number changes violate mass conservation. It jointly learns a velocity field and a growth function by tying a two-period dynamic view of semi-relaxed OT to a practical joint flow-matching objective, enabling stable high-dimensional learning. The method combines a flow-matching loss with a distribution-fitting loss based on Wasserstein distance, and demonstrates superior performance on synthetic and real scRNA-seq datasets in reconstructing dynamics and mass changes, surpassing unbalanced baselines. VGFM's ability to capture branch-specific growth and state transitions offers a scalable tool for trajectory inference and growth-rate estimation in complex cellular systems, with potential impact on developmental biology and disease modeling. Limitations include reliance on a partially simulation-free framework and potential entanglement of biological growth with observational effects; future work proposes incorporating priors and pursuing fully simulation-free alternatives.

Abstract

Learning the underlying dynamics of single cells from snapshot data has gained increasing attention in scientific and machine learning research. The destructive measurement technique and cell proliferation/death result in unpaired and unbalanced data between snapshots, making the learning of the underlying dynamics challenging. In this paper, we propose joint Velocity-Growth Flow Matching (VGFM), a novel paradigm that jointly learns state transition and mass growth of single-cell populations via flow matching. VGFM builds an ideal single-cell dynamics containing velocity of state and growth of mass, driven by a presented two-period dynamic understanding of the static semi-relaxed optimal transport, a mathematical tool that seeks the coupling between unpaired and unbalanced data. To enable practical usage, we approximate the ideal dynamics using neural networks, forming our joint velocity and growth matching framework. A distribution fitting loss is also employed in VGFM to further improve the fitting performance for snapshot data. Extensive experimental results on both synthetic and real datasets demonstrate that VGFM can capture the underlying biological dynamics accounting for mass and state variations over time, outperforming existing approaches for single-cell dynamics modeling.
Paper Structure (33 sections, 5 theorems, 38 equations, 18 figures, 11 tables, 1 algorithm)

This paper contains 33 sections, 5 theorems, 38 equations, 18 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Related Works
  3. Optimal transport.
  4. Flow matching for learning single cell dynamics.
  5. Unbalancedness-aware distribution learning methods.
  6. Method
  7. Unbalanced Dynamics of Single Cell
  8. Building Unlabalnced Dynamics Based on Semi-Relaxed Optimal Transport
  9. Dynamic understanding of semi-relaxed optimal transport.
  10. Building unbalanced dynamics.
  11. Velocity and Growth Flow Matching
  12. Training Process
  13. Total loss and algorithm.
  14. Experiments
  15. Results and Analysis
  16. ...and 18 more sections

Key Result

Proposition 1

Assume $c(\mathbf{x}_0, \mathbf{x}_1) = \|\mathbf{x}_0 - \mathbf{x}_1\|^2$ and if we enforce ${\rm P}_{\#}^0\pi$ and $p_0$ to share the same support for admissible solution $\pi$ to problem eq:semikp, then we have $\min_\pi \mathcal{J}_{\rm sot}(\pi) = \min_{v_t,g_t}\mathcal{J}^\lambda_{\rm tpt}$$(v

Figures (18)

  • Figure 1: The goal of this paper is to learn a joint state transition (controlled by $v_\theta(\mathbf{x},t)$) and mass growth (controlled by $g_\omega(\mathbf{x},t)$) dynamics for single-cell evolution.
  • Figure 2: Illustration of our proposed VGFM, consisting of the velocity and growth flow matching deduced by the dynamic reformulation of the semi-relaxed optimal transport.
  • Figure 3: (a) Predicted dynamics with trajectories of VGFM on Simulation Gene. (b) Predicted and (c) true growth rates on Simulation Gene. Note that for Simulation Gene, the true growth rates could be accessed. (d), (e) and (f) respectively compare predicted relative mass by UDSB, TIGON, DeepRUOT, and VFGM, and observed relative mass from data, on three datasets. Note that in (f), on 50D dataset, the mass predicted by TIGON deviates significantly from the observed trend, because of its difficulty in handling high-dimensional data as discussed in their paper sha2024reconstructing.
  • Figure 4: Visualization of predicted dynamics by (a) VFGM (w/o $\mathcal{L}_{\rm OT}$) and (b) VGFM, and growth rate by (c) VFGM, on EB (5D) dataset, where the hold-out time is the first intermediate timepoint.
  • Figure A-5: Empirical evidence of Theorem \ref{['thm:thm1']}.(a) Original two-period transport dynamics. (b) Joint dynamics defined via reparameterization \ref{['eq:reparam']}. (c) The two dynamics ended in the same distribution.
  • ...and 13 more figures

Theorems & Definitions (7)

  • Proposition 1
  • Theorem 1
  • Theorem A-2: Brenier-Benamou formula
  • Proposition A-2
  • proof
  • Theorem A-3
  • proof