1 Introduction↩︎

Spatio-temporal forecasting serves as a foundational pillar within the AI4Transport ecosystem. It drives intelligent transportation technologies [1]–[3]. By predicting future road traffic conditions from historical data, AI agents enable services such as autonomous traffic management, route optimization, and road safety support [4]–[6]. Traffic signals exhibit spatial heterogeneity and complex temporal dynamics. Spatially, interactions between regions are intricate and non-Euclidean. Temporally, signals exhibit strong periodicity entangled with abrupt changes and non-stationary patterns. These properties motivate forecasting models that can capture both spatial dependencies and temporal variations.

To address these challenges, researchers have developed sophisticated architectures to explicitly model these dependencies. As illustrated in Figure 1, Graph-based methods perform forecasting by assuming predefined or learned graph structures. These structures capture spatial correlations among nodes [7]. Representative works such as STGCN [8] and Graph WaveNet [9] leverage graph convolutions and adaptive adjacency matrices. Meanwhile, Transformer-based methods utilize attention mechanisms. They model long-range dependencies across spatio-temporal tokens [10]. Additionally, spectral decomposition techniques like StemGNN [11] and STWave [12] focus on frequency-domain transforms. Frameworks like STID [13] utilize identity-based embeddings. These decoupled methods aim to obtain latent spatio-temporal representations that support more accurate predictions.

Despite diverse structural designs, many spatio-temporal architectures employ MLPs with fixed activation functions as their default nonlinear mapping modules. Although MLPs are expressive general-purpose approximators, their nonlinearities are usually selected in advance and shared across units. This observation motivates a complementary research question: beyond architectural design, to what extent does the choice of nonlinear function approximator affect spatio-temporal forecasting performance?

In this work, we investigate whether the choice of nonlinear function approximator constitutes an underexplored factor in spatio-temporal forecasting performance. We propose STKAN, a Kolmogorov–Arnold Network (KAN)-based architecture that decomposes spatial and temporal dependencies while introducing TaylorKAN layers into the spatial and temporal token-mixing mappings. STKAN provides a controlled architectural setting for examining whether Taylor-polynomial KAN token mixers can complement conventional spatio-temporal modeling components. The channel mixer and Transformer feed-forward network remain MLP-based, so the model does not replace all MLP components.

Our main contributions are summarized as follows:

• We investigate the role of nonlinear function-approximator design in spatio-temporal forecasting and introduce STKAN, which incorporates Taylor-polynomial KAN mappings into spatial and temporal token-mixing modules.

• We develop a learnable soft node-group assignment mechanism together with dedicated spatial and temporal mixing blocks. The grouping mechanism constructs compact spatial representations, while the two token mixers model interactions along the spatial-group and temporal dimensions, respectively.

• We evaluate STKAN on five traffic forecasting benchmarks. The results demonstrate competitive forecasting accuracy, while the existing ablation results suggest that the TaylorKAN token mixers, adaptive spatial grouping, and attention components each contribute to the overall model in the evaluated configurations.

2 Related Work↩︎

2.1 Spatio-Temporal Forecasting↩︎

Spatio-temporal forecasting extends traditional time-series forecasting by incorporating both temporal dynamics and spatial dependencies, such as in traffic management, where multiple traffic sensors’ data is used to predict future conditions. Early deep learning approaches combined Convolutional Neural Networks (CNNs) and Recurrent Neural Networks (RNNs) to capture spatial and temporal dependencies [14]–[16]. However, grid-based CNNs may not effectively handle non-Euclidean spatial relationships, leading to the development of Graph Convolutional Networks (GCNs) [17], [18] and spatio-temporal Graph Neural Networks (STGNNs) [8], [19], [20]. These models, such as DCRNN [19], ST-MetaNet [21], and DGCRN [20], integrate GCNs with RNNs [22], while others like Graph WaveNet [9] and STGCN [8] combine GCNs with gated Temporal Convolutional Networks (TCNs). Attention mechanisms have also been widely adopted in STGNNs [23]. Some studies criticize the reliance on pre-defined graphs, suggesting alternatives like AGCRN [24] and MTGNN [25], which learn latent graph structures. Recent non-GCN solutions, such as STNorm [26] and STID [13], further highlight the need for a better understanding of spatial dependencies in forecasting tasks.

2.2 Kolmogorov-Arnold Network↩︎

KANs are motivated by the Kolmogorov–Arnold representation theorem [27], which represents multivariate functions through compositions of univariate functions. Building upon this paradigm, several variants have been proposed, including wavelet-based WavKAN [28], Taylor polynomial-based TaylorKAN [29], fKAN with trainable Jacobi basis functions [30], and FastKAN [31], which employs radial basis function approximations of B-spline bases.

Although KANs are motivated by the Kolmogorov–Arnold representation theorem, the theorem itself does not imply that a finite KAN model will universally outperform an MLP in empirical forecasting tasks. In this work, KANs are therefore treated as an alternative nonlinear parameterization with a different inductive bias, rather than as a theoretically guaranteed replacement for MLPs.

KANs have recently attracted attention in time series forecasting. Prior work has explored KAN-based models for time series analysis [32], mixture-of-experts KAN designs [33], symbolic-regression-oriented T-KAN and MT-KAN models [34], collaborative time-frequency learning with iTFKAN [35], and frequency-decomposition learning with TimeKAN [36]. BiLSTM-KAN integrates KAN layers with a bidirectional LSTM backbone for traffic flow forecasting [37]. STKAN differs from these studies by replacing the token-mixing mappings in its spatial and temporal mixer blocks with Taylor-polynomial KAN layers, while retaining conventional MLPs in the channel-mixing and Transformer feed-forward components.

3 Preliminary↩︎

Spatio-temporal forecasting is a specialized multivariate time-series forecasting problem. Given a historical input sequence \(\mathbf{X}_{t-L_{\mathrm{in}}+1:t}=[\mathbf{X}_{t-L_{\mathrm{in}}+1},\ldots,\mathbf{X}_{t}]\in\mathbb{R}^{L_{\mathrm{in}}\times N\times C_{\mathrm{in}}}\), the goal is to predict \(\widehat{\mathbf{Y}}_{t+1:t+H}=[\widehat{\mathbf{Y}}_{t+1},\ldots,\widehat{\mathbf{Y}}_{t+H}]\in\mathbb{R}^{H\times N\times C_{\mathrm{out}}}\). Here, \(L_{\mathrm{in}}\) denotes the historical input length, \(H\) denotes the prediction horizon, \(N\) is the number of spatial nodes, and \(C_{\mathrm{in}}\) and \(C_{\mathrm{out}}\) are the input and output channel dimensions.

4 Methodology↩︎

In this paper, we propose STKAN to model spatial interactions and temporal dynamics by introducing TaylorKAN token mixers into a spatio-temporal forecasting architecture. The overall architecture of STKAN is shown in Figure 2.

4.1 Embedding Layers↩︎

To capture spatio-temporal dependencies in traffic sequences, we embed the raw input \(\mathbf{X}\in\mathbb{R}^{L_{\mathrm{in}}\times N\times C_{\mathrm{in}}}\) with a fully connected layer as \(\mathbf{E}_f=\mathrm{FC}(\mathbf{X})\in\mathbb{R}^{L_{\mathrm{in}}\times N\times d_f}\). We use a learnable spatial embedding \(\mathbf{E}_s\in\mathbb{R}^{N\times d_s}\). The spatial embedding is shared across time steps and is broadcast along the temporal dimension before concatenation. To incorporate temporal periodicity, we use two embedding dictionaries: a time-of-day embedding with \(T_{\mathrm{tod}}=288\) slots for five-minute data and a day-of-week embedding with \(T_{\mathrm{dow}}=7\) categories. The resulting embeddings are denoted as \(\mathbf{E}_{\mathrm{tod}}\) and \(\mathbf{E}_{\mathrm{dow}}\) after lookup and broadcasting to the node dimension.

4.2 Temporal Convolution Blocks↩︎

All embeddings are concatenated along the feature dimension to form the initial representation: \[\mathbf{X}^{(0)} = \mathbf{E}_f \Vert \mathbf{E}_s \Vert \mathbf{E}_{\mathrm{tod}} \Vert \mathbf{E}_{\mathrm{dow}} \in \mathbb{R}^{L_{\mathrm{in}}\times N\times d_h}, \label{eq:initial95embedding}\tag{1}\] where \(d_h=d_f+d_s+d_{\mathrm{tod}}+d_{\mathrm{dow}}\). To capture local temporal context while reducing the temporal resolution, we define a patch-extraction operator that applies a \(1\times w\) convolution along the time axis with stride \(s\): \[\mathbf{X}_p = \mathrm{PatchConv} \left( \mathbf{X}^{(0)};w,s \right) \in \mathbb{R}^{L_p\times N\times d_h}, \label{eq:patch95conv}\tag{2}\] where \(L_p=\left\lfloor\frac{L_{\mathrm{in}}-w}{s}\right\rfloor+1\). Thus, \(\mathbf{X}_p\) serves as a patch-level spatio-temporal feature map for subsequent mixer and attention blocks.

4.3 Spatio-Temporal KAN Blocks↩︎

The STKAN blocks combine learnable spatial grouping, spatial TaylorKAN token mixing, and temporal TaylorKAN token mixing. STKAN first uses a soft assignment matrix to aggregate raw nodes into macro-level spatial tokens. The aggregated representations are then processed by the Spatial KAN (SKAN) block. The spatial output is mapped back to node-level features and passed to the Temporal KAN (TKAN) block, which models dependencies over the compressed temporal dimension.

Learnable Spatial Grouping. To transform the original spatial nodes into macro-level spatial tokens, we introduce an unnormalized learnable parameter matrix \(\mathbf{A}\in\mathbb{R}^{N\times G}\). The assignment matrix is obtained by applying softmax along the group dimension: \[\mathbf{G}_{n,g} = \frac{\exp(\mathbf{A}_{n,g})}{\sum_{g'=1}^{G}\exp(\mathbf{A}_{n,g'})}. \label{eq:group95softmax}\tag{3}\] Thus, \(\sum_{g=1}^{G}\mathbf{G}_{n,g}=1\) for each node \(n\). Given \(\mathbf{X}_p\in\mathbb{R}^{L_p\times N\times d_h}\), the group-level representation \(\widetilde{\mathbf{X}}\in\mathbb{R}^{L_p\times G\times d_h}\) is computed as \[\widetilde{\mathbf{X}}_{\tau,g,c} = \sum_{n=1}^{N} \mathbf{G}_{n,g} \mathbf{X}_{p,\tau,n,c}. \label{eq:spatial95grouping}\tag{4}\]

Spatial KAN Block. The SKAN block follows a mixer-style design. Although MLPs are expressive general-purpose approximators, their activation functions are usually fixed before training. KAN-based mappings provide a different parameterization in which learnable univariate functions are associated with network connections. The transmission from the \(j\)-th neuron in layer \(\ell+1\) to all neurons in layer \(\ell\) is formulated as: \[z_{\ell+1, j} = \sum_{i=1}^{n_\ell} \phi_{\ell, j, i}(z_{\ell, i}),\] where \(z_{\ell, i}\) is the \(i\)-th neuron in layer \(\ell\), \(n_\ell\) is the number of neurons in that layer, and \(\phi_{\ell, j, i}(\cdot)\) is a learnable univariate mapping.

In our model, \(\phi\) is instantiated by a TaylorKAN layer. A Taylor expansion motivates the polynomial basis: \[f(x) \approx \sum_{r=0}^{K} \frac{f^{(r)}(0)}{r!}x^r.\] The implemented layer uses learnable coefficients: \[\phi_q(\mathbf{x}) = \sum_{p=1}^{C} \sum_{r=0}^{K} \theta_{q,p,r}x_p^r+b_q. \label{eq:taylor95layer}\tag{5}\] The coefficients \(\theta_{q,p,r}\) are optimized directly and are not constrained to equal the analytical derivatives of an underlying function. Accordingly, the adopted layer is more precisely interpreted as a Taylor-inspired polynomial-basis KAN layer. The polynomial basis allows the mapping to combine nonlinear terms of different orders. Its effectiveness for spatial and temporal mixing is evaluated empirically in the subsequent experiments.

To avoid ambiguity when mixing tensor dimensions, we define a spatial permutation operator \(\mathcal{P}_s:\mathbb{R}^{L_p\times G\times d_h}\rightarrow\mathbb{R}^{L_p\times d_h\times G}\). The spatial token-mixing stage is then \[\mathbf{U}_s = \mathcal{P}_s(\widetilde{\mathbf{X}}) + \mathrm{TaylorKAN}_s \left( \mathrm{LN} \left( \mathcal{P}_s(\widetilde{\mathbf{X}}) \right) \right), \label{eq:skan95token}\tag{6}\] where \(\mathbf{U}_s\in\mathbb{R}^{L_p\times d_h\times G}\). We restore the original group-feature order by \[\overline{\mathbf{U}}_s = \mathcal{P}_s^{-1}(\mathbf{U}_s) \in \mathbb{R}^{L_p\times G\times d_h}. \label{eq:skan95restore}\tag{7}\] The channel mixer remains an MLP applied to the feature dimension: \[\mathbf{V}_s = \overline{\mathbf{U}}_s + \mathbf{W}_2 \sigma \left( \mathbf{W}_1 \mathrm{LN} (\overline{\mathbf{U}}_s) + \mathbf{b}_1 \right) + \mathbf{b}_2, \label{eq:skan95channel}\tag{8}\] where \(\mathbf{V}_s\in\mathbb{R}^{L_p\times G\times d_h}\) and \(\sigma\) denotes the GELU activation. The node-level spatial output is \[\mathbf{S}_{\tau,n,c} = \mathbf{X}_{p,\tau,n,c} + \sum_{g=1}^{G} \mathbf{G}_{n,g} \mathbf{V}_{s,\tau,g,c}. \label{eq:skan95node95restore}\tag{9}\]

Temporal KAN Blocks. The TKAN block applies TaylorKAN token mixing along the compressed temporal dimension. We define a temporal permutation operator \(\mathcal{P}_t:\mathbb{R}^{L_p\times N\times d_h}\rightarrow\mathbb{R}^{N\times d_h\times L_p}\). The temporal token mixer is \[\mathbf{U}_t = \mathcal{P}_t(\mathbf{S}) + \mathrm{TaylorKAN}_t \left( \mathrm{LN} \left( \mathcal{P}_t(\mathbf{S}) \right) \right), \label{eq:tkan95token}\tag{10}\] where \(\mathbf{U}_t\in\mathbb{R}^{N\times d_h\times L_p}\). The channel mixer is written as \[\mathbf{V}_t = \mathcal{P}_t^{-1}(\mathbf{U}_t) + \mathrm{MLP}_t \left( \mathrm{LN} \left( \mathcal{P}_t^{-1}(\mathbf{U}_t) \right) \right), \label{eq:tkan95channel}\tag{11}\] and the final temporal representation is \[\mathbf{T} = \mathbf{V}_t \in \mathbb{R}^{L_p\times N\times d_h}. \label{eq:tkan95output}\tag{12}\]

4.4 Spatio-Temporal Fusion Blocks↩︎

To complement the TaylorKAN token mixers, we introduce Transformer layers along both spatial and temporal axes. Given a hidden representation \(\mathbf{T}\in\mathbb{R}^{L_p\times N\times d_h}\), the attention block projects it into query, key, and value matrices: \[\mathbf{Q} = \mathbf{T}\mathbf{W}_Q,\quad \mathbf{K} = \mathbf{T}\mathbf{W}_K,\quad \mathbf{V} = \mathbf{T}\mathbf{W}_V,\] where \(\mathbf{W}_Q,\mathbf{W}_K,\mathbf{W}_V\in\mathbb{R}^{d_h\times d_h}\). The scaled dot-product attention is \[\mathrm{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}) = \mathrm{Softmax} \left( \frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d_k}} \right) \mathbf{V}. \label{eq:attention}\tag{13}\] Spatial attention is computed independently at each time step along the node dimension, so the softmax operates over the \(N\) spatial nodes. Temporal attention is computed independently at each node along the compressed temporal dimension, so the softmax operates over the \(L_p\) temporal tokens. Figure 2 indicates that the attention block contains LayerNorm, FFN, and residual connections in addition to the attention operation.

4.5 Prediction Head↩︎

After feature extraction and fusion via STKAN blocks, the prediction head aggregates temporal information across compressed time steps for each node and applies a linear projection to generate multi-step forecasts. Let \(\mathbf{Z}\) denote the output representation after the fusion block. The prediction head is expressed as:

\[\widehat{\mathbf{Y}}_{t+1:t+H} = f_{\mathrm{out}}(\mathbf{Z}), \label{eq:stid95prediction}\tag{14}\] where \(\widehat{\mathbf{Y}}_{t+1:t+H}\in\mathbb{R}^{H\times N\times C_{\mathrm{out}}}\). The output function \(f_{\mathrm{out}}\) denotes the reshape and linear projection operations used to map the fused representation to the future horizon.

5 Experiment↩︎

5.1 Experimental Setup↩︎

Datasets. We evaluate our model on five traffic forecasting datasets, including PEMS04, PEMS07, PEMS08, PEMS-BAY and METR-LA. Following previous work, we divide the PEMS04, PEMS07 and PEMS08 dataset into training, validation, and test sets in a ratio of 6:2:2. For the remaining datasets, we adopt a split ratio of 7:1:2. Detailed statistics of these datasets are shown in Table 1.

Baselines. We compare 11 representative baselines with our proposed STKAN. (i) Non-spatial modeling-based: STID [13], which adopts identity spatio-temporal embeddings and avoids explicit spatial dependency modeling. (ii) Static spatial-based methods: STGCN [8], GWNet [9], AGCRN [24], GMAN [23], MTGNN [25] and STDN [38] combine pre-defined or learned static graph structures with temporal modeling modules. (iii) Dynamic spatial-based methods: STAEformer [10] and STWave [12] capture time-varying spatial dependencies through adaptive or attention-based mechanisms. (iv) Spatio-temporal decomposition-based: StemGNN [11] and STNorm [26] decompose spatio-temporal series into separate components for modeling, focusing on disentangling spatial and temporal patterns.

Settings. In the experiments, we use the traffic flow of the last 12 time steps to predict the traffic flow of the next 12 time steps, and record the prediction performance of the 3rd, 6th, 12th steps and the average. We set the Adam optimizer with an initial learning rate of 0.002, where the learning rate follows a step-wise decay strategy, and the batch size is set as 64. During the training phase, we employ the early stopping strategy with tolerance 30 for 200 epochs. For performance evaluation, we adopt three widely used metrics to quantify the accuracy of traffic forecasting results: Mean Absolute Error (MAE), Root Mean Squared Error (RMSE), and Mean Absolute Percentage Error (MAPE). All baseline methods are implemented within the BasicTS [39] framework to ensure consistent training pipelines and fair performance comparison.

5.2 Performance Comparisons↩︎

Table [tbl:tab:normal] reports the forecasting results on the five benchmark datasets. STKAN achieves the best average MAE and RMSE on PEMS04, although its average MAPE is slightly higher than that of STWave. On PEMS07 and PEMS08, STKAN obtains the best average results across the three reported metrics. On PEMS-BAY, STKAN ties with STAEformer in average MAE and obtains the lowest average RMSE and MAPE. On METR-LA, STKAN remains competitive but does not outperform the strongest baseline in average MAE or RMSE. Overall, the results indicate that STKAN is particularly effective on the evaluated traffic-flow datasets, whereas its advantage is less evident on the METR-LA speed dataset.

The gains over the best-performing baselines are generally modest. Since the current evaluation reports single-run results, small numerical differences should be interpreted cautiously. These results demonstrate the competitiveness of the overall STKAN architecture. The comparison with the MLP variant further suggests that the adopted TaylorKAN token mixers are useful within the evaluated configuration, although the current experiments do not isolate function-approximator capacity as the sole source of the observed gains.

5.3 Ablation Study↩︎

Effectiveness of KAN Modules. To evaluate the role of KAN components within STKAN, we design three model variants:

• MLPs: replacing the TaylorKAN token mixers with MLP layers in the evaluated implementation.

• w/o SKAN: removing the spatial block while retaining TKAN to test the importance of inter-node mixing.

• w/o TKAN: removing the temporal block while retaining the spatial modeling component.

As shown in Figure 3, replacing the TaylorKAN token mixers with the evaluated MLP implementation increases the forecasting errors on PEMS04 and PEMS08. This result suggests that the Taylor-polynomial mappings contribute positively within the current architecture and hyperparameter setting. Removing either SKAN or TKAN also degrades at least part of the reported performance, indicating that spatial-group mixing and temporal mixing provide complementary information. These observations are specific to the evaluated model configurations and should not be interpreted as establishing the general advantage of KANs over all parameter-matched MLP alternatives.

Effectiveness of Attention Mechanisms. To evaluate the role of attention in STKAN, we conduct ablation studies on PEMS04 and PEMS08 with the following variants:

• w/o S-Attention: disabling spatial attention while preserving temporal modeling.

• w/o T-Attention: removing temporal attention while keeping spatial modeling.

• w/o ST-Attention: removing both attention modules, leaving only KAN-based token and channel mixing.

Table 3 shows that the complete model achieves the best result for every reported metric on PEMS04 and PEMS08. However, the relative influence of spatial and temporal attention varies across metrics. For example, removing temporal attention produces the largest MAE degradation on PEMS08, whereas removing spatial attention produces the largest RMSE degradation. Removing both branches is therefore not uniformly the worst variant for every metric. The results indicate that the two attention branches contribute differently across datasets and evaluation criteria.

The ablation results suggest that attention provides an additional refinement over the KAN-based mixer blocks, but the current experiments do not establish a strict ranking between the contributions of attention and KAN components.

3pt

5.4 Hyper-parameter Study↩︎

Table 4 reports the sensitivity of STKAN to the number of spatial groups \(G\). The results exhibit a non-monotonic pattern: using either fewer or more groups than the selected value leads to slightly higher forecasting errors. Among the tested settings, \(G=16\) performs best on PEMS04, while \(G=20\) performs best on PEMS07. This observation suggests that the group number should be selected according to the dataset rather than increased monotonically with the number of nodes.

4pt

5.5 Case Study↩︎

Figure 4 provides a qualitative visualization of the learned soft node-to-group assignment matrix on PEMS-BAY. Several groups exhibit relatively concentrated assignment patterns, and some corresponding hard assignments appear to cover spatially contiguous road segments. This observation suggests that the learned grouping mechanism may capture corridor-level regularities in the sensor network. However, the visualization should be interpreted as qualitative evidence rather than a formal validation of geographical consistency or model interpretability.

6 Conclusion↩︎

In this work, we introduced STKAN, a spatio-temporal forecasting architecture that incorporates Taylor-polynomial KAN mappings into spatial and temporal token-mixing modules. The model combines learnable soft node grouping, group-wise spatial mixing, temporal mixing, and spatial–temporal attention to model traffic dynamics. Experiments on five benchmark datasets show that STKAN achieves competitive forecasting performance, with particularly strong results on the evaluated traffic-flow datasets. The comparison with the tested MLP variant suggests that TaylorKAN token mixers can provide a useful alternative nonlinear parameterization within the proposed architecture. These findings indicate that function-approximator design is a relevant component of spatio-temporal modeling, alongside spatial and temporal architectural design. Future work may examine parameter-matched comparisons, computational efficiency, statistical variability, and broader spatio-temporal applications.

References↩︎