AI论文速读 |（Mamba×时空图预测！） STG-Mamba：通过选择性状态空间模型进行时空图学习

（来了来了，虽迟但到，序列建模的新宠儿mamba终于杀入了时空预测！）
论文标题：STG-Mamba: Spatial-Temporal Graph Learning via Selective State Space Model

作者：Lincan Li, Hanchen Wang（王翰宸）, Wenjie Zhang（张文杰）, Adelle Coster

机构：新南威尔士大学（UNSW）

论文链接：https://arxiv.org/abs/2403.12418

Cool Paper：https://papers.cool/arxiv/2403.12418

TL; DR：本文提出的STG-Mamba是一种新型的时空图学习方法，它通过选择状态空间模型和卡尔曼滤波图神经网络，有效提高了STG数据学习的准确性和计算效率。

关键词：时空图预测，选择性状态空间模型（SSSM），卡尔曼滤波

标题和作者
🌟【紧跟前沿】“时空探索之旅”与你一起探索时空奥秘！🚀
欢迎大家关注时空探索之旅时空探索之旅QRCode

摘要

时空图（STG）数据具有动态、异构和非平稳的特点，导致时空图学习面临持续的挑战。在过去的几年里，人们提出了各种基于 GNN 的方法，仅仅专注于模仿 STG 网络的节点个体之间的关系，而忽略了对 STG 系统随时间推移存在的内在特征进行建模的重要性。相比之下，现代选择性状态空间模型（SSSM）提出了一种将 STG 网络视为一个系统的新方法，并仔细探索 STG 系统跨时间维度的动态状态演化。在这项工作中，引入时空图 Mamba (STG-Mamba)，作为通过将 STG 网络视为一个系统，并采用图选择性状态空间块 (Graph Selective State Space Block, GS3B) 来首次探索SSSM进行 STG 学习的强大能力，以求精确展示 STG 网络的动态演化。 STG-Mamba 被制定为一种编码器-解码器架构，以 GS3B 为基本模块，用于高效的序列数据建模。此外，为了增强 GNN 在 SSSM 设置下对 STG 数据进行建模的能力，提出了用于自适应图结构升级的卡尔曼滤波图神经网络（Kalman Filtering Graph Neural Networks, KFGN）。 KFGN 平滑地适应选择性状态空间演化的背景，同时保持线性复杂度。对三个基准 STG 预测数据集进行了广泛的实证研究，证明了 STG-Mamba 的性能优越性和计算效率。它不仅在STG预测性能方面超越了现有的最先进的方法，而且还有效缓解了大规模图网络在减少FLOP的计算成本和测试推理时间方面的计算瓶颈。

Q: 这篇论文试图解决什么问题？

A: 这篇论文试图解决的问题是如何有效地进行时空图（Spatial-Temporal Graph, STG）数据的学习。STG数据具有动态、异构和非平稳的特性，这导致了时空图学习持续面临的挑战。以往的基于图神经网络（GNN）的方法主要关注于模仿STG网络中节点个体之间的关系，而忽略了对STG系统随时间变化的内在特征进行建模的重要性。

为了解决这一问题，论文提出了一种新的方法，即时空图曼巴（STG-Mamba），它是首个利用选择性状态空间模型（Selective State Space Models, SSSMs）对STG学习进行探索的工作。STG-Mamba将STG网络视为一个系统，并细致地探索STG系统在时间维度上的动态状态演变。具体来说，STG-Mamba通过引入图选择性状态空间块（Graph Selective State Space Block, GS3B）来精确描述STG网络的动态演变，并且提出了卡尔曼滤波图神经网络（Kalman Filtering Graph Neural Networks, KFGN）来适应图结构的升级。这些方法使得STG-Mamba在保持线性复杂度的同时，能够有效地建模STG数据，并且在三个基准STG预测数据集上的广泛实证研究中展示了其性能优越性和计算效率。

Q: 有哪些相关研究？

A: 这篇论文中提到的相关研究主要集中在时空图（STG）数据的学习领域，包括但不限于以下几个方面的工作：

基于图神经网络（GNN）的方法：
- STGCN (Yu et al., 2018)：利用图神经网络捕捉交通网络中的复杂时空依赖关系。
- STSGCN (Song et al., 2020)：通过同步机制有效处理复杂的局部时空相关性和数据异质性。
- STG-NCDE (Choi et al., 2022)：利用神经控制微分方程（NCDEs）的概念进行时空处理。
- DDGCRN (Weng et al., 2023)：通过分离正常和异常信号并采用新颖的图卷积循环网络方法动态捕捉时空特征。
基于注意力机制（Transformer）的方法：
- ASTGCN (Guo et al., 2019)：结合空间注意力和时间注意力机制与图神经网络的创新方法。
- ASTGNN (Guo et al., 2021)：引入独特的趋势感知自注意力和动态图卷积来处理时间动态性、空间相关性。
- PDFormer (Jiang et al., 2023)：引入了一种定制的Transformer模型，擅长捕捉动态空间依赖性和交通条件传播的时间延迟。
- STAEformer [CIKM23] (Liu et al., 2023)：通过引入时空自适应嵌入组件来增强传统Transformer。
- MultiSPANS [WSDM24] (Zou et al., 2024)：利用多滤波卷积模块、Transformers和结构熵优化来有效建模STG预测中的复杂多范围依赖性。
选择性状态空间模型（SSSM）：
- Mamba (Gu and Dao, 2023)：作为首个深度学习基于SSSM的方法，用于序列任务的数据中心上下文学习。

这些相关研究为STG-Mamba模型的提出提供了理论和技术基础，同时也构成了论文中进行性能比较的基准方法。STG-Mamba通过结合SSSM和GNN的优势，旨在提高STG数据学习的准确性和计算效率。

Q: 论文如何解决这个问题？

STG-Mamba

A: 论文通过提出一种名为STG-Mamba的新模型来解决时空图（STG）数据学习的问题。STG-Mamba模型的核心贡献和解决方案可以概括为以下几点：

将STG网络视为系统：STG-Mamba采用选择性状态空间模型（SSSM）来处理STG网络，将网络视为一个动态系统，并通过模型其在时间维度上的状态演变。
引入图选择性状态空间块（GS3B）：GS3B是STG-Mamba的基本模块，它通过输入依赖的边构建和动态节点特征选择，来捕捉STG网络中复杂的异构依赖关系，并准确建模STG系统的动态演变。
提出卡尔曼滤波图神经网络（KFGN）：KFGN用于自适应的时空图结构生成和升级，它平滑地适应了基于SSSM的上下文，并确保了图结构与STG系统当前状态的同步更新。
编码器-解码器架构：STG-Mamba采用编码器-解码器架构，以GS3B作为基本模块，实现高效的序列数据建模。
计算效率：STG-Mamba通过精心设计的算法，实现了线性时间复杂度O(n)，显著降低了与基于Transformer的方法相比的计算开销，包括浮点运算次数（FLOPs）和测试推理时间。
实证研究：在三个基准STG预测数据集上进行了广泛的实证研究，结果表明STG-Mamba在STG预测性能和计算效率方面均超过了现有的最先进方法。

通过上述方法，STG-Mamba能够有效地处理STG数据的动态、异构和非平稳特性，同时保持了高效的计算性能，为STG学习任务提供了一种新的有前景的解决方案。

GS3算法伪代码

Q: 论文做了哪些实验？

实验数据集

实验结果（12步预测12步）

测试推理时间比较

A: 论文进行了以下几类实验来评估和验证STG-Mamba模型的性能：

数据集选择：选择了三个真实世界的STG数据集进行实验评估，包括：
- PeMS04：加利福尼亚道路网络速度记录。
- HZMetro：杭州城市地铁系统进出记录。
- KnowAir：中国大陆主要城市的气象站记录。
基线方法：为了进行公平比较，论文选择了多种基于图神经网络（STGNN）和基于Transformer的方法作为基线，包括STGCN、STSGCN、STG-NCDE、DDGCRN、ASTGCN、ASTGNN、PDFormer、STAEformer和MultiSPANS等。
实现设置：论文详细描述了模型训练前的数据预处理、模型训练的超参数设置、损失函数和评估指标等。
结果评估与比较：在三个STG预测任务上，论文比较了STG-Mamba与其他基线方法的性能，使用均方根误差（RMSE）、平均绝对误差（MAE）和平均绝对百分比误差（MAPE）作为性能指标，并将结果汇总在表格中。
计算效率评估：为了评估计算效率，论文选择了浮点运算次数（FLOPs）和推理时间作为评估指标，并将STG-Mamba与其他基于Transformer的方法进行了比较。
消融实验：为了探究STG-Mamba中各个模型组件的有效性，论文设计了几种模型变体，并在相同的数据集上评估了它们的预测性能。

通过这些实验，论文全面地展示了STG-Mamba在不同STG预测任务上的性能表现，以及其在计算效率方面相对于其他方法的优势。实验结果证明了STG-Mamba在处理具有动态、异构和非平稳特性的STG数据方面的有效性和高效性。

计算效率比较

不同 STG 节点数设置下，STG-Mamba 和 STAEformer 在 PeMS04/KnowAir 数据集上的 FLOPs 比较

消融实验

Q: 总结一下论文的主要内容

A: 这篇论文的主要内容可以总结如下：

问题定义：论文针对时空图（STG）数据的动态、异构和非平稳特性，提出了一种新的学习框架，旨在解决STG数据分析中的挑战。
方法介绍：提出了STG-Mamba模型，这是首个基于选择性状态空间模型（SSSM）的STG学习方法。STG-Mamba采用编码器-解码器架构，并通过图选择性状态空间块（GS3B）和卡尔曼滤波图神经网络（KFGN）来建模STG数据。
技术创新：
- GS3B：用于STG网络与SSSM的集成，进行输入依赖的边构建和动态节点特征选择。
- KFGN：用于自适应的时空图结构生成和升级，保持与SSSM上下文的平滑适应性。
计算效率：STG-Mamba实现了线性时间复杂度，降低了与基于Transformer的方法相比的计算开销，包括FLOPs和测试推理时间。
实验验证：
- 数据集：使用了三个真实世界的STG数据集进行实验，包括PeMS04、HZMetro和KnowAir。
- 基线方法：与多种现有的STGNN和Transformer基线方法进行了比较。
- 结果：STG-Mamba在预测性能和计算效率方面均超过了现有方法。
消融实验：通过模型变体验证了STG-Mamba中各个组件的有效性，包括KFGN和GS3B。
结论：STG-Mamba为STG学习领域提供了一种新的有前景的方法，它在保持竞争力的模型性能的同时，显著降低了计算成本。