Spatiotemporal Imputation with Graph-Informed Flow Matching

Missing data is a common challenge in spatiotemporal systems, arising in applications such as air quality monitoring and urban traffic management. Traditional machine learning approaches, like recurrent and graph neural networks, rely on iterative propagation, which tends to accumulate errors over time and space. Recent diffusion-based methods mitigate error propagation but require iterative sampling and often depend on problem-agnostic Gaussian priors, limiting both efficiency and effectiveness. To address these limitations, we propose GiFlow, a Graph-Informed Flow Matching framework for spatiotemporal imputation. GiFlow replaces the typical Gaussian prior with a graph-informed prior constructed via spatiotemporal filtering of observable signals, which better aligns the source distribution to the target and thereby simplifies the generation trajectory. The flow field is parameterized by a hybrid vector field model that integrates spatial attention, temporal attention, and spatiotemporal propagation, enabling joint modeling of spatial and temporal dependencies. Extensive experiments on both synthetic and real-world datasets demonstrate that the proposed GiFlow outperforms the state-of-the-art approaches in spatiotemporal imputation. The code is available at https://github.com/zepengzhang/GiFlow.

Spatiotemporal data characterizes both spatial and temporal information and is ubiquitous in domains such as environmental science, urban systems, and climate forecasting. In practice, spatiotemporal data is often incomplete due to sensor failures, transmission errors, or system instability. The incompleteness of spatiotemporal data compromises the reliability of subsequent analyses, motivating the need for robust spatiotemporal imputation techniques. Our paper introduces GiFlow, a novel flow matching framework for spatiotemporal imputation with a graph-informed prior. Unlike the mostly used diffusion-based approaches that rely on a problem-agnostic Gaussian prior, the proposed graph-informed prior better aligns the source distribution to the target and thereby simplifies the generation trajectory. We evaluate GiFlow on both synthetic and real-world datasets. Experimental results show that GiFlow consistently reconstructs missing data more accurately than existing state-of-the-art methods while remaining computationally efficient. These improvements could help make spatiotemporal imputation more reliable and efficient in real-world settings.