FedBM: Stealing Knowledge from Pre-trained Language Models for Heterogeneous Federated Learning

TL;DR

This paper tackles data heterogeneity in heterogeneous federated learning for medical image classification by introducing FedBM, a framework that uses linguistic knowledge to eliminate local learning bias. It consists of two modules: LKCC, which constructs a high-quality fixed classifier from concept embeddings derived from prompts and pre-trained language models, and CGDE, which trains a conditional generator to synthesize pseudo data that calibrates local feature extractors. The approach yields state-of-the-art performance across five public medical datasets under varying non-IID settings, with extensive ablations validating the necessity and effectiveness of both modules. The work advances privacy-preserving, robust FL for medical imaging by grounding classifier alignment and data augmentation in linguistic priors, potentially impacting how heterogeneous FL systems are designed in practice.

Abstract

Federated learning (FL) has shown great potential in medical image computing since it provides a decentralized learning paradigm that allows multiple clients to train a model collaboratively without privacy leakage. However, current studies have shown that data heterogeneity incurs local learning bias in classifiers and feature extractors of client models during local training, leading to the performance degradation of a federation system. To address these issues, we propose a novel framework called Federated Bias eliMinating (FedBM) to get rid of local learning bias in heterogeneous federated learning (FL), which mainly consists of two modules, i.e., Linguistic Knowledge-based Classifier Construction (LKCC) and Concept-guided Global Distribution Estimation (CGDE). Specifically, LKCC exploits class concepts, prompts and pre-trained language models (PLMs) to obtain concept embeddings. These embeddings are used to estimate the latent concept distribution of each class in the linguistic space. Based on the theoretical derivation, we can rely on these distributions to pre-construct a high-quality classifier for clients to achieve classification optimization, which is frozen to avoid classifier bias during local training. CGDE samples probabilistic concept embeddings from the latent concept distributions to learn a conditional generator to capture the input space of the global model. Three regularization terms are introduced to improve the quality and utility of the generator. The generator is shared by all clients and produces pseudo data to calibrate updates of local feature extractors. Extensive comparison experiments and ablation studies on public datasets demonstrate the superior performance of FedBM over state-of-the-arts and confirm the effectiveness of each module, respectively. The code is available at https://github.com/CUHK-AIM-Group/FedBM.

Figures (4)

• Figure 1: Data heterogeneity causes local learning bias, including classifier bias and feature extractor bias. (Best viewed in color)
• Figure 2: The overview of the proposed FedBM framework. FedBM contains Linguistic Knowledge-based Classifier Construction (LKCC) and Concept-guided Global Distribution Estimation (CGDE). LKCC uses class concepts, prompts and PLMs to build latent concept distributions, which are sent to clients as local classifiers. CGDE samples probabilistic concept embeddings from the distributions to train a conditional generator. The generator is shared by all clients and produces pseudo data to calibrate updates of local feature extractors. (Best viewed in color)
• Figure 3: The performance of our method with different batch sizes of generated samples on OCT-C8 and Kvasir-v2 datasets.
• Figure 4: The performance of our method with different client numbers on Kvasir-v2 and OCT-C8 datasets.