STAMImputer: Spatio-Temporal Attention MoE for Traffic Data Imputation

2025-06-09

AI Summary

To address the challenges of rigid static graph structures, weak modeling capability for non-stationary distributions, and imbalanced local-global spatiotemporal correlations under block-wise missing traffic data, this paper proposes a novel imputation framework integrating dynamic graph learning with a spatiotemporal Mixture of Experts (MoE). We innovatively design a Low-Rank Guided Sampling Graph Attention (LrSGAT) mechanism to enable data-driven, adaptive dynamic graph construction. Moreover, we are the first to introduce MoE into traffic imputation, enabling adaptive weighting of spatiotemporal features tailored to block-wise missing patterns. Extensive experiments on four public benchmarks demonstrate significant improvements over state-of-the-art methods, particularly in block-missing imputation accuracy. The source code is publicly available.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Data Mining & Knowledge Management: Mining of Spatial, Temporal or Spatio-Temporal Data

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Traffic data imputation is fundamentally important to support various applications in intelligent transportation systems such as traffic flow prediction. However, existing time-to-space sequential methods often fail to effectively extract features in block-wise missing data scenarios. Meanwhile, the static graph structure for spatial feature propagation significantly constrains the models flexibility in handling the distribution shift issue for the nonstationary traffic data. To address these issues, this paper proposes a SpatioTemporal Attention Mixture of experts network named STAMImputer for traffic data imputation. Specifically, we introduce a Mixture of Experts (MoE) framework to capture latent spatio-temporal features and their influence weights, effectively imputing block missing. A novel Low-rank guided Sampling Graph ATtention (LrSGAT) mechanism is designed to dynamically balance the local and global correlations across road networks. The sampled attention vectors are utilized to generate dynamic graphs that capture real-time spatial correlations. Extensive experiments are conducted on four traffic datasets for evaluation. The result shows STAMImputer achieves significantly performance improvement compared with existing SOTA approaches. Our codes are available at https://github.com/RingBDStack/STAMImupter.

Problem

• Handles block-wise missing data in traffic datasets
• Dynamic graph structure for nonstationary traffic data
• Improves spatial-temporal feature extraction in imputation

Innovation

• Mixture of Experts captures spatio-temporal features
• Low-rank guided Sampling Graph ATtention balances correlations
• Dynamic graphs capture real-time spatial correlations