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

2025-02-24

AI Summary

To address classifier and feature extractor bias induced by data heterogeneity in medical federated learning, this paper proposes FedBM. First, it leverages a pre-trained language model (PLM) to capture semantic distributions of medical concepts and constructs a language-guided unbiased classifier. Second, it designs a concept-conditional generative adversarial network to collaboratively calibrate feature extractors across clients via synthetically generated pseudo-samples. FedBM is the first framework to integrate PLMs, concept embeddings, and generative data augmentation into a dual-module architecture, enhanced by prompt engineering and triple regularization—enabling frozen-classifier optimization and enforcing feature-space consistency. Extensive experiments on multiple public medical imaging benchmarks demonstrate significant improvements over state-of-the-art methods. Ablation studies validate the efficacy of each component, and the source code is publicly available.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Addresses local learning bias in federated learning
• Leverages pre-trained language models for classification
• Improves global model performance with pseudo data

Innovation

• Exploits pre-trained language models
• Constructs high-quality classifier
• Generates pseudo data calibration