[论文解读] STG-Mamba: Spatial-Temporal Graph Learning via Selective State Space Model

Lincan Li, Hanchen Wang+2 more|arXiv (Cornell University)|Mar 19, 2024

Advanced Graph Neural Networks被引用 10

一句话总结

STG-Mamba 通过选择性状态空间模型用于时空图学习，结合 GS3B 与 KFGN 实现高效的 Seq2Seq 预测，线性复杂度与最先进的准确性。

ABSTRACT

Spatial-Temporal Graph (STG) data is characterized as dynamic, heterogenous, and non-stationary, leading to the continuous challenge of spatial-temporal graph learning. In the past few years, various GNN-based methods have been proposed to solely focus on mimicking the relationships among node individuals of the STG network, ignoring the significance of modeling the intrinsic features that exist in STG system over time. In contrast, modern Selective State Space Models (SSSMs) present a new approach which treat STG Network as a system, and meticulously explore the STG system's dynamic state evolution across temporal dimension. In this work, we introduce Spatial-Temporal Graph Mamba (STG-Mamba) as the first exploration of leveraging the powerful selective state space models for STG learning by treating STG Network as a system, and employing the Spatial-Temporal Selective State Space Module (ST-S3M) to precisely focus on the selected STG latent features. Furthermore, to strengthen GNN's ability of modeling STG data under the setting of selective state space models, we propose Kalman Filtering Graph Neural Networks (KFGN) for dynamically integrate and upgrade the STG embeddings from different temporal granularities through a learnable Kalman Filtering statistical theory-based approach. Extensive empirical studies are conducted on three benchmark STG forecasting datasets, demonstrating the performance superiority and computational efficiency of STG-Mamba. It not only surpasses existing state-of-the-art methods in terms of STG forecasting performance, but also effectively alleviate the computational bottleneck of large-scale graph networks in reducing the computational cost of FLOPs and test inference time. The implementation code is available at: \url{https://github.com/LincanLi98/STG-Mamba}.

研究动机与目标

• Motivate STG forecasting on dynamic, heterogeneous spatial-temporal graphs.
• Propose a graph-selective state space architecture to model STG dynamics.
• Develop adaptive graph learning to handle evolving STG structures.
• Demonstrate scalability and efficiency advantages over Transformer-based STG models.

提出的方法

• Formulate STG learning as an Encoder-Decoder with stacked Graph Selective State Space Blocks (GS3B).
• Integrate Kalman Filtering Graph Neural Networks (KFGN) to generate adaptive, time-varying graph adjacency matrices.
• Employ Graph State Space Layer (GSS-Layer) for input-dependent feature learning.
• Discretize continuous state-space dynamics via zero-order hold to enable efficient deep learning implementations.
• Show linear time complexity in sequence length O(L) and reduced FLOPs compared to Transformer baselines.

实验结果

研究问题

• RQ1 Can selective state space models effectively capture dynamic STG dependencies?
• RQ2 How can adaptive graph learning be integrated with SSSMs for STG data?
• RQ3 Do GS3B/KFGN-based STG-Mamba achieve superior forecasting accuracy with lower computational cost compared to Transformer-based STG models?

主要发现

• STG-Mamba consistently outperforms baselines on PeMS04, HZMetro, and KnowAir across RMSE, MAE, and MAPE, with minor exception in one metric.
• STG-Mamba achieves notable computational efficiency, exhibiting linear FLOP growth with respect to STG size and shorter inference times than comparator Transformers.
• Across datasets, STG-Mamba attains the best RMSE and MAE in most tasks (e.g., PeMS04: 29.53 RMSE, 18.09 MAE; KnowAir: 18.05 RMSE, 11.73 MAE).
• STG-Mamba demonstrates robust performance under varying external conditions (rush vs non-rush hours, weekend vs weekday).
• The model reduces FLOPs significantly relative to STAEformer at larger STG scales (e.g., PeMS04 with 300 nodes: 36.85G vs 92.49G).