STG-Mamba: Spatial-Temporal Graph Learning via Selective State Space Model

📄 arXiv: 2403.12418v4 📥 PDF

作者: Lincan Li, Hanchen Wang, Wenjie Zhang, Adelle Coster

分类: cs.LG, cs.AI

发布日期: 2024-03-19 (更新: 2024-05-18)

🔗 代码/项目: GITHUB

💡 一句话要点

提出STG-Mamba以解决时空图学习中的动态特征建模问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 时空图学习 选择性状态空间模型 图神经网络 动态特征建模 卡尔曼滤波

📋 核心要点

  1. 现有的GNN方法主要关注节点间的关系，未能有效建模时空图系统的动态特征，导致预测性能不足。
  2. 本文提出STG-Mamba，利用选择性状态空间模型，系统性地探索时空图的动态状态演化，并引入卡尔曼滤波机制以增强模型的时序特征整合能力。
  3. 在三个基准时空图预测数据集上的实验表明，STG-Mamba在预测性能上超越了现有方法，并显著降低了计算成本。

📝 摘要（中文）

时空图（STG）数据具有动态性、异构性和非平稳性，给时空图学习带来了持续挑战。以往的图神经网络（GNN）方法主要关注节点间关系，忽视了时空系统内在特征的建模。本文提出STG-Mamba，首次将选择性状态空间模型（SSSM）应用于STG学习，利用时空选择性状态空间模块（ST-S3M）精确聚焦于STG潜在特征。此外，提出的卡尔曼滤波图神经网络（KFGN）能够动态整合不同时间粒度的STG嵌入，显著提升了STG预测性能，并有效降低了大规模图网络的计算成本。实验结果表明，STG-Mamba在多个基准数据集上超越了现有最先进方法。

🔬 方法详解

问题定义：本文旨在解决时空图学习中动态特征建模不足的问题。现有方法多集中于节点关系，忽视了时空图系统的内在动态特性，导致预测效果不佳。

核心思路：STG-Mamba通过将时空图视为一个系统，利用选择性状态空间模型（SSSM）来精确建模时空图的动态状态演化，进而提升预测性能。

技术框架：整体架构包括时空选择性状态空间模块（ST-S3M）和卡尔曼滤波图神经网络（KFGN）。ST-S3M专注于提取时空图的潜在特征，而KFGN则负责动态整合不同时间粒度的嵌入。

关键创新：最重要的创新在于将选择性状态空间模型引入时空图学习，形成了新的建模思路，显著区别于传统的GNN方法。

关键设计：模型设计中采用了可学习的卡尔曼滤波机制，以实现对不同时间粒度的嵌入进行动态整合，优化了模型的参数设置和损失函数，提升了计算效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，STG-Mamba在三个基准时空图预测数据集上表现出色，超越了现有的最先进方法，具体提升幅度达到了XX%。同时，模型在计算效率上也显著降低了FLOPs和测试推理时间，展示了其在大规模图网络中的应用潜力。

🎯 应用场景

该研究具有广泛的应用潜力，尤其在交通预测、环境监测和社交网络分析等领域。通过有效建模时空图的动态特征，STG-Mamba能够为决策支持系统提供更准确的预测，进而推动智能城市和智慧交通的发展。未来，随着数据规模的扩大，该方法的计算效率和预测能力将显得尤为重要。

📄 摘要（原文）

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}.