ContextFlow: Context-Aware Flow Matching for Trajectory Inference from Spatial Omics Data

ICLR 2026 Conference Submission20626 Authors

Keywords: Flow-matching, optimal transport, spatial omics

TL;DR: We propose a novel approach to incorporate biological prior knowledge into flow-matching framework for spatial omics

Abstract: Inferring trajectories from longitudinal spatially-resolved omics data is fundamental to understanding the dynamics of structural and functional tissue changes in development, regeneration and repair, disease progression, and response to treatment. We propose ContextFlow, a novel context-aware flow matching framework that incorporates prior knowledge to guide the inference of structural tissue dynamics from spatially resolved omics data. Specifically, ContextFlow integrates local tissue organization and ligand-receptor communication patterns into a transition plausibility matrix that regularizes the optimal transport objective. By embedding these contextual constraints, ContextFlow generates trajectories that are not only statistically consistent but also biologically meaningful, making it a generalizable framework for modeling spatiotemporal dynamics from longitudinal, spatially resolved omics data. Evaluated on three datasets, ContextFlow consistently outperforms state-of-the-art flow matching methods across multiple quantitative and qualitative metrics of inference accuracy and biological coherence.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 20626