FedDBL: Communication and Data Efficient Federated Deep-Broad Learning for Histopathological Tissue Classification

TL;DR

FedDBL is a novel federated learning framework that achieves high-performance histopathological tissue classification with limited data and only one communication round, significantly reducing data and communication costs while ensuring privacy.

Contribution

The paper introduces FedDBL, a lightweight federated learning framework that requires only one communication round and limited training data, improving efficiency and privacy in medical image classification.

Findings

Outperforms competitors with one-round communication and limited data

Reduces communication from 4.6GB to 276.5KB per client

Achieves comparable performance to multi-round methods

Abstract

Histopathological tissue classification is a fundamental task in computational pathology. Deep learning-based models have achieved superior performance but centralized training with data centralization suffers from the privacy leakage problem. Federated learning (FL) can safeguard privacy by keeping training samples locally, but existing FL-based frameworks require a large number of well-annotated training samples and numerous rounds of communication which hinder their practicability in the real-world clinical scenario. In this paper, we propose a universal and lightweight federated learning framework, named Federated Deep-Broad Learning (FedDBL), to achieve superior classification performance with limited training samples and only one-round communication. By simply associating a pre-trained deep learning feature extractor, a fast and lightweight broad learning inference system and a classical federated aggregation approach, FedDBL can dramatically reduce data dependency and improve communication efficiency. Five-fold cross-validation demonstrates that FedDBL greatly outperforms the competitors with only one-round communication and limited training samples, while it even achieves comparable performance with the ones under multiple-round communications. Furthermore, due to the lightweight design and one-round communication, FedDBL reduces the communication burden from 4.6GB to only 276.5KB per client using the ResNet-50 backbone at 50-round training. Since no data or deep model sharing across different clients, the privacy issue is well-solved and the model security is guaranteed with no model inversion attack risk. Code is available at https://github.com/tianpeng-deng/FedDBL.

Figures

Tables

Table 1. TABLE I: Statistics of MC-CRC. #1 denotes TCGA, #2 denotes Kather, #3 denotes Guangdong Provincial People’s Hospital and #4 denotes Yunnan Cancer Hospital.

A DI

B ACK

D EB

L YM

M UC

M US

N ORM

S TR

T UM

Equations

$$W_{client}=\underset{W_{init}}{\mathrm{argmin}}{\Vert \mathbf{B}{W}_{init}-Y\Vert }_2^2+\gamma {\Vert W_{init}\Vert }_2^2$$

$$W_{client}={\mathbf{B}}^{+}Y=\underset{\lambda \to 0}{\lim }{\left(\mathbf{B}{\mathbf{B}}^{\mathbf{T}}+\lambda E\right)}^{-1}{\mathbf{B}}^{\mathbf{T}}Y$$

$$W_{global}=\sum _{k=1}^K\frac{n_k}{N}{W}_{client}^k$$

Full text

I Introduction

Tissue classification [gurcan2009histopathological, FUCHS2011515], also known as tissue phenotyping, aims to use computer algorithms to automatically recognize different tissue types in the Whole Slide Images (WSIs). It is one of the fundamental tasks in computational pathology [srinidhi2021deep, wang2019weakly] which can parse the landscape of tumor microenvironment for precise predictions of cancer diagnosis [bulten2022artificial], prognosis [fu2020pan, pages2018international] and treatment response [vanguri2022multimodal]. With the advancements of deep learning algorithms and the growing number of open data [kather2019predicting, zhao2020artificial, yang2022two], this problem has been well studied with outstanding classification performance [hatami2021deep]. In clinical practice, however, it still faces ethical, regulatory and legal obstacles where centralized data collection may lead to privacy leakage, especially the RAW data.

Federated Learning (FL) [yang2019federated, le2021federated] framework provides a promising solution to protect user privacy by only sharing the intermediate results or the model parameters instead of the raw data, which has been widely studied in medical image analysis [pati2022federated, sheller2020federated]. But only very few attempts [saldanha2022swarm, shen2022tmi, ke2021isbi] have been made in computational pathology and the research progress still lags behind other medical image modalities [rauniyar2022federated] due to the following two obstacles.

The first one is the data dependency problem. Since most of the existing FL frameworks are constructed based on deep learning models. They are data-hungry and commonly require a large amount of well-annotated samples. However, labeling histopathological images is time-consuming, expertise-dependent and expensive [greenwald2022whole, pati2021reducing]. When without enough training samples, existing models may not achieve favorable performance. Another obstacle is the communication overhead. The training procedure of traditional FL models needs multiple cloud-client iterations to achieve global convergence. However, deep learning models are with tens of millions of parameters, which greatly increases the communication burden when with multiple communication rounds. Although some recent works use self-supervised learning SSL methods with unlabelled data to reduce the demand of labelled data [yan2023label, kassem2022federated], they still need multiple communication rounds to train a stable domain-specific pre-trained model. Lack of training samples may further amplify the communication burden because deep learning models commonly require more iterations to converge when with limited training samples. Moreover, frequent communications may increase the chance of being attacked, such as man-in-the-middle attacks [wang2020man].

Therefore, it is urged to construct a data-efficient and communication-efficient FL model for histopathological tissue classification. In this paper, we proposed a simple and effective solution for histopathological tissue classification, which considers not only the data sharing problem, but also the data dependency, communication efficiency, model robustness and model inversion attack. Our proposed model Federated Deep-Broad Learning, FedDBL in short, contains three integrated components, including a common federated learning framework, a pre-trained deep learning (DL) backbone and a broad learning (BL) inference system [BLS, gong2022research]. The federated learning framework serves for decentralized training to avoid data sharing across different medical centers or institutions. The pre-trained DL backbone can provide stable and robust deep features when there are not enough training labels even with domain-irrelevant pre-trained DL backbone. It can also effectively avoid the model inversion attack since no back-propagation is calculated for gradients [zhu2019deep]. The BL system is a lightweight classifier with good approximation capability which can greatly shorten the transmission time and overcome the data dependency problem. Fig. 1 comprehensively demonstrates the strengths of FedDBL compared with the centralized learning and the conventional federated learning ways.

Extensive experiments with five-fold cross-validation are conducted to demonstrate the superiority of FedDBL in several aspects, including data dependency, communication efficiency, flexibility and the practicability of the model encryption. When with enough training data, FedDBL can mostly outperform conventional FL strategies and achieve comparable or even better classification performance compared with centralized learning strategy. When reducing the training samples in the data dependency experiment, FedDBL still maintains a high-level performance and greatly outperforms both centralized learning and conventional FL frameworks, even with only 1% training samples. FedDBL is also flexible to any deep learning architectures to support data- and communication-efficient histopathological tissue classification. Another spotlight of FedDBL is communication efficiency. Compared with the conventional FL frameworks, FedDBL’s one-round training manner reduces the upload workload from 4.609GB to 276.5KB (over 17,000 times faster) with ResNet-50 backbone compared to traditional 50-round iterative training. Thanks to the tiny model size, FedDBL is also computationally efficient in model encryption which can further upgrade the privacy protection level. The main contributions of this paper can be summarized as follows:

• •

We propose a novel federated learning approach (FedDBL) for histopathological tissue classification to preserve patients’ privacy.

• •

To the best of our knowledge, FedDBL is the first study that considers communication efficiency and data efficiency simultaneously which reduces the communication overhead of each client by around 17,000 times on ResNet-50 with extremely limited training samples (only 1%).

• •

FedDBL is a simple, effective and easy-to-use algorithm that associates three classical modules, including a robust pre-trained deep learning feature extractor, a fast broad learning inference system and a simple federated learning framework. It is highly extendable that allows to replace any module with a more advanced one.

• •

Extensive experiments demonstrate that FedDBL drastically relieves the dependence on training data and reduces the communication overhead while maintaining outstanding classification performance, which promotes its clinical practicability.

II Related Works

II-A Histopathological Tissue Classification

High-resolution WSIs offer a wide range of tissue phenotypes where the pixel-level annotation is time-consuming and requires a great deal of biomedical knowledge ([srinidhi2021deep]), making patch-level histopathological tissue classification an alternate solution for automated analysis in computer-aided tumor diagnosis ([kather2019predicting, xue2021selective, abdeltawab2021pyramidal]).

Due to the rapid development of computer vision, the most popular natural image classification models can be transferred into histopathological tissue phenotyping. However, it still suffers from the data dependency problem with a huge annotation burden ([ayyad2021role]). Thus, various approaches have been proposed to reduce the annotation effort. [han2022multi] proposed a multi-layer pseudo-supervision approach with a progressive dropout attention mechanism to convert patch-level labels into pseudo-pixel-level labels. An extra classification gate mechanism was presented which reduced the false-positive rate for non-predominant category classification and improved the segmentation performance in return. [xue2021selective] utilized a generative adversarial network (GAN) to generate pseudo samples to expand the training data. [dolezal2022uncertainty] cropped WSIs into tiles for training the uncertainty quantization model and solved the problem of domain shift in external validation data. In order to get rid of lacking image annotations, [wang2022transformer] employed unsupervised contrastive learning to obtain a robust initialized model with moderate feature representation of the histopathological feature space, with no annotation burden. Our previous study ([lin2022pdbl]) introduced pyramidal deep-broad learning (PDBL) as a pluggable module for any CNN backbone to further improve histopathological tissue classification performance.

Besides that, another unexplored challenge is the patient privacy issue. Only a few attempts ([saldanha2022swarm, saldanha2022direct]) have been made in federated learning for computational pathology, which will be discussed in the following subsection. And to the best of our knowledge, we are the first study to consider privacy protection in histopathological tissue classification.

II-B Federated Learning

II-B1 Federated Learning in Medical Image Analysis

Because of the ethical issue, federated learning (FL) has been widely adopted in medical applications to preserve the patients’ privacy ([pati2022federated, warnat2021swarm, sheller2020federated]). In medical imaging, FL has witnessed a boost in interest ([kaissis2020secure]), such as MRI reconstruction ([guo2021multi, li2020multi]), CT lesion segmentation ([yang2021federated]) and etc. In the COVID-19 pandemic, COVID-19-related applications with data from different medical centers or even from different countries become the most urgent demand in the real-world clinical scenario while FL greatly advances the diagnostic performance ([bai2021advancing]). [dayan2021federated] used 20 institutes’ data across the global for predicting the future oxygen requirements of symptomatic patients suffering from COVID-19. [dou2021federated] proposed a federated model to detect COVID-19 lung abnormalities with good generalization capability on unseen multinational datasets.

II-B2 Federated Learning in Computational Histopathology

In histopathological images, a swarm learning architecture with blockchain protocols has been proposed to predict the mutational status ([saldanha2022swarm]). However, compared with other medical imaging modalities, there are few studies ([saldanha2022direct]) that adopt federated learning in histopathological images for the following reasons. First, the digitalization of pathology is unpopular. Pathological diagnosis still relies on observing specimens under a microscope. Second, image annotation is also an obstacle for the histopathological image process since only pathologists are capable to label WSIs which greatly increases the difficulties of acquiring well-annotated data. Third, due to the gigapixel resolution of WSIs, the size of the deep learning model is generally large, which increases the communication burden in networking.

There are technical solutions in FL to the high communication overhead problem, such as compressing the model size ([reisizadeh2020fedpaq, jhunjhunwala2021adaptive]). [reisizadeh2020fedpaq] proposed FedPAQ to reduce the interactive overhead of FL by compressing the model with lower bit-precision and [jhunjhunwala2021adaptive] proposed an adaptive quantization strategy to achieve communication efficiency.

However, the underlying assumption of existing studies is that there should be enough samples for model training where they may not be able to take into account both communication efficiency and limited data issue ([kamp2021federated, zhang2023two]). In this study, we fully consider the specialty of histopathological images, the difficulties of data labeling and the communication efficiency in the real-world clinical scenario, which has never been discussed in decentralized computational pathology.

III Methodology

In this section, we introduce our framework Federated Deep-Broad Learning (FedDBL). This framework is designed for privacy-preserving tissue classification with limited training samples and extremely low communication overhead. In the following subsections, we first describe the intuitive thinking and problem setting in Section III-A. The overall framework and the methodology of FedDBL are shown in Section III-B. Finally, we demonstrate the implementation details in Section III-C.

III-A Problem Setting

As a classical upstream task in computational pathology, existing tissue classification approaches have achieved outstanding performance under an ideal condition with enough training samples by centralized learning. However, they might face the following obstacles in the real-world clinical scenario.

Annotation burden: Collecting enough well-labeled training samples is expensive and time-consuming because it requires labelers with medical background.

Privacy preservation: The raw data should not be shared across different medical institutions (or clients) to preserve the patient’s privacy. Transmitting raw data may break the principle of medical ethics.

Communication cost: The communication overhead has always been a challenge in federated learning models affected by many compound factors, such as the model size, the communication rounds, the model convergence speed, the network bandwidth and etc.

To resolve the aforementioned challenges, we propose a simple and effective FL-based framework, demonstrated in Fig. 2. First, we abandon conventional end-to-end training manner since limited training samples may harm the robustness of the deep learning model and decrease the convergence speed. Therefore, we separate feature extraction and inference for local training in each client. A pre-trained deep feature extractor (CNN backbone) is introduced to avoid the feature extractor being affected by the training sample bias from different clients in order to guarantee the robustness of extracted features. Then an independent broad learning inference system [BLS, lin2022pdbl] serves for fast inference. Finally, we apply a classical weighted averaging as in FedAvg [mcmahan2017communication], to fuse the broad learning inference systems from all the clients.

III-B FedDBL Architecture and Formulation

As shown in Fig. 2, FedDBL consists of three modules, deep feature extraction module (DL-module), broad inference module (BL-module) and federated decentralized module (Fed-module). DL-module together with BL-module serves for local training on the client side. Fed-module is executed on the server side. Algorithm 1 provides the details of the entire FedDBL pipeline.

Let $\mathcal{D}_{1},\mathcal{D}_{2},\cdots,\mathcal{D}_{k},\cdots,\mathcal{D}_{K}$ denote the local training sets from $K$ clients with the dataset size of $n_{k}$ for each client $\mathcal{D}_{k}$. The total number of training samples is denoted as $N=\sum_{k=1}^{K}{n_{k}}$. For each sample $X$ with ground truth $Y$ in $\mathcal{D}$, DL-module with pre-trained parameters $\Theta$ extracts the features and stores them in the local feature bank $\mathbf{B}$. Then BL-module calculates the weights $W_{client}$ of broad learning system. By the federated aggregation approach, we can obtain the global weight $W_{global}$. The workflows of the server and the clients are demonstrated in Algorithm 1 and Algorithm 2, respectively.

III-B1 Deep Feature Extraction Module

A large number of samples and repeated backpropagation are required in standard DL training to achieve a good feature representation ability. When suffering from the insufficient data problem, the model training procedure might be unstable which leads to poor feature representation and model overfitting. Our previous study [lin2022pdbl] reveals that directly adopting a stable pre-trained model for feature extraction is more favorable to the model performance than training the model with limited samples, even the pre-trained model was trained by an irrelevant image domain (ImageNet111https://image-net.org/). Inspired by this idea, we use a pre-trained CNN model with no further training to extract the deep features. Notice that, the selection of the pre-trained models is flexible, and can be from any image domain. We have conducted an experiment to justify the flexibility in Section IV. Another advantage of using pre-trained models is to avoid model inverse attacks since the training samples are all unseen. To enrich the feature representation, we extract multi-stage features from low-level to high-level, details as below.

As illustrated in DL-module of Algorithm 2, each client $k$ $(k\in\left[1,\cdots,K\right])$ downloads the pre-trained DL backbone as feature extractor $\Theta$ and extracts multi-stage deep features $\mathbf{b}_{X}$ of training sample $X$ locally (we omit $k$ for simplicity), where $\mathbf{b}_{X}$ consists of multiple stages’ features $\Theta^{s}(X)(s\in\left[1,\cdots,S\right])$. The features of the entire dataset $\mathcal{D}_{k}$ are stored in the local feature bank $\mathbf{B}$. Then the local feature bank $\mathbf{B}$ will be passed to broad inference module. Since neither the training data nor the feature bank is shared across different clients, there is no privacy leakage risk for the RAW data in deep feature extraction module.

III-B2 Broad Inference Module

With the local feature bank $\mathbf{B}$, each client $k$ can conduct a local BL system [BLS] through BL-module (Algorithm 2) for fast inference. By solving the Eq. (1) optimization problem, an optimal BL model $W_{client}$ can be obtained rapidly through the pseudo-inverse method (Eq. (2)).

[TABLE]

[TABLE]

where $Y$ represents the ground-truth label matrix, $\mathbf{B}$ is local feature bank in the form of matrix. $W_{init}$ is the initialized broad learning weights. $E$ is the identity matrix, $\lambda$ is a constant parameter and $\gamma$ is the regularization parameter. The pseudo-inverse method of solving BL model considerably reduces the computational burden while achieving high communication efficiency. For the inference process, the predicted results can be calculated by $Y_{test}=\mathbf{B}_{test}W_{client}$ after extracting test samples’ deep features with the largest probabilistic value.

Thanks to the lightweight broad learning model $W_{client}$, the communication efficiency is drastically improved compared with the conventional DL-based FL frameworks.

III-B3 Federated Decentralized Module

In this module, we conduct a federated learning framework for decentralized learning. Given the broad learning model $W_{client}^{k}$ of each client $k$, we first upload the models from all the clients to the central server. And then general federated aggregation methods can be applied to aggregate them. Here, we use the most common weighted averaging way for model aggregation as adopted in FedAvg [mcmahan2017communication], FedProx [li2020federated] and FedPAQ [reisizadeh2020fedpaq].

[TABLE]

where $W_{global}$ is the global model from the server, $n_{k}$ is the number of training samples in client $k$ and $N$ is the total number of training samples. A larger training dataset will contribute more to the global model. Since we only share the broad learning model for once, the communication efficiency and the patient’s privacy are guaranteed.

III-C Implementation Details

All of our experiments are implemented in Pytorch on a workstation with an NVIDIA RTX 3090 and the i9-11900K CPU with 16 cores. We use the cross-entropy loss for the baseline centralized training with batch size $20$. The SGD optimizer is set as follows: the learning rate is $1e^{-3}$, the weight decay is $1e^{-4}$ and the momentum is 0.9. The patches are $224\times 224$ under $20\times$ WSIs. Different client numbers are used depending on the datasets.

We adopt three well-known federated aggregation methods, FedAvg [mcmahan2017communication], FedProx [li2020federated] and FedPAQ [reisizadeh2020fedpaq], for comparison. And the centralized model is trained as the baseline. FedProx has the parameter $\mu$ to adjust the effect of the proximal term on the loss function. Here we set $\mu$ as $1$ which has better performance.

IV Experiments

In this section, we present the details of the datasets and conduct various experiments to demonstrate the performance and efficiency of the proposed FedDBL. Section IV-A shows two open datasets and the experimental settings in the federated learning framework. In Section LABEL:sub:one-round, we compare FedDBL with centralized learning baselines, conventional federated learning baselines and one-round federated learning baselines. The effectiveness is comprehensively discussed in Section LABEL:sub:multi-round. We use Matthews Correlation Coefficient (MCC), Accuracy and F1-score as the evaluation metrics in all the experiments.

IV-A Datasets and Experimental Settings