Traffic Flow Prediction Based on Federated Learning and Spatio-Temporal Graph Neural Networks

Abstract

In response to the insufficient consideration of spatio-temporal dependencies and traffic pattern similarity in traffic flow prediction methods based on federated learning, as well as the neglect of model heterogeneity and objective heterogeneity, a traffic flow prediction model based on federated learning and spatio-temporal graph neural networks is proposed. The model is divided into two stages. In the road network division stage, the traffic road network is divided into subnetworks by the dynamic time warping algorithm and the K-means algorithm, to ensure the same subnetwork has the similar traffic flow pattern. The federated learning stage is divided into two sub-stages. In the local training phase, the spatio-temporal graph neural network with an attention mechanism is utilized to create personalized models and meme models to capture the spatio-temporal dependencies of each subnetwork. At the same time, deep mutual learning is utilized to address model heterogeneity and objective heterogeneity through knowledge distillation. In the global aggregation phase, a multi-factor weighted aggregation strategy is designed to measure the contribution of each local model to the global model, to enhance the fairness of aggregation. Three sets of experiments were conducted on two real datasets, and the experimental results demonstrate that the proposed model outperforms the baseline models in three common evaluation metrics.