FedBA: A Traffic Prediction Approach Based on Bi-LSTM and Federated Learning

Artificial intelligence (AI) traffic prediction technology is key to achieving 6G network management and optimization, enhancing network efficiency, and automating processes to meet diverse service demands. Deep learning models have improved traffic prediction accuracy through time series modeling and spatial feature extraction. However, centralized training methods face issues like bandwidth consumption and network capacity depletion. Federated Learning (FL) offers a distributed training architecture for mobile network prediction models. However, FL faces challenges related to data heterogeneity and model generalization. To address these issues, we propose a traffic prediction model named FedBA, which integrates a Bidirectional Long Short-Term Memory (Bi-LSTM) network with an attention mechanism and Federated Learning. To manage data heterogeneity, a transfer strategy based on historical data is employed. The model trained on historical datasets is treated as a quasi-global model, transferred to a central server, and shared among all base stations. An iterative clustering algorithm is also introduced to group base stations based on correlations between different traffic patterns. A dual-attention model aggregation scheme is applied to enhance global model generalization by focusing on the unequal contributions of different local models and the quasi-global model. Experimental results demonstrate that FedBA outperforms existing Federated Learning models like FedDA. For example, on the Milano and Trento datasets, FedBA reduces prediction errors for SMS, Internet, and Call traffic by 16.9%, 11.5%, 21.2%, and 7.7%, 44.1%, 40%, respectively.