Clustered federated learning with weighted model aggregation for imbalanced data

As a promising edge learning framework in future 6G networks, federated learning (FL) faces a number of technical challenges due to the heterogeneous network environment and diversified user behaviors. Data imbalance is one of these challenges that can significantly degrade the learning efficiency. To deal with data imbalance issue, this work proposes a new learning framework, called clustered federated learning with weighted model aggregation (weighted CFL). Compared with traditional FL, our weighted CFL adaptively clusters the participating edge devices based on the cosine similarity of their local gradients at each training iteration, and then performs weighted per-cluster model aggregation. Therein, the similarity threshold for clustering is adaptive over iterations in response to the time-varying divergence of local gradients. Moreover, the weights for per-cluster model aggregation are adjusted according to the data balance feature so as to speed up the convergence rate. Experimental results show that the proposed weighted CFL achieves a faster model convergence rate and greater learning accuracy than benchmark methods under the imbalanced data scenario.