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

TL;DR

This work addresses data heterogeneity in federated learning for medical imaging by introducing FedMRL, which combines a MARL-driven adaptive proximal term $\mu$, a fairness loss $L_{fair}$ to balance cross-client performance, and SOM-based server-side aggregation to handle distribution shifts. The proposing components enable personalized regularization and distribution-aware weight updates, with QMIX guiding cooperative decisions across hospitals. Empirical results on ISIC-2018 and Messidor show FedMRL outperforms state-of-the-art baselines (FedAvg, FedProx, FedNova, FedBN) in ACC and AUC under non-IID conditions. While effective, the framework may face scalability and computational overhead in large-scale deployments, motivating future work on distributed and resource-sharing strategies across diverse domains.

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 $(μ)$ 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}}.

Figures (2)

• Figure 1: The proposed architecture comprises clients and a server: (a) Clients transmit their model weights to the global server, while agents corresponding to clients retrieve the corresponding state $s_i$ from the global state $s$ and compute the respective $\mu_i$ value, subsequently sharing it with the corresponding hospitals. (b) Global weight aggregation is facilitated using SOM, where $\alpha_i$ denotes the weight adjustment factor, and $w_i$ represents the local model weights utilized to derive the final global model $w_{t+1}$ for subsequent communication rounds.
• Figure 2: The plot of the proposed loss function considering the 2 clients. It is clear from the graph that the minimum of the loss function occurs when both clients have the same loss value, i.e., tending to zero.