Integrated probabilistic clustering and Deep Reinforcement Learning for bias mitigation and device heterogeneity of Federated Learning in edge networks

Abstract

Federated Learning (FL) enables decentralized and collaborative training on resource-constrained Edge Devices (EDs) while preserving data privacy by avoiding raw data transmission. However, traditional FL approaches face challenges such as non-independent and identically distributed (non-IID) data, biased model aggregation due to device heterogeneity, and inefficiencies caused by stragglers during model updates. We propose a novel Hierarchical Deep Reinforcement Learning-based Probabilistic Federated Learning (Hier-FedDRL) strategy to address these limitations. This framework combines local and central Deep Reinforcement Learning (DRL) agents with a probabilistic clustering approach to manage heterogeneous devices and optimize resource allocation dynamically. Local DRL agents optimize intra-cluster operations, including training and resource distribution, while the central DRL agent oversees global model updates and inter-cluster coordination.

To ensure balanced aggregation and mitigate biases, the proposed framework employs Gaussian Mixture Models (GMMs) for clustering EDs based on their data distributions and resource characteristics. Additionally, a dynamic contribution-based aggregation technique is introduced to fairly weigh updates from diverse EDs, reducing biases in the global model. The performance of Hier-FedDRL is evaluated in a cloud-based setup, where Docker containers are used to simulate EDs and Google Kubernetes Engine clusters for cloud orchestration. Experimental results over benchmark datasets demonstrate that the proposed Hier-FedDRL achieves 4%–6% higher accuracy, reduces convergence time by 7%–10%, and lowers bias in the global model by 25%, outperforming state-of-the-art FL approaches while effectively addressing data and resource heterogeneity.