Federated Model Aggregation via Self-Supervised Priors for Highly Imbalanced Medical Image Classification

TL;DR

This paper introduces Fed-MAS, a federated learning method that uses self-supervised priors to improve model aggregation in highly imbalanced medical image classification tasks, capturing intra-class variations across clients.

Contribution

It proposes a novel dynamic aggregation approach leveraging self-supervised auxiliary networks to address intra-class variations in federated medical imaging.

Findings

Effective handling of class imbalance in federated learning.

Robust and unbiased global model achieved across diverse medical datasets.

Utilizes self-supervised priors for improved model aggregation.

Abstract

In the medical field, federated learning commonly deals with highly imbalanced datasets, including skin lesions and gastrointestinal images. Existing federated methods under highly imbalanced datasets primarily focus on optimizing a global model without incorporating the intra-class variations that can arise in medical imaging due to different populations, findings, and scanners. In this paper, we study the inter-client intra-class variations with publicly available self-supervised auxiliary networks. Specifically, we find that employing a shared auxiliary pre-trained model, like MoCo-V2, locally on every client yields consistent divergence measurements. Based on these findings, we derive a dynamic balanced model aggregation via self-supervised priors (MAS) to guide the global model optimization. Fed-MAS can be utilized with different local learning methods for effective model aggregation toward a highly robust and unbiased global model. Our code is available at \url{https://github.com/xmed-lab/Fed-MAS}.