Federated Learning Based on Mutual Information Clustering for Wireless Traffic Prediction

Wireless traffic prediction can help operators accurately predict the usage of wireless networks, and it plays an important role in the load balancing and energy saving of base stations. Currently, most traffic prediction methods are centralized learning strategies, which need to transmit a large amount of traffic data and have timeliness and data privacy issues. Federated learning, as a distributed learning framework with no client data sharing and multi-client collaborative training, can solve such problems. We propose a federated learning wireless traffic prediction framework based on mutual information clustering (FedMIC). First, a sliding window scheme is used to construct the raw data into adjacent and periodic dual-traffic sequences and capture their traffic characteristics separately to enhance the client model learning capability. Second, clients with similar traffic data distributions are clustered together using a mutual information-based spectral clustering algorithm to facilitate the capture of the personalized features of each clustered model. Then, models are aggregated using a hierarchical aggregation architecture of intra-cluster model aggregation and inter-cluster model aggregation to address the statistical heterogeneity challenge of federated learning and to improve the prediction accuracy of models. Finally, an attention mechanism-based model aggregation algorithm is used to improve the generalization ability of the global model. Experimental results show that our proposed method minimizes the prediction error and has superior traffic prediction performance compared to traditional distributed machine learning methods and other federated learning methods.