Mitigating client heterogeneity in federated learning: a reinforcement learning approach for efficient client selection

Abstract

Federated learning (FL) is a highly promising distributed machine learning method that enables collaborative model training while protecting participants’ data. In Internet of Things (IoT) applications, due to the large number of devices and high communication overhead, selecting a subset of clients for each training round has become a common requirement. However, a key challenge with this strategy lies in balancing the trade-off between system heterogeneity and statistical heterogeneity during client selection to ensure the efficiency and stability of global training. To address this, we propose the FLASH-DRM client selection framework based on Double Deep Q-Learning (DDQL). This framework employs a utility-based reputation function to determine the contribution of individual clients to the global model. Additionally, it introduces a reward path that directly links the agent’s client selection to the global accuracy, enabling the agent to comprehensively consider the impact of both individual and collective client utilities on the global model. Notably, the combined utility integrated into FLASH-DRM incurs virtually no additional communication overhead while effectively mitigating the heterogeneity challenges during training. In various heterogeneous dataset environments, we trained FLASH-DRM on CIFAR-10, MNIST, and EMNIST test suites, and compared it with state-of-the-art algorithms. Experimental results show that FLASH-DRM generally outperforms FedAvg, FAVOR, FedGCS, FLASH-RL, FedProx and FedMSC in terms of accuracy, latency, and energy consumption, and the proposed FLASH-DRM is a promising technique for addressing system heterogeneity and statistical heterogeneity in the Federated Learning setting.