Learning Continuous-Time Dynamics of Tissue Remodeling from Spatial Transcriptomic Snapshots with Biology-Aware Flow Matching
Abstract
Understanding the molecular mechanisms of tissue regeneration requires reconstructing continuous spatiotemporal dynamics from sparse, irregularly sampled data. We introduce a computational framework that learns these dynamics using a multi-marginal flow matching approach. Our model explicitly accounts for cell proliferation and signaling, and constrains inferred trajectories to the data manifold to ensure biological plausibility. To capture the influence of the microenvironment, we incorporate cell-cell signaling representations. For efficient training, we use semi-balanced Fused Gromov-Wassertein couplings to pick multi-marginal samples for training and define conditional reference paths based on cubic Hermite splines. We then decompose the learned velocity field into spatial and gene expression components, disentangling the factors behind tissue migration from molecular state transitions. Applied to an axolotl brain regeneration and ulcerative colitis datasets, our method successfully reconstructs tissue composition at unseen time points, particularly under irregular sampling. It also recovers cell type-specific signaling influences such as Wnt and TGF-(\beta) signaling in neural maturation and interleukin signaling in inflammation. This work provides a scalable and interpretable framework for studying dynamic tissue remodeling from limited spatial transcriptomic observations.