STG-Mamba: Spatial-Temporal Graph Learning via Selective State Space Model

Figure 1 illustrates the comprehensive architecture of the proposed STG-Mamba. Specifically, we formulate the overall architecture following the Residual Encoder fashion for efficient sequential data modeling and prediction. In STG-Mamba, we leverage the Graph Selective State Space Block (GS3B) as the basic Encoder module, and repeat it $N$ times. GS3B consists of several networks and operations, including Layer Norm, Kalman Filter Graph Neural Networks (KFGN) which consists of DynamicFilter-GNN followed by Kalman Filter Upgrading (KF-Upgrading), Spatial-Temporal Selective State Space Module (ST-S3M) which consists of Linear & Split, Conv1D, SiLU, Grapg State Space Selection Mechanism (GSSSM), Element-Wise Concatenation, Graph Information Feed-Forward, as well as the Residual Connection. The Graph Information Feed-Forward structure here is specifically designed for coordinating different modules’ information transmission and updating, ensuring the latest STG information is available for every module.

The motivation for employing Kalman Filtering-based methods in constructing GNNs stems from the need to enhance the reliability and accuracy of spatial-temporal forecasting. STG big data (i.e., traffic sensor records, weather station records, etc) often contains inherent biases and noises that existing methods typically overlook. By integrating Kalman Filtering-based optimization and upgrading, KFGN effectively addresses these inaccuracies by dynamically weighing the reliability of data streams from different temporal granularity, optimizing the fusion of these data streams based on the estimated variances. This approach not only corrects dataset inherent errors but also significantly improves the model’s capacity to capture the complex inter-dependencies within STG patterns.

As depicted in Figure 2, the KFGN pipeline consists of two key steps. In the first step, embeddings generated from the model inputs of different temporal granularities (i.e., recent steps/periodic steps/trend steps) are sent into a DynamicFilter-GNN module for processing. Then, in the second step, the outputs of DynamicFilter-GNN module are integrated and optimized through the Kalman Filtering Upgrading module. In the initial stage of DynamicFilter-GNN, a learnable Base Filter parameter matrix of size $\mathbb{R}^{\text{in\_fea}\times\text{in\_fea}}$ is defined. It’s employed to transform the graph adjacency matrix, thereby dynamically adjusting the connection degrees between nodes. Following that, we uniformly initialize the weights and biases, with an initialization standard deviation stdv calculated as the reciprocal of the square root of the number of input features: $\text{stdv}=\frac{1}{\sqrt{\text{in\_fea}}}$. This uniform initialization is critical for ensuring the module starts from a neutral point.

Subsequently, we transformer the Base Filter via a linear transformation layer, and combine it with the original adjacency matrix $A_{\text{ini}}^{c_{i}}$, resulting in a dynamically adjusted adjacency matrix $A_{\text{dyn}}^{c_{i}}$. Let $c_{i,i=\{1,2,3\}}\in\{r,p,q\}$ represents the input data from different temporal granularities, (i.e., $c_{1}=r$ denotes historical recent, $c_{2}=p$ denotes historical period, $c_{3}=q$ denotes historical trend). Leveraging $A_{\text{dyn}}^{c_{i}}$, input embedding $h_{in}^{c_{i}}$, weights $W_{DF}^{c_{i}}$ and bias $b_{DF}^{c_{i}}$, the input embedding undergoes a graph convolution:

This design enables the model to dynamically adjust the strengths of connections between nodes within the graph based on the learned adjustments.

Following DynamicFilter-GNN, the next step is Kalman Filtering Upgrading. Since the input data are subjected to the same dataset but within different temporal granularities, they are assumed to follow Gaussian Distributions:

where subscript $q=c_{3},p=c_{2},r=c_{1}$ denote historical trend/period/recent dataset, respectively.

We employ Kalman Filtering method to derive accurate information from the multi-granularity observation sets. Specifically, their individual probability distribution functions can be integrated together by multiplication:

To simplify the representation of $\mu_{fuse}$ and $\sigma_{fuse}^{2}$, we hereby introduce parameters $\omega_{q}=\frac{1}{\sigma_{q}^{2}},\omega_{p}=\frac{1}{\sigma_{p}^{2}},\omega_{r}=\frac{1}{\sigma_{r}^{2}}$. Then, $\mu_{fuse}$ and $\sigma_{fuse}^{2}$ can be re-written as:

ST-S3M plays the significant role of adaptive STG feature selection, serving as the counterpart of Attention mechanism. The architecture of ST-S3M is shown in Figure 1. After being processed by KFGN module for dynamic spatiotemporal dependency modeling and statistical-based integration & upgrading. The generated embedding $\hat{h}_{fuse}$ is sent to ST-S3M for input-specific dynamic feature selection. In ST-S3M, the input STG embedding passes through a Linear Layer, followed by a splitting operation. Let $b$ denotes batch size, $l$ denotes sequence length, $d_{model}$ be the feature dimension, $d_{inner}$ be the model inner feature dimension, we have:

where $\hat{h}_{fuse}\in\mathbb{R}^{b\times l\times d_{model}}$, $W_{in}\in\mathbb{R}^{2d_{inner}\times d_{model}}$, $b_{in}\in\mathbb{R}^{2d_{inner}}$, $h_{\text{main-res}}\in\mathbb{R}^{b\times l\times 2d_{inner}}$, $h_{\text{main}},res\in\mathbb{R}^{b\times l\times d_{inner}}$

Then, $h_{\text{main}}$ flows into a 1D convolution layer, followed by SiLU activation function:

The GSSSM in ST-S3M plays a main role in adaptive spatiotemporal feature selection. We detail the parameter computation and update process of the Graph State Space Selection Mechanism in Algorithm 2. Graph Selective Scan Algorithm (Algorithm 1) which receives graph information feed-forward is the most significant step in the state space selection process. We detail it in the following.

We start with obtaining the dimension of the input tensor $u\in\mathbb{R}^{(b,l,d_{in})}$, where $b$ is the batch size, $l$ denotes the sequence length, $d_{in}$ is the input features dim; and the second dim of $\mathbb{A}$ is denoted as $n$.

Next, we highlight the main novelty of the graph selective scan algorithm–the integrated upgrade of graph information Feed-Forward and param $\Delta^{*}$, which corresponds to line 2-line 5 of Algorithm 1. We retrieve the dynamic graph adjacency matrix $\alpha_{t}$from the DynamicFilter-GNN: $\alpha_{t}=\text{DynamicFilter-GNN}.\text{get\_transformed\_adjacency()}$. In order to let the graph information to participate in the state space selection and update process, an intuitive and natural way is to fuse the parameter $\Delta^{*}$, and let the fused parameter flow into the following calculation of $\text{delta}\mathbb{A}$ and $\text{delta}\mathbb{B}_{u}$. However, the dimension of $\alpha_{t}$ and $\Delta^{*}$ may not be identical. Thus, we initialize a padding matrix $\text{adj\_padded}=\mathbf{1}^{d_{\text{in}}\times d_{\text{in}}}$, then we fill the padding matrix with the graph information from $\alpha_{t}$ as: $\text{adj\_padded}[:\alpha_{t}.\text{size}(0),:\alpha_{t}.\text{size}(1)]=\alpha_{t}$. After the dimension adjustment, we therefore integrate $\Delta^{*}$ and adj_padded by performing matrix multiplication: $\Delta^{{}^{\prime}}=\text{matmul}(\Delta^{*},\text{adj\_padded})$.

Following that, we Discretize the continuous parameters $\mathbb{A}$ and $\mathbb{B}$ as:

Then, Iterative State Update is executed on state $x\in\mathbb{R}^{b\times d_{in}\times n}$. For time step from $i$ to $l$, we have:

The Graph Selective Scan algorithm offers several advantages including: enhanced spatiotemporal dependency modeling, adaptability to changing graph structures, and improved forecasting accuracy, making it particularly suitable for spatiotemporal graph learning tasks.

STGNN-based Methods: (I) STGCN Yu et al. (2018), (II) STSGCN Song et al. (2020), (III) STG-NCDE Choi et al. (2022), (IV) DDGCRN Weng et al. (2023).

Attention (Transformer)-based Methods: (V) ASTGCN Guo et al. (2019), (VI) ASTGNN Guo et al. (2021), (VII) PDFormer Jiang et al. (2023), (VIII) STAEformer Liu et al. (2023), (IX) MultiSPANS Zou et al. (2024).

We split all the datasets with a ratio of 6:2:2 for Training/Validation/Testing sets along the time axis, respectively. Before starting model training session, all the data samples are normalized into range $[0,1]$ with MinMax normalization method. We set the batch size as 48 and employ AdamW as the Optimizer. The learning rate is initialized as $1e^{-4}$, with a weight decay of $1e^{-2}$. CosineAnnealingLR Loshchilov and Hutter (2016) is adopted as the learning rate scheduler, with a maximum of 50 iterations, and a minimum learning rate $1e^{-5}$. Model training epochs is set as 100. The number of GS3B Encoder Layer is set as $N=4$. We use the historical ground-truth data from the past 12 steps (for historical recent/period/trend scales) to forecast the future 12 time steps. Mean Squared Error (MSE) is employed as the Loss Function for optimization. All experiments are conducted under one NVIDIA A100 80GB GPU for neural computing acceleration.

We employ RMSE/MAE/MAPE/$std_{\text{MAE}}$ as the deterministic evaluation metrics, all the indicators are smaller the better; $\rm{\Delta_{\text{VAR}}}$ (Change in Variance Accounted For) and $\rm{R^{2}}$ (Coefficient of Determination) as statistic-based evaluation metric, where $\rm{\Delta_{\text{VAR}}}$ is the smaller the better, $\rm{R^{2}}$ is the larger the better. The reported results in Experiment section are the average value of 10 runs. The detailed indicator calculation methods are provided in Appendix B because of limited space.

A comprehensive evaluation between STG-Mamba and the proposed baseline methods is conducted, and we report the RMSE/MAE/MAPE/$std_{\text{MAE}}$ performance results in Table 2. STG-Mamba consistently outperforms other benchmark methods, except for the MAPE criteria of PeMS04 (Flow). It can also be easily observed that Transformer-based methods which integrate GNNs inside the architecture always gain superior performance than sole GNN-based methods. For instance, PDFormer, STAEformer, and MultiSPANS always show better performance than STGCN, STSGCN, and STG-NCDE. The reason is that although GNN is an effective approach of modeling STG tasks, Transformer-based methods have stronger ability to capture local and long-term dynamic dependencies. Within the GNN-based methods, STG-NCDE and DDGCRN stand out for their competitive performance. STG-NCDE models continuous-time dynamics through neural controlled differential equations (NCDE), offering a more precise representation of spatiotemporal feature changes over time compared to other discrete GNN-based methods.

As a new contender poised to challenge Transformer architecture, the selective state space mechanism of Mamba diverges from Attention by utilizing a continuous-time model that captures the dynamic evolution of spatial-temporal dependencies more naturally and efficiently. Unlike Attention, SSSM directly captures the temporal evolution of features, enabling it to scale more gracefully with the sequence length and complexity of STG data. Furthermore, we facilitate STG-Mamba with Kalman Filtering-based adaptive STG evolution, the elaborately designed graph selective scan algorithm, and the Feed-Forward connection, making it inherit the advantages of modern SSSM, and highly suitable for STG learning.

STG data exhibit notable periodicity and diversity, with urban mobility/traffic data showing clear differences between peak travel times in the morning/evening and non-peak periods, as well as between weekdays and weekends. Given these variations due to external environmental changes, it is significant to determine whether deep learning models can effectively model the spatiotemporal dependencies under different conditions. As such, we establish four distinct external scenarios: (a) rush hours on weekdays from 8:00-11:00 AM and 4:00-7:00 PM; (b) non-rush hours during weekdays; (c) weekends (all hours); and (d) non-weekend (all hours). Extensive experiments were conducted on the PeMS04 dataset. Table 3 presents the forecasting results of the four traffic scenarios.

An ideal deep learning-based STG system should be robust to external disturbances, maintaining consistent performance across different scenarios. Compared methods like ASTGNN, PDFormer, and STAEformer show significant performance differences between Rush hour and Non-rush Hour scenarios. Even in the Weekend/Non-Weekend scenario, STG-Mamba demonstrates the smallest performance difference, showing maximum robustness and achieving the best RMSE/MAE/MAPE metrics. In summary, STG-Mamba exhibits superior robustness and outperforms existing baselines in diverse traffic conditions.

Statistical-based evaluation in spatiotemporal forecasting is indispensible as they provide a quantifiable measure of how well a model captures and predicts complex data dynamics over time and space. Specifically, $R^{2}$ and $\Delta_{\text{VAR}}$ help in assessing the accuracy and robustness of models under varying conditions, ensuring the reliability of predictions and effectively inform decision-making processes. The Statistical evaluation results are presented in Table 4. Here, STG-Mamba demonstrates superior performance across datasets, achieving the highest $R^{2}$ values and the lowest $\Delta_{\text{VAR}}$ scores, indicating the advancement of accuracy and efficiency in handling spatiotemporal dependencies.

To investigate the effectiveness of each model component within STG-Mamba, we further design five kinds of model variants, and evaluate their forecasting performance on PeMS04/HZMetro/KnowAir dataset: (I) STG-Mamba: The full STG-Mamba model without any modification. (II) STG-Mamba w/o KF-Upgrading: We replace the KF-Upgrading module in KFGN with the simple sum average operation to fuse the three temporal branches. (III) STG-Mamba w/o Dynamic Filter Linear: We replace the proposed Dynamic Filter Linear in KFGN with the basic static graph convolution. (IV) STG-Mamba w/o GSSSM: The Graph State Space Selection Mechansim (GSSSM) is replaced by the basic SSSM Gu and Dao (2023). (V) STG-Mamba w/o ST-S3M: The whole ST-S3M module is removed from GS3B Encoder.

As illustrated in Figure 3, STG-Mamba w/o KF-Upgrading shows worse performance than the full model, demonstrating the necessity and suitability of employing statistical learning-based GNN state space upgrading and optimization for SSSM-based methods. The decrease of performance due to the absence of Dynamic Filter Linear demonstrate its effectiveness, and the necessity of designing a suitable adaptive GNN for STG feature learning. Furthermore, removing the State Space Selection Mechansim in ST-S3M results in a substantial degradation of model capability, proving the feasibility of using SSSM as the substitute for Attention mechanism. Finally, the removal of ST-S3M makes the STG-Mamba degrade to a plain GNN-based model, resulting in the lowest performance.