STG-Mamba: Spatial-Temporal Graph Learning viaSelective State Space Model

原创已于 2025-10-15 22:18:04 修改 · 1.3k 阅读

18 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #深度学习 #车流量预测 #图论 #神经网络

于 2025-10-11 11:45:53 首次发布

交通流预测论文阅读专栏收录该内容

2 篇文章

订阅专栏

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

研究背景

时空图数据动态、异质、非平稳：典型如交通速度/客流/天气站数据。以往很多方法要么只在图上做卷积，要么用 Transformer 做全局注意力，但时序长度 L 增长时注意力的二次复杂度会成为瓶颈。选择性状态空间模型提供了线性时间的序列建模能力，并强调“基于输入的选择”以聚焦关键信息——非常适合长序列、大图的 STG 学习。

STG-Mamba的核心思想是：用“选择性状态空间模型（Mamba）”做时间上的线性扫描，同时把“动态图结构”直接注入扫描内核，让“时”和“空”在每一步状态更新里耦合起来；再用多时间尺度 + 反方差融合增强鲁棒性与稳定性。

状态空间模型：

用隐状态刻画系统随时间演化，再由输入→状态→输出的一阶线性微分关系建模序列。深度学习里的 SSM 用参数 (A,B,C,D)描述“输入如何驱动状态、状态如何生成输出”。

这里 A,B,C 是把连续系统离散化后的矩阵，用于高效的序列建模。

选择性状态空间模型：

在 SSM 之上加入选择机制（和注意力的核心思想相似），让模型对当前输入自适应地挑选或放大更相关的时序特征。

为什么时间复杂度为O(L)：

计算结构是“选择性扫描”而不是注意力的两两点乘。Mamba把 SSM 离散化后得到一个长度为 L 的结构化卷积核 K=(CB, CAB, …, CAL−1B)，对序列做一次扫描式卷积/递推；没有 L×L的成对相似度计算，所以随序列长度只线性增长。

Notation

图与序列：

（历史）

（目标预测）

任务（历史12步预测未来12步）：

总体框架

STG-Mamba 用残差式编码器堆叠若干个 GS3B（Graph Selective State-Space Block） 作为基本单元。每个 GS3B 依次包含：LayerNorm → KFGN（动态图 + Kalman 式融合）→ ST-S3M（选择性状态空间模块）→ 残差；并配一条“Graph Information Feed-Forward”把最新的动态图信息传递给后续块与扫描核，保证模块间时空信息同步更新。整体目标是学习一个基于 SSSM 的映射

模块A：KFGN（Kalman Filtering Graph Neural Networks）

动机：STG 数据噪声与不确定性强，需要输入自适应的动态图与统计鲁棒融合。KFGN 由两部分组成

1.DynamicFilter-GNN（动态图）

基于可学习滤波器变换原始邻接得到 Adyn，并进行一次线性“图卷积”（空间域图卷积，输入相关的邻接加权），使连边强度随输入时刻/场景而变。它生成的动态图αt会被下游 ST-S3M 的扫描核直接使用。

2.KF-Upgrading（卡尔曼滤波）

STG 数据（交通、气象等）噪声与不确定性强，不同时间尺度的信号“可信度”不同。卡尔曼滤波被用作“统计加权融合器”，把三个时间尺度（近期/周期/趋势）的图表示按“方差的倒数”加权融合，抑制噪声与偏差，得到更稳、更准的时空图特征，再交给选择性状态空间模块（Mamba）做后续选择/更新。

将近/周期/趋势三路时序表征按反方差权进行融合（谁不确定性小谁权重大），并引入可学习系数 ϵ,ϕ：

实际训练里为简化数值稳定性，去掉了公共分母（常数项）。方差通过样本统计近似估计，保证高效可扩展。

模块 B：ST-S3M（Spatial-Temporal Selective State-Space Module）

ST-S3M 把动态图信息（来自前面的 KFGN）通过Feed-Forward送进选择性状态空间的选择/更新里；具体做法是用一个图选择性扫描算法（Graph Selective Scan, GSSSM），把邻接矩阵融入 Mamba 的选择性扫描参数与离散化的 A,B计算，从而完成图感知的状态迁移与控制。而普通 Mamba 只对时序做选择性卷积，不具备这种图信息注入与图结构感知的参数更新。

作为 SSSM 版的“注意力”，根据输入做动态特征选择，并以线性时间复杂度进行时序建模。流程：
Linear & Split → 1D Conv → SiLU → GSSSM →（与残差 res 点乘门控）→ Linear 输出

GSSSM图嵌入

1.准备参数：选择机制产出的 Δ 与 KFGN 的动态图 αt 维度通常不同，先做padding 对齐后相乘：

2.离散化更新：用爱因斯坦求和计算离散矩阵（与 Mamba/S4 的离散化一致）：

3.时序递推（长度L）：

这使每一步状态转移都被动态图调制，实现真正的时空耦合；同时保持对序列长度的线性复杂度 O(L)。

数据集

结果、性能分析

鲁棒性分析

统计评价结果

消融实验

总结

STG-Mamba首次引入了基于深度学习的选择性状态空间模型（SSSM）用于时空图学习任务。我们提出了STG-Mamba，它利用现代SSSM进行准确有效的STG预测。在STG-Mamba中，ST-S3 M模块促进了输入依赖图进化和特征选择，成功地将STG网络与SSSM集成在一起。

通过卡尔曼滤波图神经网络（Kalman Filtering Graph Neural Networks，KFGN），学习的STG嵌入实现了基于统计学习的平滑优化升级，与整个STG选择性状态空间建模过程保持一致。与基于注意力的方法相比，STG-Mamba实现了线性时间复杂度，并大幅降低了FLOP和推理时间。我们相信，STG-Mamba提出了一个全新的有前途的方法，一般STG学习领域，提供有竞争力的模型性能和输入依赖的上下文特征选择下负担得起的计算成本。