FedMRL: Data Heterogeneity Aware Federated Multi-agent Deep Reinforcement Learning for Medical Imaging

TL;DR

FedMRL is a federated multi-agent deep reinforcement learning framework that effectively handles data heterogeneity in medical imaging by incorporating fairness, personalized objectives, and adaptive weighting, leading to improved performance over existing methods.

Contribution

This paper introduces FedMRL, a novel framework combining MARL and adaptive weighting to address data heterogeneity in federated learning for medical imaging.

Findings

FedMRL outperforms state-of-the-art methods on real-world datasets.

The framework effectively mitigates data distribution shifts.

Incorporates a fairness loss to prevent model bias.

Abstract

Despite recent advancements in federated learning (FL) for medical image diagnosis, addressing data heterogeneity among clients remains a significant challenge for practical implementation. A primary hurdle in FL arises from the non-IID nature of data samples across clients, which typically results in a decline in the performance of the aggregated global model. In this study, we introduce FedMRL, a novel federated multi-agent deep reinforcement learning framework designed to address data heterogeneity. FedMRL incorporates a novel loss function to facilitate fairness among clients, preventing bias in the final global model. Additionally, it employs a multi-agent reinforcement learning (MARL) approach to calculate the proximal term $(\mu)$ for the personalized local objective function, ensuring convergence to the global optimum. Furthermore, FedMRL integrates an adaptive weight adjustment method using a Self-organizing map (SOM) on the server side to counteract distribution shifts among clients' local data distributions. We assess our approach using two publicly available real-world medical datasets, and the results demonstrate that FedMRL significantly outperforms state-of-the-art techniques, showing its efficacy in addressing data heterogeneity in federated learning. The code can be found here~{\url{https://github.com/Pranabiitp/FedMRL}}.