Efficient Large-Scale Traffic Forecasting with Transformers: A Spatial Data Management Perspective

TL;DR

PatchSTG tackles large-scale traffic forecasting by addressing the quadratic cost of dynamic spatial modeling with an irregular patching strategy guided by a leaf KDTree. It combines spatio-temporal embedding, patch-based attention (depth for local, breadth for global), and a projection decoder to deliver efficient yet faithful spatial forecasting. Key contributions include the novel leaf KDTree-based patching, a dual attention encoder that preserves interpretability, and empirical results showing state-of-the-art accuracy with substantial speed and memory improvements on four large-scale datasets. The approach offers practical impact for city-scale traffic management by enabling fast, scalable, and interpretable spatio-temporal forecasting.

Abstract

Road traffic forecasting is crucial in real-world intelligent transportation scenarios like traffic dispatching and path planning in city management and personal traveling. Spatio-temporal graph neural networks (STGNNs) stand out as the mainstream solution in this task. Nevertheless, the quadratic complexity of remarkable dynamic spatial modeling-based STGNNs has become the bottleneck over large-scale traffic data. From the spatial data management perspective, we present a novel Transformer framework called PatchSTG to efficiently and dynamically model spatial dependencies for large-scale traffic forecasting with interpretability and fidelity. Specifically, we design a novel irregular spatial patching to reduce the number of points involved in the dynamic calculation of Transformer. The irregular spatial patching first utilizes the leaf K-dimensional tree (KDTree) to recursively partition irregularly distributed traffic points into leaf nodes with a small capacity, and then merges leaf nodes belonging to the same subtree into occupancy-equaled and non-overlapped patches through padding and backtracking. Based on the patched data, depth and breadth attention are used interchangeably in the encoder to dynamically learn local and global spatial knowledge from points in a patch and points with the same index of patches. Experimental results on four real world large-scale traffic datasets show that our PatchSTG achieves train speed and memory utilization improvements up to $10\times$ and $4\times$ with the state-of-the-art performance.

Figures (8)

• Figure 1: Sketch of our PatchSTG. We first partition the equilibrium number of points into small regions. Then we merge small regions into patches after padding points into regions where the maximum number of points has not been reached. Finally, we perform the depth and breadth attention on points in the same patch and points in other patches with the same index to capture local and global spatial knowledge.
• Figure 2: The workflow of our PatchSTG.
• Figure 3: (a)-(b): left part draws the spatial partition using the original and leaf KDTree, and the right part is the corresponding tree. (c): KDT and LKDT are abbreviations of KDTree and leaf KDTree.
• Figure 4: Structure of depth and breadth attention.
• Figure 5: Leaf KDTree on the GLA dataset.