PriSTI: A Conditional Diffusion Framework for Spatiotemporal Imputation

TL;DR

PriSTI introduces a novel conditional diffusion framework that effectively imputes missing spatiotemporal data by leveraging enhanced prior modeling and attention mechanisms, outperforming existing methods especially in high missing rate scenarios.

Contribution

This paper presents PriSTI, a new diffusion-based model that improves spatiotemporal imputation by incorporating conditional feature extraction and spatiotemporal attention, addressing error accumulation issues of prior autoregressive models.

Findings

Outperforms existing imputation methods across various missing data patterns.

Effectively handles high missing rates and sensor failures.

Utilizes conditional features and attention for improved accuracy.

Abstract

Spatiotemporal data mining plays an important role in air quality monitoring, crowd flow modeling, and climate forecasting. However, the originally collected spatiotemporal data in real-world scenarios is usually incomplete due to sensor failures or transmission loss. Spatiotemporal imputation aims to fill the missing values according to the observed values and the underlying spatiotemporal dependence of them. The previous dominant models impute missing values autoregressively and suffer from the problem of error accumulation. As emerging powerful generative models, the diffusion probabilistic models can be adopted to impute missing values conditioned by observations and avoid inferring missing values from inaccurate historical imputation. However, the construction and utilization of conditional information are inevitable challenges when applying diffusion models to spatiotemporal imputation. To address above issues, we propose a conditional diffusion framework for spatiotemporal imputation with enhanced prior modeling, named PriSTI. Our proposed framework provides a conditional feature extraction module first to extract the coarse yet effective spatiotemporal dependencies from conditional information as the global context prior. Then, a noise estimation module transforms random noise to realistic values, with the spatiotemporal attention weights calculated by the conditional feature, as well as the consideration of geographic relationships. PriSTI outperforms existing imputation methods in various missing patterns of different real-world spatiotemporal data, and effectively handles scenarios such as high missing rates and sensor failure. The implementation code is available at https://github.com/LMZZML/PriSTI.