Modeling Microenvironment Trajectories on Spatial Transcriptomics with NicheFlow

TL;DR

This work tackles the challenge of modeling tissue-scale spatiotemporal dynamics from time-resolved spatial transcriptomics by shifting focus from single cells to cellular microenvironments, i.e., neighborhoods represented as point clouds. It introduces NicheFlow, a flow-based generative model that combines entropic optimal transport with Variational Flow Matching, conditioned on source microenvironments, and implemented via a Microenvironment Transformer backbone. The method uses mixed-factorized posteriors for coordinates and gene expression, enabling joint generation of spatial and transcriptional changes across two consecutive timepoints, and demonstrates superior reconstruction of spatial architecture and cell-type composition across embryonic development, brain development, and aging datasets compared with baselines. The results are supported by a suite of metrics (PSD/SPD, 1NN-F1, GW/FGW, Wasserstein) and qualitative analyses, underscoring the potential of principled, region-level generative modeling for perturbation prediction and multi-omics integration in spatial biology.

Abstract

Understanding the evolution of cellular microenvironments in spatiotemporal data is essential for deciphering tissue development and disease progression. While experimental techniques like spatial transcriptomics now enable high-resolution mapping of tissue organization across space and time, current methods that model cellular evolution operate at the single-cell level, overlooking the coordinated development of cellular states in a tissue. We introduce NicheFlow, a flow-based generative model that infers the temporal trajectory of cellular microenvironments across sequential spatial slides. By representing local cell neighborhoods as point clouds, NicheFlow jointly models the evolution of cell states and spatial coordinates using optimal transport and Variational Flow Matching. Our approach successfully recovers both global spatial architecture and local microenvironment composition across diverse spatiotemporal datasets, from embryonic to brain development.

Figures (16)

• Figure 1: Overview of NicheFlow. At time $t_1$, we generate a target microenvironment $\mathcal{M}^1$ by transforming Gaussian noise $\mathcal{M}^z$ using a Variational Flow Matching model with a posterior $\mu_t^{\theta}$ conditioned on a source microenvironment $\mathcal{M}^0$ at $t_0$. Source-target pairs are identified via entropic OT over pooled microenvironment coordinates and gene expression profiles.
• Figure 2: Qualitative comparison of generated samples on the embryonic development dataset (9.5-11.5 days). We show source and target samples alongside predictions from SPFlow and NicheFlow with different objectives.
• Figure 3: Left and right panels show the mapping of spinal cord (E10.5 to E11.5) and head neural crest cells (E9.5 to E10.5). In each panel, the left column shows source cells and expected targets, and the right column shows density contours of the most likely mapped regions. Bar plots display transition probabilities to the most likely descendant cell types. For NicheFlow, contours represent the proportion of samples in generated point clouds assigned to real cell coordinates across 10 samples.
• Figure 4: Comparison of NicheFlow and moscot on the prediction of the anterior neural crest cells' fate. For both models, we take source facial neural crest cells at E9.5, push them to time point E10.5, and show the compositional and density predictions in the middle panel. Then, the predictions at 10.5 are used as a source for a second trajectory prediction operation from 10.5 to 11.5, for which we inspect again the cell density over the target slide and the cell type probabilities.
• Figure 5: Comparison between moscot and NicheFlow on mapping the liver structure from E10.5 to E11.5. The liver at time E10.5 is used as a source for trajectory prediction using the different models. The left column shows the source and expected target regions highlighted on the respective E10.5 and E11.5 embryos. The middle column displays the density of the prediction obtained by transporting niches from the source to the target slide. On the right, the aggregated cell type proportions according to the density in the middle column (see \ref{['app: biological experiments']}).

Theorems & Definitions (1)

• Proposition 1