STAMImputer: Spatio-Temporal Attention MoE for Traffic Data Imputation

1 Introduction

Intelligent Transportation Systems (ITS), which rely on comprehensive and high-quality data to perform their tasks effectively [Liu et al., 2023; Yin et al., 2021], play a significant role in maintaining urban traffic order. However, in realworld traffic data collection, limitations in resources such as sensors (e.g., vehicles, drones) and disruptions in data acquisition (e.g., network interruptions, extreme weather conditions) can lead to data missing in both spatio-temporal dimensions [Chan et al., 2023; Zhang et al., 2024b]. Therefore, traffic data imputation, which aims to reconstruct the missing traffic data by leveraging potential data correlations, has attracted rising research attention recently.

* Corresponding authors.

† Complete Appendix is available for download via the code link.

Existing research on traffic data imputation can be broadly categorized into time-series imputation and spatio-temporal imputation. Time-series imputation primarily involves algorithms based on statistical analysis [Yu et al., 2016; Chen et al., 2020; Chen and Sun, 2021] or deep learning models that concentrate efforts exclusively on the temporal dimension of the traffic data [Cao et al., 2018; Du et al., 2023; Zhang et al., 2024a]. However, these methods are not effective in addressing the issue of block data missing, where data is absent in consecutive sequences. Spatio-temporal imputation [Xu et al., 2022; Wang et al., 2022; Li et al., 2022; Nie et al., 2024] accounts for the extra spatial dimension with static graphs or spatial attentions, allowing for more effective handling of block missing data.

Although considerable effort has been made on traffic data imputation, there are still two major challenges that hinder existing works from achieving a more accurate imputation result. First, most current time-to-space sequential learning frameworks [Marisca et al., 2022; Wei et al., 2024; Zou et al., 2024], which process spatio-temporal data by first using a temporal module followed by a spatial module, may extract and propagate invalid features when block data are missing. The block data missing scenario can occur in both spatial and temporal dimensions. This challenge is even more pronounced in scenarios with spatio-temporal blockwise missing data, for example, due to prolonged power outages in physical sensors. Second, static graph-based methods struggle to capture global spatial dependencies beyond local communities, which are essential for making accurate predictions in practical applications. For example, the traffic volume or average traffic speed of a crossroad during non-peak hours is closely related to its adjacent intersections. However, during the evening peak, these observations are also affected by nearby traffic hub nodes, which are key sensor locations that have the most significant impact on the entire traffic network. Many static graph-based methods [Ye et al., 2021; Xu et al., 2024] are limited in their ability to capture the global spatial correlations. Without static graphs, global spatial attention methods [Zhang et al., 2022; Nie et al., 2024] can propagate spatial features across traffic networks but struggle to capture spatial correlations in highly sparse data due to the lack of guidance from prior knowledge. Hence, we argue that spatio-temporal imputation requires dynamically adjustment according to the real-time missing data. Furthermore, combining the strengths of static graph-based and global attention-based methods will effectively handle complex imputation scenarios.

To address the above challenges, we propose a SpatioTemporal self-Attention Mixture of experts Imputer network named STAMImputer. We perform traffic data imputation by deploying specialized experts networks, including attention experts and observation experts. Attention experts are self-attention transformers that capture temporal or spatial features, while observation experts act as arbitrators, overseeing and controlling the broad spatio-temporal attention. The MoE framework dynamically adjusts real-time spatiotemporal attention weights by utilizing the routing control of the observation experts. If the data is sparse in terms of the spatial demension but rich in terms of temporal dimension, the observation experts guide the model to prioritize the temporal experts' learned representations. Moreover, inspired by [Fang et al., 2023] and [Nie et al., 2024], we design a novel Low-rank guided Sampling Graph ATtention mechanism, namely LrSGAT. The spatial low-rank guided matrix decoupling identifies key features of traffic hubs while filtering out the redundant global relationships. Attention projection is then applied to propagate the relevant features of these hubs. Finally, the sampled attention vectors contributes to generating dynamic graphs which fully capture the real-time spatial correlations and help improve downstream traffic-related tasks. Our contributions are summarized as follows.

• We propose a novel traffic data imputation model based on the MoE framework. Our model dynamically adjusts attention weights of spatio-temporal expert networks to balance the contributions of both spatio-temporal dimensions. To the best of our knowledge, this is the first application of the MoE framework to traffic data imputation tasks.
• We propose a low-rank guided sampling graph attention mechanism that enables the imputation model to capture global spatial dependencies beyond local communities.
• We build real-time dynamic graphs by utilizing the attention vectors sampled by the spatial expert network, which accurately represent the current traffic conditions.
• Extensive comparative experiments on four real-world benchmark datasets demonstrate that our STAMImputer model outperforms others in traffic data imputation.

2 RelatedWorks

Traditionally, statistical based approaches [Nelwamondo et al., 2007; Van Buuren and Groothuis-Oudshoorn, 2011; Yi et al., 2016] are used to study the mathematical characteristics and functional properties of spatio-temporal sparse data, such as low-rank approximation [Cichocki and Phan, 2009]. In terms of spatio-temporal dual-dimensions, [Li et al., 2013] proposed an algorithm based on spatio-temporal probabilistic principal component analysis (PPCA). Both [Deng et al., 2021] and [Xu et al., 2023] proposed low-rank guided tensor completion algorithms. The former used graph spectral regularization, while the latter introduced a Hessian regularization spatio-temporal low-rank method.

More recently, deep learning approaches have been proposed for traffic data imputation. Early works include Autoencoder, RNN and CNN networks [Cao et al., 2018; Benkraouda et al., 2020; Zhao et al., 2020], and they rely on a keen sense of temporal variation to impute the missing data. As development progressed, graph-based methods are introduced [Chen and Chen, 2022; Liang et al., 2022; Kong et al., 2023], contributing spatial dimensionality correlations to the imputation work. Notably, [Cini et al., 2021] proposed bidirectional autoregressive architecture which conducted spatio-temporal imputation through message passing graph neural networks. Nevertheless, GANs [Yoon et al., 2018; Chen et al., 2019] are also used in traffic data imputation. For example, STGAN [Yuan et al., 2022] implemented a spatio-temporal GAN model to perceive local and global spatio-temporal distribution.

In more recent work, the self-attention mechanism provides a fresh and excellent choice for imputation work in spatio-temporal dimensions [Vaswani, 2017]. [Ma et al., 2019] proposed a cross-dimensional attention completion method, which, for the first time, applied the self-attention mechanism to multivariate geotagged time series data to achieve a cross-dimensional joint capture of self-attention. The SPIN model in [Marisca et al., 2022] implemented an spatio-temporal attention based framework, whitch can handle sparse data without propagating prediction errors or requiring a bidirectional model to encode forward and backward time dependencies. [Nie et al., 2024] adopted attention on both spatio-temporal dimensions. Combining low-rankinduced temporal attention and spatial embedding attention, it demonstrates state-of-the-art performance in traffic spatiotemporal data imputation. However, we argue that none of the above-cited methods can dynamically balance the spatiotemporal dual-dimensional learning features.

3 ProblemDefinition

The traffic completion problem aims to infer unobserved or unrecorded data using known traffic data recorded by sensors. Specifically, Xt:t+T ∈ RT×N×C , where C = 1, represents the observed traffic flow or speed value of N nodes in traffic network from a fixed time slice t to t+T , andMt:t+T sharing the same size with Xt:t+T represents the missing position of the observed data as follows:

Given Xt:t+T , Mt:t+T and the topology graph GK which represents a K-nearest neighbours graph structure of traffic networks, the purpose of our paper is to learn a function Imputer to infer the missing values in Xt:t+T , that is, to obtain the imputed sequence data ˆ Yt:t+T ∈ RT×N×C . The task can be formulated as:

where Θ denotes the learnable parameters in our model.

4 Methodology

In this section, we introduce our model STAMImputer. The overview of STAMImputer is shown in Figure.1 The left side of the figure shows the framework of MoE, which depicts the collaboration mode between the observation and the attention experts. The right side of the figure shows the workflow of the spatial experts, which we introduce in detail in Section 4.2.

4.1 Observation Expert andMoE Framework

As an emerging architecture, MoE has been widely used and discussed in deep learning fields. Unlike the well-known MoE in [Fedus et al., 2022; Zhang et al., 2021], we extend the expert networks to the outer layer of the framework. The advantage of this design is its ability to decouple data feature in a controlled way, enabling the temporal and spatial experts to effectively capture feature propagation relationships in both dimensions. It is called controlled decoupling because there exist a macro-control expert (i.e., Observation Expert in MoE) who can evaluate the trust of the learning results of functional experts by focusing on raw data and sparse features. Through the tasks and cooperation of experts, the system can decouple dimensional attention while ensuring cohesion in the final result, which is the original intention of the design of our MoE framework. In summary, the use of MoE enhances the model's stability and generalization. Regardless of the sparsity in time or space dimensions, the expert network can effectively guide the model based on real-time data characteristics before producing the output.

Observation Expert The evaluation ability of observation expert is directly related to the quality of the learning results and the robustness of the model. The observation expert can learn downstream-oriented evaluation feedback and make trust decisions on the spatio-temporal attention output in the completion task based on the real data (including missing) and sparsity features. Compared with the structural design of the attention expert, the observation expert can not only achieve the effect of a residual network but also improve the dynamics of residual connections, thereby improving the stability when handling sparse patterns of real complex data. In addition, it is worth noting that different dimensions of attention have different strengths and weaknesses under different missing patterns. For example, in the case of random data missing, temporal attention performs better than spatial attention in completing the task, while the opposite is true in the case of block missing patterns [Marisca et al., 2022]. Therefore, observing the expert's dynamic evaluation and sparse sensitivity of inputs improves the model's flexible perception of complex spatio-temporal networks.

MoE Framework As shown in Figure 1.1, MoE framework is structured into three main stages: 1) attention stage: to extract spatio-temporal correlations with MSAT and LrSGAT networks; 2) observation stage: to score the credibility of attention experts and weight the outputs of the attention stage; 3) readout stage: to produce the final predictions through Multilayer Perceptron (MLP).

Firstly, a vectors group X MoE = ⟨X oe ∥X in ⟩ , in which X oe ∈ RT×N×Doe is as inputs for observation expert while X in ∈ RT×N×Din is as inputs for attention experts (refer to Section.4.3 for details). Then, hierarchical iterative correlation learning and feature reconstruction by spatio-temporal experts in the attention processing stage:

X i represents the inputs or outputs of the attention stage in layer i and X outi will be the input of layer i + 1. Obviously, X tin0 = X sin0 = X in . The observation expert performs attention confidence assessment on node n at time t:

where W oe ∈ RDoe×N ae is a learnable parameter matrix and N ae is the number of attention experts. Finally, the learning feedback from the experts is weighted and concatenated according to the confidence scores E ST ∈ RT×N×N ae from the observation expert and input into the readout stage to obtain the imputation outputs.

4.2 Attention Experts of Dual-dimensions

Self-attention extraction and influence propagation are the core of STAMImputer. Whether the potential interactive relationship can be extracted is the key to improving the imputation effects. In this section, we describe in detail the design of attention experts in STAMImputer and how they effectively extract latent patterns.

Temporal Attention Expert Networks

In the temporal dimension, we introduce the encoder layer of the classic self-attention Transformer [Vaswani, 2017] as the temporal expert network unit, that is, the Multi-head SelfATtention (MSAT) network shown in Figure 1.1. In long or short time slices, the interaction patterns of time points are relatively stable (relative to the spatial dimension), which is the main reason why simple time series imputation has a weak ability to perceive high-frequency and struggles to detect abnormal events. However, temporal self-attention still has a strong advantage in restoring stable sequence trends (i.e., low-frequency recovery). Using the classic selfattention encoder as the initial temporal processing layer of the MoE framework can quickly pre-reconstruct missing input features, especially for data with low missing rates or simple missing patterns (such as random point missing). Moreover, its output vectors can provide richer information for the

(1) MoE Framework

(2) LrSGAT Workflow

spatial attention expert network, enabling the dynamic discovery of potential spatial associations.

Spatial Attention Expert Networks

The spatial patterns that can be learned will more closely align with the unpredictable nonlinear correlations observed in real-world networks. In the spatial attention expert network, as shown in Figure 1.2, we propose a Low-rank guided Sampling Graph ATtention mechanism, namely LrSGAT. This innovative design employs hybrid sampling strategies to effectively balance local and global correlations across different time periods. Additionally, the sampled attention weights can also serve as projection vectors for the input to the re-attention layer and be used to construct semi-adaptive dynamic graphs offering dynamic spatial dependencies for downstream traffic intelligence tasks.

Sampling Projector We state that the design is to optimize the balance between local and global correlations of spatial nodes while maintaining manageable complexity. We propose a sampling-based projector design, where a reduceddimensional attention vector is sampled through a general Graph Attention Network (GAT) guided by a static topology. The output vector is derived by using the sampled attention as the projection message. Specifically, an original input X Sin t ∈ RN×Din enters the LrSGATnetwork and extracts the attention weights on the local graph through the static topological adjacency matrix as formula:

where QGt ∈ RN×1×Din and KG t ∈ RN×E×Din are expanded query and key vectors calculated with input X Sin t and the static topology with E denoting the number of neighbors for each node. In addition, the canonical vectors Qt, Kt, and Vt are also obtained in this process. E G t ∈ RN×E denotes the feedback of local attention, and its role is to reversely inform the global influence of the node at the current time t. We set up a trainable scorer vector Wsc ∈ RE×1 to generate a significance score matrix E W t ∈ RN×1 from local attention and perform node sampling based on it:

N S t denotes the set of nodes sampled at time t, where N T t represents the S = ⌈logN⌉ nodes with the highest significance scores, and N U t represents the S nodes of the probability distribution sampling from the remaining parts. The combination of these two sampling strategies is designed to focus attention on different score domains, ensuring that potential traffic hubs are not overlooked due to data missing.

The abstract concept of N S t refers to the selection of nodes with the highest global influence as well as probabilistically selected random nodes with incremental effects. Finally, the projection vector consists of the mapping of the attention query or key to the samples, and the projection message is calculated using the sampled query vector via the self-attention mechanism:

Through sample mapping, the total size of the obtained projection vector is significantly reduced compared to the original input vector since S ≪N. Furthermore, due to the evaluation of the significance scores and hybrid sampling, representative information is preserved in the projection vector, allowing P S t and Mt to pass to the next layer of LrSGAT to preserve the potential local and global correlations of the original input vector.

Low-rank Guided Re-attention Low-rank factorization is to approximate a vector X ∈ Rm×n by a vector product representation as X = UV ⊤, where U ∈ Rm×k and V ∈ Rn×k are rank matrices k ≪ min(m, n). Factorization seeks to simplify information by breaking down complex network structures into abstract, compressed features that effectively capture the overall structure. At this layer, as shown in Figure 1.2(b), we apply Re-ATtention (ReAT) learning by utilizing the projection vector and the projection message received from the previous layer of the network:

P S t encapsulates the current spatial network information into a concise and compact representation of its state and generates the key vector to guide the global nodes to pay attention to the critical node information sampled. M S t corresponds to the reference for recovering compressed information P S t and generates a value vector to recover the attention guided by the compressed information in the global reconstructed hidden state. The key to approximate low-rank factorization of the original vector is the projector, which realizes information cohesion through a sampling mechanism and generates corresponding projection messages to guide feature reconstruction. The design of repeated attention compresses and restores the high-order, low-rank spatial matrix and completes the missing data features in the process.

Semi-adaptive Dynamic Graph The pattern of spatial sampled attention and the low-rank factorization approach bring to mind a classic method for constructing adaptive adjacency matrices ˜ Aadp ∈ RN×N [Wu et al., 2019], which represent fully connected adjacency relationships learned entirely by the neural network. The matrices ˜ Aadp are formulated as:

where E1 and E2 ∈ RN×C are learnable parameters. With C ≪ N, the constructed adaptive adjacency relationship is achieved by abstracting a cohesive representation and an extroversion representation, compressing and reconstructing the global network for correlation mining and coupling it with the projector model in our two upper layers of the network. However, it is not difficult to see that each of Ei is completely abstract and static, and its compact approximation does not have any actual nodes to map. And they are susceptible to low-quality and unreliable structure at data missing conditions [Zou et al., 2023].

In contrast, the sampling matrix constructed by the projector is dynamic and node-supported. If we similarly set an E to calculate ˜ Aadp = AS E⊤, where AS ∈ RN×S is the sampled attention vector and E⊤ ∈ RS×N is the adaptive matrix, the reconstruction task that E needs to perform is faced with compressed information that changes over time. Therefore, simply reconstructing the sampled projections is not of practical significance. Taking the above factors into account, the Dynamic Graph Structure Learning (DGSL) function is defined as follows:

where Qt ∈ RN×Din is the query vector from the Sampling Projector layer, KS t ∈ RS×Din is the sampled key vector as in Formula (8), Mt ∈ RS×Din is the projection message, Eref ∈ RN×Din is a learnable parameter matrix and toph(·) is the function setting terms less than the median to zero.

4.3 Spatio-temporal Representation Learning

In this section, we summarize the spatio-temporal representation learning stage in STAMImputer. Incorporating additional feature dimensions can accelerate the convergence of expert networks. In general, the spatio-temporal and frequency features of the original input are embedded before entering the MoE networks.

Discrete Wavelet Transform Discrete wavelet transform (DWT) and inverse wavelet transform (IWT) are powerful mathematical tools that decompose signals into different frequency components, effectively capturing sequence characteristics. Specifically, given the original observed sequence vector Xt:t+T ∈ RN×T , and selecting the wavelet basis w and decomposition level j, DWT decomposes the sequence X in into low-frequency and high-frequency components:

where C0 denotes low-frequency components, and the rest denote high-frequency components. IWT is formulated as:

where X l ∈ RN×T and X h ∈ RN×T share the same shape with original vector X in .

Spatio-temporal Embedding Firstly, the temporal feature, as one of the fundamental characteristics, encompasses both the data collection time and the associated weekday information, denoted as P u ∈ RN×T×2 . Secondly, the spatiotemporal position encoding employs a learnable parameter matrix P st ∈ RT×N×Dpe , where the depth of its dimensions offers richer feature references for the data. Thirdly, statistical learning is used to calculate the sparsity rate in both time and space dimensions as the sparse features denoted P sp ∈ RN×T×2 , which enable the observation expert to evaluate the attention confidence scores.

Finally, We adopt a dimension expansion MLP layer to initialize spatio-temporal and frequency features:

where X MoE is the initialized vector that will act as embedded inflow of our STAMImputer models.

5 Experiment

Experimental Setup The detailed spatio-temporal feature information of the selected benchmark datasets is shown in Table 2. For missing patterns, we consider two scenarios discussed in previous work [Cini et al., 2021], point missing and block missing. Since we aim to perform missing imputations with spatio-temporal attention on sparse traffic data, we picked ImputerFormer as the state-of-the-art baseline method. Imputerformer introduces a low-rank projection in the time dimension for the first time, then adopts a learnable embedding representation in the space dimension. We then consider other representative deep learning methods: 1) SPIN [Marisca et al., 2022]: an attention-based model with sparse spatio-temporal graphs. 2) SAITS [Du et al., 2023]: a diagonally masked self-attention Transformer with a weighted combination of representations. 3) Transformer [Vaswani, 2017]: canonical Transformer with selfattention mechanism. 4) BRITS [Cao et al., 2018]: An imputation method that directly learns missing values in a bi-recurrent dynamical system without any specific assumptions. 5) rGAIN [Yoon et al., 2018]: A temporal imputation model based on a generative adversarial net architecture. Finally, we consider classic statistical and optimization models for comparison: 1) Mean: imputation of observation average values. 2) KNN: K neighbor mean value imputation. 3)

LATC [Chen et al., 2021]: completion of the low-rank autoregressive tensor. 4) VAR: a vector autoregressive one-stepahead predictor. All experiments are conducted on a server with an Intel(R) Xeon(R) Platinum 8336C CPU operating at 2.30GHz and an NVIDIA A800 GPU with 80GB of memory for the above models and datasets.

5.1 Performance Results Comparison Imputation Task

The imputation results on the four benchmarks are shown in Table 1. For the point-missing pattern, we choose two random missing rates of 25% and 60%, and for the blockmissing pattern, we choose two failure probabilities of 0.2% and 1% for the S ∼ U(12, 48) steps, both with a random missing rate of 5%. STAMImputer generally achieves the best performance, except in a few cases where it ranks second in all experiments. Especially in urban travel benchmarks our method demonstrates more significant advantages. ImputeFormer and SAITS show a strong perception of temporal attention. Although ImputeFormer has a corresponding attention module in the spatial dimension, both methods are weak for spatially more complex urban traffic or datasets with longer temporal steps. SPIN focuses on spatial sparse graph attention and has also shown competitiveness in complex urban traffic speed data imputation. In the case of data from the urban traffic network with many nodes, such as DiDi-SZ, the imputation error of SPIN is even better than that of ImputeFormer. Still, its spatio-temporal complexity is higher when faced with complex road networks. Transformer with a classic self-attention model is ideal for completing point-missing patterns but is relatively ineffective for blocking missing patterns. Compared with the above deep learning models, the statistical and optimization methods are less competitive in imputation effect due to their limited capabilities.

* The LATC method is invalid for training on SZ-Taxi, so its result is empty.

Downstream Task with Dynamic Graph

To verify the effectiveness of the DGSL layer results, we conduct additional experiments with a classic Graph-Wavenet network [Wu et al., 2019] for downstream prediction tasks. Table 3 shows the effects of the downstream prediction task in multiple experimental scenarios on the NYC-Taxi benchmark. To verify the information gain of dynamic graphs compared to static topologies, we disabled the fully adaptive adjacency matrix in the Graph-Wavenet network. We perform traffic prediction experiments under different random point missing rates and block missing scenarios by comparing the unmasked model learning results. Graph-Wavenet can cope to a certain extent with its excellent spatio-temporal perception in lower random missing rates. Still, when the random missing rate increases or blocks are missing, the prediction performance is greatly disturbed. We test the improvement effect of STAMImputer on the model in two steps: 1) adding a pre-trained imputation model; 2) adding a pre-trained imputation model with its learnable dynamic graph module. As shown in Table 3, our imputation model enhances the downstream prediction in failed sparse scenarios (i.e. high sparsity or missing block).

5.2 Robustness Analysis

To evaluate the robustness of the model on highly sparse data, we add imputation experiments with varying degrees of sparsity. The results are shown in Figure 2. Since the Transformer is ineffective in blocking missing patterns, its performance curve is not presented. In general, deep learning completion models suffer performance degradation as sparsity increases. However, it can be seen that the combination of attention matrices and low-rank constraints in the STAMImputer and ImputeFormer models makes them more robust in the face of high sparsity. In comparison, STAMImputer exhibits lower imputation errors.

5.3 Ablation Study

We also perform ablation studies to verify the significance of the STAMImputer framework designs. The imputation results of the ablation experiments are shown in Figure 3. We first remove the external framework of MoE, in which case we do not consider the macro-control of spatio-temporal attention by observing experts. When we implement time-tospace sequential framework instead of MoE's regulation of spatio-temporal attention in the ablation study, the model's ability to perceive spatio-temporal associations decreases. MoE plays a consistent role in the imputation task in both missing patterns. Secondly, we compare attention expert networks with four variants: MSAT with MLP as a spatial expert, LrSGAT with MLP as a temporal expert, MSAT with K neighbours GAT as a spatial expert and MSAT with full GAT as a spatial expert. As shown in Figure 3, attention expert networks play the most critical role in STAMImputer, especially in the block missing pattern, where the performance degenerates substantially after replacing LrSGAT with MLP. We finally remove the wavelet variable from the initial input to verify that the additional dimension added to the embedding helps to extract valuable information for STAMImputer.

6 Conclusion

This paper presents STAMImputer, a traffic data imputation model that leverages MoE framework and dynamically balances spatio-temporal attention. We introduce a novel spatial attention sampling mechanism that combines graph attention with low-rank decomposition to capture spatial correlations. Experiments across multiple benchmarks demonstrate that STAMImputer outperforms existing models, and also improves performance in downstream tasks. Future research will focus on developing more efficient methods for constructing dynamic graphs.

Acknowledgments

This work has been supported in part by NSFC through grants 62441612 and 62322202, Local Science and Technology Development Fund of Hebei Province Guided by the Central Government of China through grant 246Z0102G, the "Pionee" and "Leading Goose" R&D Program of Zhejiang through grant 2025C02044, National Key Laboratory under grant 241-HF-D07-01, Hebei Natural Science Foundation through grant F2024210008, and CCF-DiDi GAIA collaborative Research Funds for Young Scholars through grant 202422.