Train a central traffic prediction model using local data: A spatio-temporal network based on federated learning

Abstract

With the growing complexity of onboard sensors and the widespread deployment of road sensors, deep learning enables fine-grained traffic prediction using massive amounts of raw traffic data, which facilitates accurate analysis of traffic information in the Internet of Vehicles (IoV). However, most existing studies focus on using all the local data to jointly build a prediction model, facing severe challenges of data security and privacy concerns as well as substantial communication overhead. To address these challenges, in this paper, we propose the Spatial-Temporal Traffic Prediction Network based on federated learning (F-STTP-Net), which only updates the model parameters to the centralized server without any private data. Firstly, we design a sub-area division method, which divides the road network into sub-areas with different macroscopic fundamental diagram properties. Then, we propose a local training model for each sub-area, which uses the graph attention network (GAT) and the long short-term memory (LSTM) to capture the spatio-temporal dependence of the road network. The model uses the branch structure to predict the traffic volume of each intersection in the sub-area. Finally, the local models are aggregated based on federated learning to form a powerful central model, which bridges the constraints on global data sharing and privacy guarantee. We conduct experiments on the real-life dataset in Xuchang Lotus Lake 5G automated vehicles demonstration area to demonstrate that F-STTP-Net can achieve excellent prediction performance without the interaction of sub-area raw data. In addition, the proposed model has a strong generalization ability and can be quickly transferred to a new sub-area.