FedMLP: Federated Multi-Label Medical Image Classification under Task Heterogeneity

TL;DR

This work tackles multi-label medical image classification in a federated setting with task heterogeneity due to partial annotations. It proposes FedMLP, a two-stage framework that combines a warm-up using the $L^k_{WPC}$ loss with logit adjustment and a prototype-guided missing-label detection pipeline, followed by a global-consistency regularization via a teacher model. Key mechanisms include dual-class prototypes $P^{k,c}_n$, missing-label sets $\,\mathbb{P}^k_{c,0}$ and $\,\mathbb{P}^k_{c,1}$, and self-adaptive thresholds $ au_0$, $ au_1$ derived from per-class learning degrees $d^k_c$ and $d^G_c$. Experiments on RSNA ICH and ChestXray14 show FedMLP outperforming state-of-the-art federated semi-supervised and noisy-label methods across varying missing-label rates, highlighting its practical potential for real-world clinical FL with heterogeneous data.

Abstract

Cross-silo federated learning (FL) enables decentralized organizations to collaboratively train models while preserving data privacy and has made significant progress in medical image classification. One common assumption is task homogeneity where each client has access to all classes during training. However, in clinical practice, given a multi-label classification task, constrained by the level of medical knowledge and the prevalence of diseases, each institution may diagnose only partial categories, resulting in task heterogeneity. How to pursue effective multi-label medical image classification under task heterogeneity is under-explored. In this paper, we first formulate such a realistic label missing setting in the multi-label FL domain and propose a two-stage method FedMLP to combat class missing from two aspects: pseudo label tagging and global knowledge learning. The former utilizes a warmed-up model to generate class prototypes and select samples with high confidence to supplement missing labels, while the latter uses a global model as a teacher for consistency regularization to prevent forgetting missing class knowledge. Experiments on two publicly-available medical datasets validate the superiority of FedMLP against the state-of-the-art both federated semi-supervised and noisy label learning approaches under task heterogeneity. Code is available at https://github.com/szbonaldo/FedMLP.

Figures (2)

• Figure 1: Illustration of two decentralized data settings.
• Figure 2: Overview of FedMLP. Colors red and green represent missing and labeled classes respectively.