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

TL;DR

FedBM is a novel federated learning framework that leverages pre-trained language models and concept-guided distribution estimation to mitigate local bias and improve performance in heterogeneous medical image classification tasks.

Contribution

The paper introduces FedBM, a new framework combining linguistic knowledge and distribution estimation to reduce local bias in heterogeneous federated learning, with theoretical and empirical validation.

Findings

FedBM outperforms state-of-the-art methods on public datasets.

The modules LKCC and CGDE effectively reduce local bias.

The approach improves classification accuracy in heterogeneous settings.

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.