Graph Neural Networks for Diffusion and Aggregation in Wireless Federated Learning

Abstract

User devices (UDs) with non-independent and identically distributed (non-IID) data will worsen accuracy performance of the global model in federated learning (FL). Therefore, the implementation of diffusion strategies in machine learning (ML) models can enhance the effectiveness of federated learning with non-IID data. However, in a device-to-device (D2D) wireless federated learning (WFL) system, limited wireless resources and severe wireless channel interference become the important bottleneck to restrict the diffusion performance and model aggregation so as the global model of WFL with non-IID suffers from the weight divergence challenge. Thus, we propose a novel joint over-the-air computation (OAC) aggregation and diffusion framework by using a graph neural network (GNN) for WFL, termed an OAC-GNN-Dif framework. By integrating the OAC with message passing neural network (MPNN) of GNN, we further develop the OAC-MPNN-Dif algorithm based on the OAC-GNN-Dif framework. To further reduce communication costs, we designed an OAC message recurrent neural network (OAC-MPRNN-Dif) algorithm, where each UD propagates local models via D2D communications to refresh the graph embedding in the current frame based on the graph feature extraction and localization state of the previous frame to reduce communication costs. Additionally, we introduce dynamic time-varying MPNN for federated diffusion within evolving D2D network topologies. The experimental results indicate that our approach significantly performs well in communication overhead, with a 30%-60% decreasing in wireless resources overhead and 1.2-3.5 times decreasing in the number of model transfers compared to the FedDif methods. Moreover, our approach also improves the global model test accuracy, which is about 2.7% higher than the existing communication diffusion FL with non-IID characteristics.