PaperlayerView exact revision on arXiv

Source-linked AI summary

PriSTI: A Conditional Diffusion Framework for Spatiotemporal Imputation

Mingzhe Liu, Han Huang, Hao Feng, Leilei Sun, Bowen Du, Yanjie Fu

arXiv:2302.09746v1cs.LG

TL;DR

Spatiotemporal imputation must recover values lost through sensor failures or transmission loss while modeling complex dependencies. PriSTI uses conditional diffusion with an extracted global context prior, spatiotemporal attention, and geographic relationships, and it outperforms existing methods across missing-data settings.

  • Problem

    Spatiotemporal data are frequently incomplete, and diffusion-based imputation still faces challenges in constructing and using conditional information for spatiotemporal dependencies.

  • Method

    PriSTI is a conditional diffusion framework that extracts conditional features as a global context prior and uses geographic relationships in noise estimation.

  • Results

    PriSTI outperforms existing imputation methods across various missing patterns and real-world spatiotemporal datasets, including high-missing-rate and sensor-failure scenarios.

  • Takeaways & Limitations

    Constructing and using conditional information with spatiotemporal global correlations and geographic relationships improves PriSTI’s spatiotemporal imputation performance.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Spatiotemporal data are often incomplete, while existing imputation approaches struggle to model complex dependencies and conditional information. PriSTI addresses this with conditional diffusion, enhanced prior modeling, and geographic relationships, achieving strong performance across missing-data settings.

  • Problem motivation: Spatiotemporal data become incomplete through sensor failures and transmission loss, making missing-value imputation important for downstream analysis.Missing positions and patterns vary, so imputation must exploit dependencies in available observations.
  • Problem motivation: Diffusion models avoid autoregressive error accumulation but make constructing and introducing consistent conditional information challenging.The resulting inconsistency increases the difficulty of learning spatiotemporal dependencies.
  • Proposed approach: PriSTI extracts conditional features from observations as a global context prior and uses them to calculate spatiotemporal attention weights for noise prediction.The framework also incorporates geographic relationships when constructing and using conditional information.
  • Scope and contribution: The method is designed to handle high missing rates and sensor-failure scenarios while supporting downstream tasks through imputation.Its contributions include constructing conditional information from spatiotemporal global correlations and geographic relationships.

III. METHODOLOGY

PriSTI formulates spatiotemporal imputation as conditional generation, using observed data, interpolated conditional information, and geographic structure. Its reverse diffusion process progressively converts noise into missing values while a specialized noise-prediction model learns spatiotemporal dependencies.

  • Diffusion model for imputation: PriSTI treats imputation as conditional generation and uses a reverse diffusion process to convert random noise into missing values.The reverse process is conditioned on interpolated observations and geographic information.
  • Diffusion model for imputation: The model predicts added Gaussian noise from a noisy imputation target together with conditional information X and adjacency matrix A.The noise-prediction network is flexible in architecture, allowing specialization for spatiotemporal imputation.
  • Training process: During training, PriSTI masks observed values, interpolates the remaining observations, and samples diffusion steps and Gaussian noise.Point, block, and hybrid mask strategies create different training targets.
  • Imputation process: During imputation, the model constructs interpolated conditional information from all observed values and generates missing-value samples from noisy inputs.The trained noise-prediction model receives the noisy target and conditional information during reverse diffusion.
  • Framework pipeline: PriSTI’s pipeline first interpolates observations and extracts a global context prior, then uses conditional information to guide noise estimation.This specialized design addresses the difficulty of learning spatiotemporal dependencies from noisy information.

B. Design of Noise Prediction Model

PriSTI’s noise prediction design uses conditional information to model spatiotemporal dependencies during diffusion-based imputation. Its two modules share temporal attention, spatial attention, and message passing components while serving different architectural roles.

  • PriSTI first extracts spatiotemporal dependencies from interpolated conditional information before estimating noise for imputation.The conditional feature extraction module provides the prior used by the noise estimation module.
  • The conditional feature extraction and noise estimation modules both use temporal attention, spatial attention, and a message passing neural network.
  • The conditional feature extraction module uses a wide architecture, whereas the noise estimation module uses a deep architecture.

1) Conditional Feature Extraction Module:

The conditional feature extraction module constructs a global context prior from interpolated observations and geographic structure. This prior captures dependencies without the Gaussian noise added to the diffusion sample.

  • 1) Conditional Feature Extraction Module:: The module addresses the difficulty of mixing conditional and noisy information in the noise prediction model.
  • 1) Conditional Feature Extraction Module:: Linear interpolation creates coarse conditional information that fills observed-value gaps before feature extraction.The interpolation is intended to provide an effective starting representation for denoising.
  • 1) Conditional Feature Extraction Module:: The learnable module γ models spatiotemporal information from interpolated conditional data because interpolation alone omits nonlinear temporal relationships and spatial correlations.
  • 1) Conditional Feature Extraction Module:: Hpri = γ(H, A) combines spatiotemporal global correlations with geographic dependency and supplies the global context for attention-weight calculation.Here, H is produced by a 1 × 1 convolution from the interpolated conditional information.
  • 1) Conditional Feature Extraction Module:: The extracted conditional feature contains temporal dependencies, spatial global correlations, and geographic dependencies without added Gaussian noise.It therefore serves as a coarse prior for subsequent spatiotemporal dependency learning.

2) Noise Estimation Module:

The noise estimation module uses noisy interpolated inputs together with the conditional prior and adjacency matrix to learn temporal and spatial dependencies. It derives attention weights from the prior and reduces spatial-attention cost through virtual nodes.

  • 2) Noise Estimation Module:: The noise estimation inputs combine interpolation information, a noise sample, the conditional feature Hpri, and the adjacency matrix A.
  • 2) Noise Estimation Module:: The module first learns temporal features and then aggregates them through spatial dependency learning to model global correlations and geographic relationships.
  • 2) Noise Estimation Module:: Conditional information Hpri supplies attention weights while the module converts noisy inputs toward Gaussian-noise estimates, reducing noise interference in dependency learning.This design avoids calculating attention weights directly from the increasingly noisy sample.
  • 2) Noise Estimation Module:: The module retains the same attention and message-passing components as the conditional feature extractor but uses different inputs and architecture.
  • 2) Noise Estimation Module:: Mapping N nodes to k virtual nodes reduces spatial-attention complexity from O(N^2d) to O(Nkd), with k < N.

3) Auxiliary Information and Output:

PriSTI supplements both modules with temporal and spatial auxiliary information, then aggregates layered outputs to produce the noise prediction for the imputation target.

  • 3) Auxiliary Information and Output:: Auxiliary information combines sine-cosine temporal encoding with learnable node embeddings through an MLP.The resulting representation is supplied to both the conditional feature extraction and noise estimation modules.
  • 3) Auxiliary Information and Output:: The noise estimation module stacks layers whose residual connections feed later layers and whose skip connections are combined to produce the noise prediction.
  • 3) Auxiliary Information and Output:: The final output retains only the imputation target value, and the loss is calculated using Equation (4).

IV. EXPERIMENTS

The experiments evaluate PriSTI across real-world datasets, missing rates, missing patterns, conditional-information designs, and spatial-temporal dependency questions.

  • Research questions: The evaluation compares PriSTI with established methods under multiple missing patterns and asks whether it improves imputation performance.The experiments also examine performance across missing rates and the construction, utilization, and learning of spatiotemporal dependencies.
  • Research questions: A further research question tests whether PriSTI can impute time series for unobserved sensors using only geographic location.
  • Datasets: Experiments test PriSTI on AQI-36, METR-LA, and PEMS-BAY, covering air-quality and traffic-speed data.AQI-36 contains 36 Beijing stations; METR-LA and PEMS-BAY contain 207 and 325 highway sensors, respectively.

B. Baselines

The study compares PriSTI with statistical, machine-learning, matrix-factorization, autoregressive, and generative imputation baselines using deterministic and probabilistic metrics.

  • Baseline categories: Baselines span statistic, classic machine-learning, low-rank factorization, deep autoregressive, and deep generative methods.Named examples include MEAN, KNN, MICE, VAR, TRMF, BATF, BRITS, GRIN, V-RIN, GP-VAE, rGAIN, and CSDI.
  • Implementation: The reported baseline results combine reproduced implementations with results retained from GRIN for several methods after consistency checks.
  • Evaluation metrics: The evaluation uses MAE and MSE for deterministic imputation and CRPS for compatibility between estimated probability distributions and observed values.CRPS is computed from quantile losses over distributions approximated with generated samples.

D. Experimental settings

The experimental setup defines dataset splits, missingness simulations, training masks, optimization settings, and evaluations using deterministic, probabilistic, downstream, and comparative results.

  • Dataset and imputation settings: Traffic-data evaluation includes block and point missingness, while training uses point, block, or hybrid masking matched to each dataset and missing pattern.AQI-36 uses historical missing patterns with hybrid masking; traffic block missingness uses hybrid-block masking, and point missingness uses point masking.
  • Implementation settings: The experiments use batch size 16, scheduled learning-rate decay, a quadratic diffusion noise schedule, sinusoidal embeddings, and five runs.
  • Overall evaluation: The primary comparison reports MAE and MSE in Table III and CRPS in Table IV, with probabilistic outputs approximated from 100 generated samples.The deterministic result for sampled methods is the median of generated samples.
  • Overall evaluation: PriSTI outperforms other baselines across missing patterns and datasets, including CSDI, while diffusion methods outperform autoregressive approaches in the reported comparisons.The text attributes the autoregressive gap to error accumulation and reports a stronger PriSTI improvement over CSDI for traffic block missingness than point missingness.
  • Downstream evaluation: On AQI-36, downstream Graph Wavenet forecasting follows the imputation ranking, showing that imputation quality is reflected in subsequent prediction performance.The evaluation uses past 12 time steps to predict the next 12 and reports MAE and RMSE.

2) Sensitivity analysis (RQ2):

Sensitivity analysis examines PriSTI under 10%-90% missing rates and ablates conditional-information and spatiotemporal-dependency components.

  • Sensitivity analysis: At 90% missingness on METR-LA, PriSTI improves MAE over comparison methods by 4.67%-34.11% for block missingness and 3.89%-43.99% for point missingness.The comparison is reported for the sparse-data sensitivity experiment and illustrated in Figure 5.
  • Sensitivity analysis: PriSTI retains better imputation performance at high missing rates and shows larger gains as the observed data become sparser.The authors relate this to interpolated conditional information retaining spatiotemporal dependencies under sparse observations.
  • Ablation study: Ablations evaluate conditional-feature extraction, noise-estimation inputs, spatial and temporal dependency modules, message passing, and spatial global attention.The variants include mix-STI, w/o CF, w/o spa, w/o tem, w/o MPNN, and w/o Attn.
  • Ablation study: Removing enhanced conditional information or conditional features reduces imputation effectiveness, while removing either spatial or temporal dependency learning also harms performance.The ablation results are reported on AQI-36 and METR-LA using MAE.

4) Case study (RQ4):

PriSTI produces probabilistic imputations that follow observed temporal trends and leverage spatial relationships, including when sensors have extensive or complete observation loss. In sensor-failure experiments, it outperforms the geographic-information baseline GRIN on unobserved nodes.

  • Probabilistic imputation: The visualization uses black crosses for observations, colored dots for missing-value ground truth, a green deterministic estimate, and a 0.05–0.95 quantile shadow.These encodings are used for AQI-36 and block-missing METR-LA cases.
  • Probabilistic imputation: PriSTI’s generated distributions cover most observations and follow the time trends of different METR-LA sensors, including sensors with continuous or complete missingness.The case study indicates that temporal dependencies are learned from available observations while spatial dependencies support imputation when temporal information is limited.
  • Sensor failure: PriSTI imputes completely unobserved sensors using geographic and cross-sensor information when only sensor locations remain available.This setting is related to reconstructing a time series at a location from other sensors’ observations.
  • Sensor failure: MAE is 10.23 for station 14 and 15.20 for station 31, with PriSTI outperforming GRIN on the unobserved-node experiment.All observations from the two selected AQI-36 stations were masked during training; Figure 7 compares the resulting imputations with GRIN.

6) Hyperparameter analysis and time costs:

The analysis selects βT = 0.2 and sets d and k to 64, balancing performance with efficiency. PriSTI’s generative modeling also incurs higher METR-LA training and inference costs than CSDI.

  • Hyperparameter analysis: βT = 0.2 is selected as optimal, while d = 64 and k = 64 are chosen for efficiency despite better performance at larger values.d and k control learned information capacity, whereas βT controls sampled noise and can hinder noise prediction when too large or too small.
  • Hyperparameter analysis: Hyperparameter sensitivity is evaluated on METR-LA for hidden-state channel size d, maximum noise level βT, and virtual-node count k.The study examines how information capacity and sampled-noise level affect PriSTI’s performance.
  • Time costs: The efficiency gap among methods is smaller on AQI-36, which has fewer nodes, than on the larger METR-LA dataset.The supplied discussion attributes the larger generative-method cost difference to the METR-LA setting.
  • Time costs: PriSTI requires 25.7% more training time and 17.9% more inference time than CSDI on METR-LA.The comparison uses total training and inference time on an AMD EPYC 7371 CPU with an NVIDIA RTX 3090.
  • Time costs: The paper identifies improving scalability and computation efficiency on larger spatiotemporal datasets as future work.It also proposes refining conditional information for imputing longer temporal dependencies.
Loading 2302.09746v1…