[論文レビュー] 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}.

研究の動機と目的

• 動的で異種な空間-時間グラフへ STG 予測を動機づける。
• STG ダイナミクスをモデル化するためのグラフ選択的状態空間アーキテクチャを提案する。
• 進化する STG 構造を扱う適応的グラフ学習を開発する。
• Transformer ベースの STG モデルに対するスケーラビリティと効率性の利点を示す。

提案手法

• STG 学習を積み重ねられた Graph Selective State Space Blocks (GS3B) を備えたエンコーダ-デコーダとして定式化する。
• 適応的で時変するグラフ隣接行列を生成するために Kalman Filtering Graph Neural Networks (KFGN) を統合する。
• 入力依存の特徴学習のため Graph State Space Layer (GSS-Layer) を採用する。
• 効率的なディープラーニング実装を可能にするために連続状態空間ダイナミクスをゼロ次ホールドで離散化する。
• シーケンス長 O(L) における線形時間計算量と Transformer ベースラインと比較した FLOPs の削減を示す。）

実験結果

リサーチクエスチョン

• RQ1 選択的状態空間モデルは動的 STG の依存関係を効果的に捉えられるか？
• RQ2 STG データに対して適応的グラフ学習を SSSM とどのように統合できるか？
• RQ3 GS3B/KFGN ベースの STG-Mamba は Transformer ベースの STG モデルと比較して、計算コストを抑えつつ予測精度が優れるか？

主な発見

モデル	PeMS04 (Flow) RMSE	PeMS04 (Flow) MAE	PeMS04 (Flow) MAPE	HZMetro RMSE	HZMetro MAE	HZMetro MAPE	KnowAir RMSE	KnowAir MAE	KnowAir MAPE
STGCN	35.55	22.70	14.59	34.85	21.33	13.47	20.64	15.27	12.92
STSGCN	33.65	21.19	13.90	33.24	20.79	13.06	20.25	15.22	12.69
STG-NCDE	31.09	19.21	12.76	32.91	20.75	12.88	20.08	14.93	11.87
DDGCRN	30.51	18.45	12.19	31.69	19.52	12.45	19.63	12.45	11.84
ASTGCN	35.22	22.93	16.56	34.36	21.12	13.50	20.37	15.07	12.11
ASTGNN	31.16	19.26	12.65	32.75	20.63	12.47	19.96	12.38	11.23
PDFormer	29.97	18.32	12.10	30.18	19.13	11.92	18.65	12.21	11.06
STAEformer	30.18	18.22	11.98	29.94	18.85	12.03	18.49	11.93	10.75
MultiSPANS	30.46	19.07	13.29	30.31	18.97	11.85	18.57	12.24	10.88
STG-Mamba	29.53	18.09	12.12	29.26	18.38	11.67	18.05	11.73	10.66

• STG-Mamba は PeMS04、HZMetro、KnowAir の各データセットで RMSE、MAE、MAPE の全指標でベースラインを一貫して上回り、1つの指標を除いてわずかな例外あり。
• STG-Mamba は STG サイズに対して線形に FLOP が増加する顕著な計算効率を達成し、比較対象の Transformer より推論時間が短い。
• データセットを問わず、多くのタスクで最良の RMSE と MAE を達成（例: PeMS04: RMSE 29.53、MAE 18.09; KnowAir: RMSE 18.05、MAE 11.73）。
• 外部条件（混雑時間帯と非混雑時間帯、週末と平日）を変えても頑健な性能を示す。
• 大規模な STG スケールでは STAEformer に比べて FLOPs を大幅に削減（例: 300 ノードの PeMS04: 36.85G 対 92.49G）。