ASTPKEFormer: Adaptive Spatiotemporal Prior Knowledge Embedding-Induced Transformers for Traffic Data Forecasting

1 Introduction

Traffic data prediction is a key task in Intelligent Transportation Systems (ITS), aiming to predict future time series of road networks based on historical observation. Accurate prediction can promote proactive decision-making and optimal resource allocation, providing data support for downstream tasks such as traffic management, congestion mitigation, and travel time estimation, thereby improving the overall efficiency [Luo et al., 2023b].

* Xiaoyun Xia and Xiangjie Kong are corresponding authors.

In recent years, with the rapid development of deep learning, numerous data-driven models have been proposed to capture the complex temporal dynamics and spatial correlations in traffic networks. Among them, Spatiotemporal Graph Neural Networks (STGNNs) [Yu et al., 2018; Diao et al., 2019; Fang et al., 2019] were among the earliest models developed to jointly learn spatial and temporal dependencies, giving rise to three main architectural paradigms. First, static-graphbased approaches [Zhao et al., 2019; Lv et al., 2020; Wang et al., 2020; Pan et al., 2020] utilize predefined road network topologies to model spatial dependencies among sensors or road segments. Second, dynamic-graph-based approaches [Diao et al., 2019; Zhang et al., 2020; Bai et al., 2020; Guo et al., 2021; Zhang et al., 2025] adaptively update graph structures over time to reflect the evolving spatial relationships driven by traffic dynamics. Finally, hybrid static-dynamic graph approaches [Wu et al., 2019; Shao et al., 2022b; Li et al., 2023; Ma et al., 2024] combine the advantages of both, enabling simultaneous modeling of explicit static and implicit dynamic dependencies. However, due to the inherent structural limitations of STGNNs, these models still struggle to capture long-range temporal dependencies and global contextual relationships. In contrast, Transformer-based models [Liu et al., 2023; Jiang et al., 2023a; Geng et al., 2024; Guo et al., 2025; Fang et al., 2025] have attracted increasing attention for their superior capability in modeling global dependencies, evolving along two primary directions. On one hand, models based on the vanilla self-attention mechanism effectively capture long-term temporal interactions and global correlations but suffer from high computational complexity. On the other hand, models employing efficient or linearized attention mechanisms aim to significantly reduce computational costs while maintaining comparable predictive performance. However, as the field progresses, it has become apparent that increasing the complexity of network architecture has not brought about significant performance improvements, which necessitates shifting our focus to the data itself.

Based on the above, this article focuses on a simple yet highly effective representation technique-input embeddings like [Liu et al., 2023]. Fig .1 presents the principal types of embeddings commonly adopted in contemporary traffic prediction models. STGNN primarily uses feature embedding Ef [Yu et al., 2018; Wu et al., 2019], which projects raw inputs into a latent representation space via a linear transformation. Transformer-based models typically introduce additional time period information (hour, day, week) embedding Et [Zheng et al., 2020; Shao et al., 2022a] and employ temporal position encoding Etpe [Jin et al., 2023] to assist the self-attention mechanism in understanding the temporal flow characteristics of the sequence. In addition, spatial embedding Es are extensively utilized, often learned as implicit node representations to encode spatial correlations across the transportation network. More recent architectures, including STAEformer [Liu et al., 2023], ASTRformer [Wang et al., 2024], and DTRformer [Chen et al., 2025b], all utilize adaptive spatiotemporal embedding Ea to improve the model's ability to model temporal sequence information and other traffic patterns in the input. However, most existing embedding strategies still exhibit notable limitations: (1) the spatial embedding component often underutilizes prior structural knowledge of the road network, limiting its expressiveness; and (2) the interactions among different embeddings remain insufficiently modeled, as current approaches commonly rely on simple concatenation or additive fusion (e.g., Ea and Et), thereby lacking a deeper mechanism for crossembedding interaction.

To address the aforementioned limitations, we incorporate spatial structural priors into the Transformer prediction framework in the form of a spatial prior embedding Eg. Our approach, referred to as Adaptive Spatiotemporal Prior Knowledge Embedding-Induced Transformers (ASTPKEformer), enhances the model's spatiotemporal representational capability by introducing temporal and spatial prior embeddings, while employing a cross-fusion mechanism to strengthen the deep interactions among different embeddings, particularly between the adaptive spatiotemporal embedding Ea and the prior-based embeddings (Et and Eg). Moreover, we design a multi-scale temporal-channel attention module to further improve the modeling of feature embeddings Ef by capturing hierarchical temporal dependencies. Extensive experiments and analyses conducted on six real-world traffic datasets demonstrate that the proposed method achieves SOTA performance in traffic prediction. The main contributions of this work are summarized as follows:

• We propose a novel Transformer-based architecture that integrates spatial structure priors and temporal periodic priors into a traffic prediction framework. By embedding spatial and temporal prior knowledge into the model, the proposed ASTPKEformer effectively enhances the model's ability to capture explicit and im-
• plicit spatiotemporal dependencies in complex traffic networks.
• We construct a multi-scale spatiotemporal channel attention module, enabling the model to jointly learn multi-level temporal dynamics and cross-channel feature interactions. Furthermore, we design an embedding cross-fusion mechanism to promote deep interaction between heterogeneous embeddings, thereby improving the model's representational richness and robustness.
• Extensive experiments conducted on six realworld benchmark traffic datasets demonstrate that ASTPKEformer consistently outperforms existing SOTA models, achieving the best overall predictive performance across multiple evaluation metrics.

2 RelatedWork

Spatiotemporal Traffic Prediction. STGNNs and Transformers, which focus on modeling spatiotemporal dependencies, are two highly representative techniques applied to solving traffic data prediction tasks [Fang et al., 2026]. For example, STGNN [Yu et al., 2018] utilize graph convolutional networks based on predefined prior graph structures to extract spatiotemporal features of multi-scale traffic networks. AGCRN [Bai et al., 2020] further learns adaptive graph structures to capture implicit spatial dependencies, while GWNet [Wu et al., 2019] and the DGCRN [Li et al., 2023] integrate static and dynamic graphs to jointly explore explicit and implicit spatial relationships. Despite their success, STGNN-based approaches still exhibit limited capability in modeling global dependencies. The self-attention mechanism in Transformers effectively captures global dependencies. Building on this, STAEFormer [Liu et al., 2023] employs a vanilla Transformer architecture for sequential temporal modeling. PDFormer [Jiang et al., 2023a] and STGAFormer [Geng et al., 2024] perform parallel spatiotemporal modeling, where PDFormer enhances microscopic spatial representation through traffic delay-aware learning, while STGAFormer strengthens local temporal feature extraction via a gating mechanism. However, these methods still have shortcomings in the global and local modeling of spatiotemporal dependencies.

Spatiotemporal Embedding Learning. In traffic spatiotemporal prediction tasks, while complex model architectures have been extensively studied, input embedding techniques have also received considerable attention. These techniques aim to efficiently represent input data and enrich its informative content, thereby improving prediction accuracy. For example, feature embedding is widely used in various types of prediction methods, including but not limited to STGNN-based [Kong et al., 2023; An et al., 2025], Transformer-based [Zhao et al., 2024; Jin et al., 2025], and Linear-based [Zeng et al., 2023] methods. Temporal and spatial embeddings are also frequently used in these methods, such as GMAN, STID, and LGSTFormer [Zhou et al., 2025]. Position embedding is typically used in Transformerbased methods because attention mechanisms cannot preserve sequential information between time points. Furthermore, adaptive spatiotemporal embedding was proposed in STAEformer, and the effectiveness of this simple structure has been proven; it is also widely used in methods such as ASTRformer, ImputeFormer, and STDHformer. Despite the achievements of these spatiotemporal embedding techniques, the interactions between different embeddings are often overlooked.

3 Preliminary

Traffic Spatio-Temporal Graph: Given a fixed traffic road network consisting of N nodes, the traffic spatiotemporal graph is formally defined as G = (V, E,A,X). Specifically, |V| = N denotes the set of nodes, and E represents the set of edges that describe the spatial connectivity among nodes. The corresponding adjacency matrix is denoted as A ∈ RN×N , where each element aij equals 1 when an edge exists between node i and node j, and 0 otherwise. Furthermore, at each time step t, the dynamic graph signal of the traffic network (e.g., traffic flow or speed) is represented as a time-varying feature matrix Xt ∈ RN×D , where D = 1 denotes the input dimension corresponding to the traffic volume.

Traffic Spatio-Temporal Graph Prediction: Given the historical traffic data Xt-T+1:t ∈ RT×N×1 in the previous T time steps, the goal of traffic sptio-temporal graph prediction is to infer the traffic data Xt+1:t+P ∈ RP×N×1 in the future P time steps by learning a project F(·) based on the traffic spatiotemporal graph G:

4 Methodology

An overview of the proposed ASTPKEformer is shown in Figure 2. Our model is composed of a multi-scale temporal-channel attention module (MTCA Module), an embedding layer (Embedding Layer), a prior knowledge embedding cross-fusion module (PKECF Module), vanilla Transformers are applied along the temporal and spatial dimensions respectively to construct the temporal and spatial Transformer layers , and finally, a regression layer for prediction output ((Transformer and Regression Layer)).

4.1 MTCAModule

As in most spatiotemporal prediction architectures, a fully connected layer is first employed to project the raw input into a d-dimension representation space, yielding the feature embedding Ef ∈ RT×N×d , where d represents the feature embedding dimnesion. Distinct from conventional designs, we further introduce a feature enhancement layer to better capture the intrinsic temporal characteristics of traffic data. Specifically, the enhancement module consists of two main components: a channel alignment block (CA) and a temporal alignment block (TA). The CA consists of two alternately executed fully connected network (FC) and layer normalization (LayerNorm), following the design principle proposed in [Luo et al., 2023a]. In the first stage, a squeeze operation is applied to the feature embedding Ef along the channel dimension using an FC-LayerNorm sequence, enabling the model to recalibrate the distribution in the low-dimensional subspace. Subsequently, an excitation operation is performed through a gated threshold learning mechanism, which adaptively evaluates and emphasizes the relative importance of each channel as:

where HCA denotes the compressed low-dimensional feature representation, KCA represents the importance coefficients associated with each channel, and ⊙ indicates the Hadamard element-wise product, with ReLU and Sigmoid serving as the nonlinear activation functions.

Similarly, to emphasize critical temporal points within the sequence, we design a multi-scale temporal convolution layer within the TA to evaluate the importance of different time steps from multiple receptive fields. Each temporal convolution branch is formulated as a squeeze-excitation module as:

where Hk TA denotes the low-dimensional temporal feature representation obtained with a 1 × k convolution kernel, Kk TArepresents the importance factors for different time steps, Θ1,k and ∗conv denote the convolution kernel and the 2D convolution operation, respectively. BatchNorm indicates the Batch Normalization function. Based on this design, given k different convolution kernel sizes {k1, k2, . . . , k}, the multi-scale temporal convolution branch produces enhanced temporal features under different receptive fields, denoted as { ˆHk1 TA, ˆHk2 TA, . . . , ˆHk TA} . Finally, a FC layer is applied to get feature embedding Ef ∈ RT×N×d :

4.2 Embedding Layer

We improve the model's learning capacity by supplying multiple types of embedding inputs. Specifically, we design three distinct categories of embeddings, including temporal prior embedding, spatial graph-structure prior embedding, and adaptive spatiotemporal embedding.

Temporal prior embedding Et: Et is composed of two temporal-information embeddings: a timestamp-day embedding Ed and a day-week embedding Ew. The former encodes the sampling timestamp within a day, while the latter identifies the specific day of the week on which the observations occur. To obtain these two embeddings, we first construct two learnable parameter matrices Wd ∈ R288×d and Ww ∈ R7×d as an embedding dict. Then, the timestamp inputs Xd ∈ RT×N and day-of-week Xw ∈ RT×N inputs are used as indices to retrieve the corresponding embedding vectors, yielding Ed ∈ RT×N×d and Ew ∈ RT×N×d . By concatenating them, we get the temporal prior embedding Et = [Ed||Ew] ∈ RT×N×2d .

Spatial graph-structure prior embedding Eg: Instead of using a learnable parameter matrix to dynamically model the spatial dependencies of the road network, we directly incorporate the adjacency matrix A of the network as the spatial embedding input. Specifically, following the procedure in [Guo et al., 2025], we first compute the forward transition matrix Af ∈ RN×N and the backward transition matrix Ab ∈ RN×N from the adjacency matrix A. These two transition matrices are then encoded through two separate multilayer perceptrons (MLPs) to obtain the graph-based representations Ef g ∈ RN×dg and Eb g ∈ RN×dg , where dg represents the graph embedding dimension. Furthermore, an expand operation is applied to align their dimensions with that of Ef , yielding the bidirectional spatial embeddings Ef g ∈ RT×N×dg and Eb g ∈ RT×N×dg . By concatenating them, we get the spatial graph-structure prior embedding Eg = [Ef g ||Eb g ] ∈ RT×N×2dg .

Adaptive spatiotemporal embedding Ea: For a given traffic node, its traffic feature at a specific timestamp is often similar to the state observed at the same timestamp on the previous day or the previous week, and it is also influenced by the traffic conditions at adjacent time points. Moreover, the spatial dependencies among nodes in a road network are typically dynamic. Inspired by [Liu et al., 2023], we introduce a learnable adaptive spatiotemporal embedding Ea ∈ RT×N×da , where da represents the adaptive spatiotemporal embedding dimension, to capture these complex spatiotemporal relationships in a unified manner.

4.3 PKECFModule

The embeddings Et, Eg and Ea correspond to temporal, spatial, and spatiotemporal information, respectively. We design an embedding interaction module to strengthen the depth of cross-type information exchange. This module comprises two interaction processes: temporal-adaptive spatiotemporal interaction and spatial-adaptive spatiotemporal interaction.

Temporal-Adaptive spatiotemporal interaction (TASI): We first concatenate Et and Ea to get Eta = [Et||Ea] ∈ RT×N×(2d+da) . Then we introduce a multilayer perceptron with residual connections (MLPRes) as the basic computational unit, where each MLPRes comprises two fully connected layers interleaved with a ReLU activation function.

By sequentially stacking 2 MLPRes blocks, the interaction between Et and Ea is achieved as:

where E2 ta ∈ RT×N×(2d+da) . It is worth noting that we reseparate E2 ta into ˆEt and ˆEa along the channel dimension, while maintaining their original dimension size.

Spatial-Adaptive spatiotemporal interaction (SASI): Same as Temporal-Adaptive spatiotemporal interaction, we concatenate Eg and Ea to get Ega = [Eg||Ea] ∈ RT×N×(2dg+da) , and use 2 MLPRes blocks to achieve the interaction between Eg and Ea as:

where E2 ga ∈ RT×N×da . Furthermore, we reduce the complexity of the E2 ga parameters by using an FCre : R2dg+da → Rda layer. Finally, By concatenating the embeddings above, we can get the hidden spatio-temporal representation H ∈ RT×N×dh as follows:

where the hidden dimension dh equals 3d + da.

4.4 Transformer and Regression Layer

Transformer Layer: To learn the inherent spatiotemporal dependencies of traffic data, we use a pure self-attention mechanism along the temporal and spatial axes of the embedded input H to calculate the importance between different time points and between different nodes [Liu et al., 2023]. By definition, we can obtain the query, key, and value matrix for the corresponding temporal attention based on H as:

where W tQ,W tK ,W tV ∈ Rdh×dh are all learable parameters. Next, we compute the temporal self-attention score in parallel by performing matrix multiplication between Qt and Kt as:

where Attnt captures the implicit temporal dependencies. Subsequently, we output the temporal feature representation Ht as:

Similarly, the spatial self-attention layer performs as:

where Hs is the spatial feature representation. selfAttention repeats the calculation process of formulas 15, 16 and 17 in the spatial axis.

Regression Layer: Finally, we used a single FC layer to output the spatio-temporal representation Hs as the future

where Xt+1:t+P is the prediction result, P is the length of prediction.

5 Experiment

5.1 Experiment Settings

Datasets: We select six widely used traffic datasets covering two different data modalities to comprehensively evaluate the performance of ASTPKEformer. The detailed statistics of each dataset are provided in Table 1.

Metrics: To evaluate the traffic prediction performance of ASTPKEformer, we adopt three commonly used evaluation metrics to quantitatively assess the experimental results: Mean Absolute Error (MAE), Root Mean Square Error (RMSE) and Mean Absolute Percentage Error (MAPE).

Baselines: We compare ASTPKEformer with a series of representative spatio-temporal models to demonstrate the model's effectiveness. These models can be categorized into three types according to their calculation patterns: 1) Linearbased methods: STID [Shao et al., 2022a], BLSTF [Chen et al., 2025a]; 2) GCN-based methods: AGCRN [Bai et al., 2020], GWNet [Wu et al., 2019], MegaCRN [Jiang et al., 2023b], HimNet [Dong et al., 2024], STDN [Cao et al., 2025]; 3) Attention-based methods: GMAN [Zheng et al., 2020], STWave [Fang et al., 2023], STAEformer [Liu et al., 2023], STHDformer [Guo et al., 2025].

Experimental Setup: Following the settings fromprevious studies [Liu et al., 2023], we divide the dataset into training, validation, and test sets in a 6:2:2 ratio for traffic flow datasets and 7:1:2 ratio for traffic speed datasets based on the time order. Our experiments were conducted on a GeForce GTX 4090 GPU. All models were implemented using PyTorch. We set the model's input and output length to 1 hour, i.e., T=P=12. The optimizer was set to Adam with a decay learning rate of 0.001 and a batch size of 16. We also employed an early stopping mechanism with a patience value of 10. The kernel size of MTCA is set to [1,3,5,7,9], and the graph structure embedding size is 64. Furthermore, for the six datasets in Table 1, the hidden embedding d and the adaptive spatio-temporal embedding da is set to {24, 24, 24, 32, 24, 24} and {80, 40, 80, 80, 80, 80}, respectively. The graph embedding dg maintains the same size across all datasets, 64.

5.2 Performance Analysis

We present the prediction results of all methods in Table 2 and Table 3, where - indicates that no result was obtained due to excessive computational resources required. The following observations can be drawn from the prediction results: (1) The proposed ASTPKEformer achieves the best results on all six datasets, demonstrating that its enhanced embedding-interaction mechanism effectively captures spatiotemporal dependencies and yields strong generalization; (2) Despite their simplicity, MLP-based methods do not perform the worst on either flow or speed prediction tasks and often surpass models such as AGCRN, STDN, and STWave, indicating that even basic embedding strategies can provide competitive accuracy; (3) Among GCNbased approaches, MegaCRN performs best on METR-LA, while HimNet achieves superior overall results on PEMSBAY and the flow datasets; (4) Attention-based models generally outperform GCN-based ones. STHDformer, which also integrates graph-structure embeddings, ranks second only to ASTPKEformer on flow datasets, and STAEformer delivers strong performance on speed datasets, underscoring the adaptability of attention mechanisms to spatiotemporal traffic forecasting.

5.3 Performance Evaluation of MTCA

As one of the most important components proposed in this study, we further evaluate the effect of MTCA module. We selected three comparison models that can be used in a plugand-play manner for MTCA, including GWNet, GMAN and STWA. Table 4 shows the performance comparison results. Based on the experimental results, it can be seen that MTCA can improve the prediction accuracy of the original model to a certain extent. This shows that MTCA is applicable to different types of models and can be used as a plug-and-play module.

5.4 Ablation Experiment

To investigate the relative importance of different modules in the prediction process, we construct six model variants: (1) w/o MTCA: the multi-scale temporal-channel attention (MTCA) module is removed; (2) w/o PKECF: the embedding interaction module is removed; (3) w/o TASI: the temporal prior embedding Et and adaptive spatiotemporal embedding Ea do not interact; (4) w/o Ea: the spatial-structure prior embedding is removed, which also eliminates the SASI process; (5) w/o SA: the spatial Transformer module is removed; (6)

w/o TA: the temporal Transformer module is removed. Corresponding ablation experiments are conducted on the traffic flow dataset PEMS08 and the traffic speed dataset METRLA. The results of these experiments are presented in Table 5. The results clearly show that each module contributes positively to the prediction performance. Removing the introduced Eg degrades prediction performance on both datasets. Furthermore, removing the interactive modules has a stronger negative impact, further demonstrating the importance of interactions between embeddings.

5.5 Hyperparameter Experiment

We conduct hyperparameter experiments on the PEMS04 and PEMS08 datasets to evaluate the impact of the feature embedding dimension d, the graph structure embedding dimension dg and the adaptive spatiotemporal embedding dimension da on the model's predictive performance. The results are shown in Figure 3. We tune d ∈ {12, 24, 36, 48}, dg ∈ {24, 32, 64, 128}, and da ∈ {20, 40, 80, 100}. The results show clear dataset-dependent behaviors: the optimal d and da vary across datasets due to differences in traffic dynamics. For PEMS04, larger d degrades performance, likely from over-parameterization, whereas for PEMS08, moderately increasing d strengthens the model's ability to capture complex temporal patterns. A similar trend holds for da, where datasets with higher variability or stronger spatiotemporal correlations benefit from larger adaptive embedding dimensions.

5.6 Visualization Study

We visualize the adaptive spatiotemporal embedding Ea and compare it with STAEformer using 2D t-SNE [Maaten and Hinton, 2008] for spatial analysis and time-step correlation for temporal analysis. As shown in Fig.4. ASTPKEformer exhibits clearer spatial clustering, indicating enhanced spatial representation from graph-structured priors. Temporally, unlike STAEformer's overly smooth correlations, ASTPKEformer shows a sharper decay that better captures short-term dependencies and long-term attenuation. These results demonstrate that interacting Et with Ea effectively improves temporal modeling and validate the benefits of graph priors and strengthened embedding interactions.

6 Conclusion

In this work, we propose ASTPKEformer, a Transformerbased traffic forecasting model guided by adaptive spatiotemporal prior knowledge. Building upon the standard spatiotemporal self-attention mechanism, the model first employs a multi-scale temporal-channel alignment module to generate high-quality feature embeddings. It then enriches the input representation by simultaneously integrating temporal prior embeddings and graph-structured spatial prior embeddings. Furthermore, a cross-fusion mechanism is designed to facilitate deep interactions between prior embeddings and adaptive spatiotemporal embeddings, thereby strengthening the model's capacity to capture complex spatiotemporal dependencies. Extensive experiments conducted on six real-world benchmark datasets demonstrate the effectiveness and superiority of the proposed framework. In the future, we will test the model's performance in other scenarios, such as those with missing data or few samples.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grant 62476247 and 61703183, and in part by the National Key Research and Development Program of China under Grant 2023YFC3305900, 2023YFC3305903, and in part by Zhejiang Provincial Natural Science Foundation of China under Grant ZCLZ25F0201, D24F020013, LGG19F030010, and in part by the Qin Shen Scholar Program of Jiaxing University..