Knowledge Distillation-Based Spatio-Temporal MLP Model for Real-Time Traffic Flow Prediction

Abstract

Real-Time Traffic Flow Prediction (RT-TFP) is one of the critical technologies for implementing the Intelligent Transportation System (ITS), enabling rapid and accurate prediction of real-time traffic flow at intersections. RT-TFP typically needs to be deployed on-site edge devices for real-time traffic flow calculation that requires low inference latency and minimal computational resources. However, the existing Traffic Flow Prediction (TFP) models are generally based on spatiotemporal graph neural networks (STGNNs), which are complex and require high computational resources and relatively high inference times that can hardly be deployed on edge devices. To this end, this work proposes a simple RT-TFP model, SpatioTemporal-MultiLayer Perceptron (ST-MLP), which requires low computational resources and inference times. The base idea of this work is to establish a spatio-temporal MLP model to replace the STGNN model for conducting the TFP, which is much faster and simpler. Specifically, first, a TempEncoder is proposed to encode the temporal information into the MLP features. Then, a Spatiotemporal Mixer is proposed to mix spatial information into the temporal-enriched MLP features. After, MLP features are distilled from a complex STGNN model to obtain a simple MLP that inherits complete Spatial-Temporal information of the traffic graph. The experimental results on four real-world datasets show the proposed model achieves competitive prediction accuracy with STGNN models in much fewer computational resources and lower prediction time costs. It is worth noting that, the proposed method is faster than the compared STGNNs by an average of 21.62 times (~10.81s $\rightsquigarrow ~\sim 0.50$ s). Interestingly, the proposed ST-MLP even has a −3.23% error rate decreasing on average compared to the corresponding STGNN model. Moreover, the error rate of the proposed ST-MLP decreases over pure MLPs by −3.92% $\sim -42.62$ %. The source code is available at: https://github.com/zhangjunfeng1234/ST-MLP