Non-parametric regularization for class imbalance federated medical image classification

TL;DR

This work tackles the challenge of robust medical image classification under severe class imbalance in federated settings by introducing non-parametric regularization (NPR). NPR models class-specific structure through per-class sub-clusters and uses Sinkhorn-based online clustering to regularize the feature extractor, forming the basis of FedNPR and a personalized variant FedNPR-Per. Across skin lesion and intracranial hemorrhage tasks, FedNPR outperforms state-of-the-art FL methods, and NPR consistently improves existing FL approaches when added as a module. The findings suggest NPR as a practical, compatible tool to enhance federated learning under clinical data heterogeneity and imbalance, enabling more robust, transferable representations.

Abstract

Limited training data and severe class imbalance pose significant challenges to developing clinically robust deep learning models. Federated learning (FL) addresses the former by enabling different medical clients to collaboratively train a deep model without sharing privacy-sensitive data. However, class imbalance worsens due to variation in inter-client class distribution. We propose federated learning with non-parametric regularization (FedNPR and FedNPR-Per, a personalized version of FedNPR) to regularize the feature extractor and enhance useful and discriminative signal in the feature space. Our extensive experiments show that FedNPR outperform the existing state-of-the art FL approaches in class imbalance skin lesion classification and intracranial hemorrhage identification. Additionally, the non-parametric regularization module consistently improves the performance of existing state-of-the-art FL approaches. We believe that NPR is a valuable tool in FL under clinical settings.

Figures (4)

• Figure 1: Variations of class distributions under the federated setting between: 1) different clients, where some clients have missing and rare classes, and 2) each individual client and the global federation.
• Figure 2: Illustration of federated learning with non-parametric regularization (FedNPR). FedNPR has two components: 1) non-parametric regularization, where each client first cluster its features into a K $\times$$\#$ class clusters, which serve as anchors to regularize the extracted features, and 2) class-prior calibration according to each client's local class distribution. The regularization components are applied during the local update process, ensuring that the learned features capture the underlying data distribution.
• Figure 3: Illustration of the steps in non-parametric regularization (NPR) module: 1) online clustering of each client's extracted features into $KC$ sub-clusters with sinkhorn iteration sinkhorn, and 2) feature regularization using the center of the sub-clusters where we push the feature of each sample to be close to the corresponding sub-clusters and away from the others.
• Figure 4: Illustrations of the differences between FedNPR and FedNPR-Pe. FedNPR-Per keeps a personalized classifier head at each client instead of sharing it. Therefore, each client is able to adapt and learn its own classifier which is more aligned with its data characteristics.