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

TL;DR

This paper introduces PatchSTG, a Transformer-based framework that efficiently models large-scale traffic data by using irregular spatial patching, significantly reducing computational complexity while maintaining high forecasting accuracy.

Contribution

The paper proposes a novel irregular spatial patching method using KDTree and a Transformer framework that improves speed and memory efficiency for large-scale traffic forecasting.

Findings

Achieves up to 10x faster training speed

Reduces memory usage by up to 4x

Maintains state-of-the-art forecasting performance

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.