FedQMIX: Communication-efficient federated learning via multi-agent reinforcement learning

Abstract

Since the data samples on client devices are usually non-independent and non-identically distributed (non-IID), this will challenge the convergence of federated learning (FL) and reduce communication efficiency. This paper proposes FedQMIX, a node selection algorithm based on multi-agent reinforcement learning(MARL), to address these challenges. Firstly, we observe a connection between model weights and data distribution, and a clustering algorithm can group clients with similar data distribution into the same cluster. Secondly, we propose a QMIX-based mechanism that learns to select devices from clustering results in each communication round to maximize the reward, penalizing the use of more communication rounds and thereby improving the communication efficiency of FL. Finally, experiments show that FedQMIX can reduce the number of communication rounds by 11% and 30% on the MNIST and CIFAR-10 datasets, respectively, compared to the baseline algorithm(Favor).