论文简介

本推文主要介绍了一种基于选择性状态空间模型的时空图学习方法，论文标题为《STG-Mamba: Spatial-Temporal Graph Learning via Selective State Space Model》，第一作者为Lincan Li。时空图数据是一类广泛存在于日常生活中的非欧几里得数据，时空图学习可以被视为理解和预测时空图网络随时间演变的复杂过程。由于时空图数据的动态性、异质性和非平稳性，如何对其进行准确预测一直是一个挑战性任务。为了解决这一问题，论文提出了STG-Mamba模型，该模型是首个基于深度学习的选择性状态空间模型，专注于数据驱动的时空图学习。实验结果表明，该方法在多个数据集上的表现均优于当前主流的时空预测模型，同时显著降低了计算成本。推文作者为朱旺，审校为李杨和陆新颖。

论文下载地址：

https://arxiv.org/abs/2403.12418

开源代码地址：

https://github.com/LincanLi98/STG-Mamba

01 背景与挑战

时空图数据是一种广泛存在于我们日常生活中的非欧几里得数据，如城市交通网络、地铁系统流入/流出、社会网络、区域能源负荷、天气观测等。时空图学习可以被视为理解和预测时空图随时间演变的复杂过程。目前在时空图预测方面仍存在以下问题：

（1）时空图数据预测的挑战：由于时空图数据具有动态性、异质性和非平稳性的特点，为时空图学习带来了持续性的挑战。

（2）现有模型的局限性：大多数现有的基于图神经网络的方法仅关注于模拟时空图网络中节点之间的关系，而忽视了对时空图系统随时间变化的内在特征进行建模的重要性。

选择性状态空间模型提出了一种新的方法，将网络视为一个系统，从而能够更细致地探索时空图系统在时间维度上的动态演化。基于上述挑战和机遇，论文首次将现代选择性状态空间模型应用于时空图学习，通过将时空图网络视为一个系统，深入探索其在时间维度上的动态状态演变。

02 方法

​图1 STG-Mamba的整体框架

论文针对时空图数据在时间维度上的动态变化、异构性和非平稳的特性，采用堆叠残差编码器的方式设计了STG-Mamba模型框架，首次将选择性状态空间模型应用于时空图学习。如图1所示，该框架堆叠了多个图选择状态空间块作为骨干网络，并引入卡尔曼滤波图神经网络、时空图选择性状态空间模块以及同步时空图前馈连接，来序列化和协调不同的内部模块。

​

图2 卡尔曼滤波图神经网络

如图2所示，在卡尔曼滤波图神经网络（Kalman Filter Graph Neural Networks，KFGNN）模块主要由两个关键步骤组成。在第一步中，动态滤波图神经网络通过预定义的可学习参数矩阵来动态调整节点之间的连接度，将不同时间粒度（近期数据、周期性数据和趋势性数据）的输入进行自适应聚合嵌入。在第二步中，通过统计理论的学习进行集成和优化，统一初始化权重和偏差，加权融合多时间粒度的数据，从而捕获更全面准确的特征。这种方法有效缓解了时空图数据中固有的偏差和噪声，提高了预测的可靠性和准确性。接着，论文设计了一个基于图扫描算法的时空选择性状态空间模块。该模块通过爱因斯坦求和来更新状态转移矩阵和控制矩阵，从而在不同的时间步上迭代更新状态。这样的设计使得选择过程能够动态适应输入特征的变化，实现对输入数据的逐步优化。

03 实验及结果

3.1 数据集

时空图数据是一种广泛存在于我们日常生活中的非欧几里得数据，如城市交通网络、地铁系统流入/流出、社会网络、区域能源负荷、天气观测等。论文在三个公开的基准时空图预测数据集上进行了广泛的实验：加利福尼亚交通数据集（PeMS04）、杭州地铁系统进出站记录（HZMetro）和中国主要城市天气记录（KnowAir）。数据集的详细描述如表1所示：

表1 三个数据集的统计和描述

​

3.2 对比试验

论文将STG-Mamba和基于图神网络以及Transformer的九种基线模型进行了对比试验。实验性能结果如表2所示。结果表明，STG-Mamba在大多数评价指标上均优于其他基准方法。此外，由于Transformer具有更强捕获局部和长期动态依赖关系的能力，集成Transformer的模型明显优于单一图神经网络方法。而作为Transformer的挑战者，Mamba的选择性状态空间机制能直接捕获特征的时间演变，在保证准确率的同时，显著减少了计算开销，尤其在大规模网络中展现了线性时间复杂度的优势。

表2 性能评估与基线比较的结果

​

3.3 鲁棒性实验

时空图数据具有明显的周期性和多样性，城市交通早晚高峰以及工作日与周末出行存在明显差异。考虑到这些外部环境引起的变化，确认所提出的模型是否能有效的模拟不同条件下的时空依赖关系，论文针对早晚高峰以及工作日与周末进行了鲁棒性实验。表3为四种情景下的预测结果。

表3 不同时段下的鲁棒性实验

通过在高峰期、非高峰期、周末和工作日等多种场景下进行测试，STG-Mamba展示了较强的鲁棒性。在外部环境变化较大的情况下，模型依然能够保持稳定的预测性能，体现了其在多种复杂时空条件下的适应能力。

04 结论

STG-Mamba首次将选择性状态空间模型引入时空图学习，通过结合选择性状态空间模型、卡尔曼滤波图神经网络和自适应图特征选择机制，提供了一种在复杂时空场景下高效、准确的时空图学习方法。相比于传统基于注意力机制的Transformer模型，STG-Mamba在计算复杂度和模型表现方面都具有明显的优势。未来的研究可以进一步扩展其在更广泛数据集上的适用性，并增强模型的可解释性。