[논문 리뷰] STG-Mamba: Spatial-Temporal Graph Learning via Selective State Space Model

Lincan Li, Hanchen Wang외 2명|arXiv (Cornell University)|2024. 03. 19.

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 학습을 Encoder-Decoder로 형식화하고 스택된 Graph Selective State Space Blocks (GS3B)로 구현한다.
• Kalman Filtering Graph Neural Networks (KFGN)을 통합하여 적응적이고 시간에 따라 변화하는 그래프 인접 행렬을 생성한다.
• 입력 의존 특성 학습을 위한 Graph State Space Layer (GSS-Layer)를 사용한다.
• 연속 상태 공간 동역학을 0차 보정(zero-order hold)으로 이산화하여 효율적인 딥 러닝 구현을 가능하게 한다.
• 시퀀스 길이 O(L)에 대한 선형 시간 복잡도와 Transformer 기준선 대비 감소된 FLOPs를 보인다.

실험 결과

연구 질문

• RQ1 선택적 상태 공간 모델이 동적 STG 의존성을 효과적으로 포착할 수 있는가?
• RQ2 적응형 그래프 학습을 SSSMs와 STG 데이터에 어떻게 통합할 수 있는가?
• RQ3 GS3B/KFGN 기반의 STG-Mamba가 Transformer 기반 STG 모델에 비해 더 낮은 계산 비용으로 우수한 예측 정확도를 달성하는가?

주요 결과

• STG-Mamba는 RMSE, MAE, MAPE 전 항목에서 PeMS04, HZMetro, KnowAir에 대해 기준선보다 지속적으로 우수하며, 한 지표에서 소폭 예외가 있다.
• STG-Mamba는 STG 규모에 대해 선형 FLOP 증가를 보이고 비교 대상 Transformer들보다 추론 시간이 짧아 주목할 만한 계산 효율성을 달성한다.
• 데이터셋 전반에서 STG-Mamba는 대부분의 작업에서 최고의 RMSE와 MAE를 달성한다(예: PeMS04: 29.53 RMSE, 18.09 MAE; KnowAir: 18.05 RMSE, 11.73 MAE).
• STG-Mamba는 외부 조건(러시 시간대 대 비러시 시간대, 주말 대 평일) 변화에도 강건한 성능을 보인다.
• 더 큰 STG 규모에서 STAEformer에 비해 FLOPs를 크게 감소시킨다(예: 300 노드인 PeMS04에서 36.85G 대 92.49G).