Federated Learning with Privacy-Preserving Ensemble Attention Distillation

In parameter-based FL methods, each local model shares its parameters/gradients with the central server after each round of local training on its local data. The central server aggregates them by averaging [2]. Then the aggregated results are shared with the local nodes, updating their corresponding local model before proceeding with the next training round. This process is repeated until the stopping criterion is met. A variety of extensions of FedAvg [2, 7, 11, 12] employ improved aggregation schemes, such as adding momentum [4] and local weighting [11, 12]. Local weighting schemes have also been investigated based on client loss [11] and client data size [12]. FedMA [7] aggregates local parameters layer-wise by matching and averaging hidden elements. FedProx [5] incorporates a proximal term to restrict local updates close to the global model. SCAFFOLD [8] introduces control variations to correct the local updates.

Federated distillation methods exchange model output rather than model parameters. Solutions that produce central models by distilling knowledge from private data incur concerns w.r.t. the leaking of private local data [24, 25]. In contrast, some works [19, 17, 18] distill the output of public data. FedMD [17] divides local training and joint distillation into two phases by adding labeled public data. Each local model is fully trained with public and private data in the first phase. In the collaboration phase, the central and local models are jointly optimized to reach a consensus via the distillation of predictions on the public data. Cronous [18] implements a similar two-phase system, but the local initialization only utilizes private data since the public data is unlabeled. Although model-agnostic, these methods select public data based on prior knowledge w.r.t. private data. The recently proposed FedDF [20] makes it more robust to distillation data selection, but it still has privacy issues due to the iterative exchange of models over hundreds of rounds. The above-mentioned exclusively require many rounds of back-and-forth communication, leading to bandwidth bottlenecks and other inefficiencies. For communication efficiency, Guha et al. [26] proposes a preliminary investigation into offline federated learning where all local models can be trained independently without inter-institutional communication. It then distills through the averaging of predictions on unlabeled data. Another similar work is PATE [27], where the central model is trained with hard pseudo-labels voted by local models rather than the soft prediction we employed.

Model ensemble [28] is a popular regularization technique that is more robust and generalizable than individual models. Classical ensemble approaches such as bagging [29], boosting [30], and Bayesian [31] are exploited to assemble partial knowledge for better prediction as a mixture of experts [32]. With the success of knowledge transfer [33], recent advancements in ensemble networks are dominated by the student-teacher learning paradigm [34]. Ensemble learning aggregates the knowledge of multiple teachers before distilling the knowledge into the student network. Supervised ensemble learning is dominated by gate learning to design the weight for aggregation [34, 35, 36]. In semi-supervised and self-supervised scenarios, [37] and [38] exploit the relative similarity between samples. Furthermore, co-distillation extends one-way transfer to bidirectional collaborative learning [39, 40, 41].

Beyond the use of softened labels for distillation [33], many improvements have been developed to transfer structure knowledge, such as intermediate representations [42]. Attention transfer relaxes feature-level knowledge to attention maps either from bottom-up activation [43] or top-down gradients [44]. Yim et al. [45] utilizes the Gram matrix as a flow of solution procedure for distillation. Other works match features such as maximum mean discrepancy [46] or mutual information [47]. More recent attempts explore structured knowledge ensemble for knowledge distillation with labeled distillation data. FEED [48] can be seen as an extension of self-distillation [49] that accumulates and assembles feature-level knowledge to train itself recursively. Knowledge flow [50] enhances the student by adding transformed and scaled intermediate representations from teacher models. In contrast, our feature-level ensemble method is label-free, model agnostic, and is used for heterogeneous federated distillation scenarios.

Let $z^{c}_{k}=f(x_{0},\theta_{k},c)$ be the logits of a public data sample $x_{0}$ corresponding to class $c\in\mathcal{C}_{k}$, produced by the model at local node $k$, and the output of the central model be $\tilde{z}^{c}=f(x_{0},\theta_{\mathrm{s}},c)$, where $c\in\{1,\ldots,C\}$. The conventional ensemble $\hat{z}^{c}=\frac{1}{K}\sum_{k=1}^{K}z^{c}_{k}$ takes an average of all teachers’ logits and then employ activation $\sigma(\cdot,\cdot)$ using softmax to represent the probabilities that the sample belongs to class $c$ for:

where $\tau$ is a temperature parameter. Taking the teachers’ ensembled soft labels as $\sigma(\hat{z}^{c},\tau)$ and the student soft label as $\sigma(\tilde{z}^{c},\tau)$ , conventional knowledge distillation employs the Kullback-Leibler divergence to update the student model:

Hinton et al. [33] has shown that minimizing Eq. 2 with a high $\tau$ is equivalent to minimizing the $\ell_{2}$ error between the logits of teacher and student, thereby relating cross-entropy minimization to matching logits.

Let the global data distribution $p_{0}(x,y)$ of image $x$ and label $y$ be the target, while the local private data distribution can be indicated as $p_{k}(x,y)$. Due to the imbalance of data distribution among locals, we come to the bias ratio of local prediction :

where we assume $p_{k}(x|y)\thickapprox p_{0}(x|y)$ as the local difference on conditional probability distribution $p_{k}(x|y)$ is minor.

To consider this aspect, during the ensemble, we introduce an importance weight $\omega$ for each local node to reflect the distribution of local private data that its model was initially trained with:

where the importance weight $\omega^{c}$ is class-specific, which means the number of samples labeled as class $c$ for training the model of local node $k$: $N_{k}^{c}=\sum_{i=1}^{|\mathcal{D}_{k}|}({y}_{k}^{i}(c)=1)$. It reflects the distribution of local private data corresponding to each particular class $c$. Without loss of generality, we denote the local model output as $z_{k}=f(x_{0},\theta_{k})$ and central model output as $\tilde{z}=f(x_{0},\theta_{\mathrm{s}})$. Specifically, for the classification task, we have $z_{k}=[z_{k}^{1},...,z_{k}^{C}]$ and $\tilde{z}=[\tilde{z}^{1},...,\tilde{z}^{C}]$ for local model and central model respectively. Following the aforementioned $\ell_{2}$ observation from [33], we consider the case of $\tau\rightarrow\infty$, hence expressing the logit loss as:

The above-mentioned distillation essentially captures the divergence between the final output from the teacher and student models. However, it provides little insight into the underlying structure knowledge or reasoning of the teacher models, which can be complementary and important to the final output, especially in the FL scenario with its high degree of heterogeneity in local data sources. Although intuitive, it challenges transferring structural and more comprehensive knowledge. Structural knowledge, such as intermediate feature representations, suffers from a high bandwidth burden and, in most cases, relies on an identical network architecture among central (student) and all the local (teacher) models. Therefore we turn to more concise attention interpretations to transfer knowledge in a more efficient (as opposed to full feature tensors) and effective (as opposed to output vectors only) way without risking privacy leakage or posting additional communication/architecture requirements.

Specifically, we propose a bounded constraint for attention ensemble distillation to achieve consensus while maintaining the local node’s’ inherent diversity. Let $\bm{A}_{k}\in\mathbb{R}^{HW}$ be the attention map produced by the $k$-th local model, where $H$ and $W$ represent the 2D size of the attention map. Given the set of local attention maps $\mathcal{A}=\{\bm{A}_{k}|k=1,...,K\}$, we take the spatial-wise minimum among them as the attention consensus $\bm{I}$, and take the spatial-wise maximum to represent the attention diversity $\bm{U}$ among $\mathcal{A}$:

where $h=1,...,H$ and $w=1,...,W$. $\bm{I}$ denotes a consensus on the high-response region, among all the local attention maps, that have a high probability to be the real attention. While $\bm{U}$ considers all the high-response regions among the local attention maps, it also preserves diversity of “expertise” among the local models.

For simplicity, we denote $\tilde{\bm{A}}$ as the attention map generated by the central model. Considering the attention consensus $\bm{I}$ as a lower bound, we enforce the response in $\tilde{\bm{A}}$ to explicitly activate at the region of consensus achieved by all locals:

$T(\cdot)$ is a soft-masking operation based on sigmoid [51]:

Considering all the high-response regions among locals $\bm{U}$ as the upper bound, we enforce the response of $\tilde{\bm{A}}$ to be explicitly lower than that of $\bm{U}$:

The intuition here is that we seek each high-response pixel in $\tilde{\bm{A}}$ to have support from at least one local model to consider model diversity.

Compared to the naive distillation that enforces precisely the same attention strength as one aggregated attention map of diverse local nodes, our designed attention bound constraint is a relaxed version, i.e., tolerating incorrect/biased local attention maps. Its high robustness to outliers enables it to handle the inherent heterogeneity among locals more efficiently.

For the classification task, the heterogeneity of FL mainly comes from its inability to cope with the more general scenario of local nodes not sharing the same target classes. We note that Eq. 4 can deal with this issue efficiently.

We employ top-down attention generated with Grad-CAM [22] to provide location cues for class activation reasoning. Feeding an image to the model obtains a raw score $z^{c}$ (before the activation layer) for each class $c$. The gradient of score $z^{c}$ is computed with respect to the feature maps in a convolutional layer $[\bm{F}_{1},\bm{F}_{2},...,\bm{F}_{J}]$, where $\bm{F}_{j}\in\mathbb{R}^{H\times W}$, $J$ is the channel size, $H$ and $W$ indicates size of a 2D feature. These gradients can be globally averaged to obtain the neuron importance $\beta_{j}^{c}$ corresponding to $\bm{F}_{j}$:

All $\bm{F}_{j}$ weighted by $\beta_{j}^{c}$ are combined and activated with ReLU to get the a class-specific attention map $\bm{A}^{c}\in\mathbb{R}^{H\times W}$:

We then normalize the attention maps to have all values lie between 0 and 1: $A_{hw}^{c}=\frac{A_{hw}^{c}}{\max_{hw}A_{hw}^{c}}$. When employing the attention-bound loss, Eq.7 and Eq.9, the class-specific attention maps are taken independently, and thus the overall loss function for classification can be written as:

Segmentation can be seen as pixelwise/voxelwise classification task while differing from the above mainly in two aspects: 1) the model’s prediction $\bm{z}^{c}\in\mathbb{R}^{HW}$ is the same shape as its input (denoted as 2D here for simplicity); 2) we directly employ $\bm{z}^{c}$ for the attention $\bm{A}^{c}=\sigma(\bm{z}^{c},\tau)$ relating to the activation bound constraint.

For image reconstruction tasks, we rewrite the weighted ensemble as:

where $|\mathcal{D}_{k}|$ denotes the number of samples used to train the model at local node $k$, and model output ${\bm{z}}$ is with 2D image size. We employ a non-local self-attention module [23] to capture spatial-wise dependencies of 2D features. For one batch, given the feature maps $\bm{F}$ with size $J\times H\times W$, we reshape $\bm{F}$ to $\bar{\bm{F}}$ (with size $J\times HW$) and then calculate the spatial-wise similarity $\bm{S}\in\mathbb{R}^{HW\times HW}$ via dot product (matrix multiplication): $\bm{S}=\bar{\bm{F}}^{\text{T}}\cdot\bar{\bm{F}}$. Then $\bm{S}$ is normalized into spatial-wise attention $\bm{A}$ using softmax along the first dimension:

The normalized similarity is used to enhance the features $\bm{F}$:

where $\text{Reshape}(\cdot)$ is to reshape the size of $J\times HW$ to $J\times H\times W$. The overall loss function for image reconstruction can be written as:

The overall process is explained in Algorithm 1.

We evaluate our method on a multi-label classification task with standard chest-x-ray datasets: NIH chestX-ray14 (NIH CXR14) [56] and CheXpert [57] as locally held private data. NIH CXR14 consists of 112,120 frontal-view x-ray images scanned from 32,717 patients labeled with 14 diseases. CheXpert contains 224,316 chest radiographs from 65,240 patients labeled with 14 common chest radiographs, including both frontal-view and lateral-view. For public data, we use 26,684 x-ray images in the RSNA Pneumonia Detection Challenge public data [58] without using their labels.

We use NIH CXR14 and CheXpert as domains where private data comes from. For samples with multiple positive labels, we randomly choose one and split the dataset across locals using the Dirichlet distribution as in most FL works [4], in which the value of $\alpha$ controls the degree of non-IID-ness. A smaller $\alpha$ indicates higher non-IID-ness. For the total of $K$ local nodes, each dataset is distributed to $K_{d}=K/2$ local nodes under the cross-domain scenario. Following the validation strategy in [59], for both datasets we randomly sample a fraction (10%) of the training data to form the validation set. For training, we use ResNet-34 with a batch size of 32 and the same data augmentation methods as in [59]. Each local model is trained individually with SGD and CosineAnnealing [60] and a decreasing learning rate from 1e-3 to 1e-6 across 15 epochs. We use SGD and CosineAnnealing for distillation and a decreasing learning rate from 1e-2 to 1e-3 across 20 epochs.

We first study the cases when local samples are from one dataset. Table 2 shows results w.r.t. when local data are within the domain with a varying number of locals $K$ and non-IID-ness $\alpha$. One can note that our method outperforms centralized training with all local data on CheXpert when $K$=3 and $\alpha$=1. Table  3 shows training results with varying public dataset sizes. The results suggest that, although unlabeled, a more extensive public dataset improves performance. To compare with the SOTA FL methods, Table 4 shows whether the method in comparison transfers parameters or employs distillation and analyzes each privacy guarantee and synchronization requirements. We can see that our method outperforms the counterparts with the best utility-privacy trade-off.

With both datasets as private data, we conduct cross-domain, cross-site evaluations with FedAD. We distribute the training datasets to $K$=6 local nodes ($K_{d}$=3 for each dataset). Table 1 shows the results of FedAD on cross-domain datasets with $\alpha$=1. On CheXpert, both single domain and cross-domain distillation achieve better performance than centralized training: cross-domain learning with FedAD obtains the best mAUC of 82.65%, slightly better than FedAD training with only data from CheXpert (82.62%). On NIH CXR14, while training with all data centrally yields the best result (79.73%), cross-domain trained FedAD still obtains better performance (75.54%) compared to FedAD trained with single domain data (75.17%) and compared to the average AUC of local nodes (75.08%). Table 1 also reports results on the six overlapping classes in each test set. It can be seen that FedAD trained with cross-domain data obtains superior performance compared to the two FedAD models trained with single domain data on both test datasets. We note that the FedAD model trained with cross-domain data can classify the 14 classes in total (12 classes from CheXpert and the 8 classes from NIH CXR14), whereas other models can only classify 8 or 12 classes.

The BraTS 2018 dataset [61] contains multi-parametric pre-operative MRI scans of 285 subjects with brain tumors. Each subject was scanned under the T1-weighted, T1-weighted with contrast enhancement, T2-weighted, and T2 fluid-attenuated inversion recovery (T2-FLAIR) modalities.

Following the same protocol as in [13], we use 242 subjects for the training set and 43 subjects for held-out test set. According to the originating institution, the training set is stratified into three federated local clients. We use the unlabeled validation set of the BraTS 2020 dataset [62] as the public data comprising 125 subjects, independent of either private dataset. We use the same model structure as [13, 63] but with half its channel numbers. We only use its segmentation branch (no reconstruction branch). The training strategy is the same as [63]. The local training takes 20,000 iterations, and local-to-central distillation takes 5,000 iterations.

Table 5 compares the segmentation performance on BraTS with the SOTA parameter based method [13] and the SOTA distillation based method [17]. Note we report both naive (non-private) and noisy (less-private) version of [13]. The comparison shows that our method achieves the best utility-privacy trade-off. We can also observe that when compared with [17], our method achieves better results with much more efficiency, i.e., lower communication bandwidth and no local synchronization requirement at the same time.

Following the prior art [64], we use fastMRI [65], IXI [66], BraTS[67] as private data distributed across local nodes and evaluate the corresponding test sets (same data split as [64]). Guo et al. [64] reports results with four brain MRI datasets, of which HPKS [68] is privately held and not available for use at the moment; so we experiment with the remaining publicly available sets as local data: fastMRI, IXI, and BraTS. Besides, we use OASIS-3 [69] as a public dataset.

Following [64], we subsample the given k-space by multiplying with a mask, where the acceleration factor (AF) is set as 4. The 2D MRI images are preprocessed with zero padding and then cropped to the size of $256\times 256$. We utilize the same U-Net [70] style encoder-decoder architecture for the reconstruction networks as the one provided in FL-MRCM [64]. A minor difference in the architecture with [64] comes from the additional residual non-local block 15 deployed on the bottom features of U-Net before forwarding into a sequence of the up-sampling layer. The size of these features $J\times H\times W$ are $512\times 16\times 16$. For local training, the network is trained with an Adam optimizer using a constant learning rate of $1e^{-4}$ with 20 epochs. For distillation, the central model is trained with an RMSprop optimizer using a constant learning rate of $1e^{-4}$ with five epochs in one communication round.

We first perform ablation studies on the impact of the output ensemble scheme, the proposed attention distillation bound, and the modality of unlabeled public data. Here we use the T1-weighted images from BraTS (B), fastMRI (F), and IXI (I) as private datasets and unlabeled T1-weighted images from OASIS-3 as public data; and we report results on the aggregated T1-weighted test images from B, F, and I. From Table 6 we can see the superiority of the importance-weighted ensemble beyond the average ensemble typically used in previous FL works [20, 12]. Our proposed attention to upper/lower bound constraint further improves the reconstruction performance with higher SSIM and PSNR. In addition, Table 6 shows the comparison of using T1-weighted and T2-weighted images of OASIS-3 as public data. The results demonstrate the robustness of our method to public data from a different domain (locally held data and data used for distillation are all from other datasets) and even different modalities (all three local nodes have T1-weighted images for training while public data is T2-weighted).

Second, we compare the performance with the prior arts [64] when taking T1/T2-weighted images from B, F, and I as local data and unlabeled T1/T2-weighted images of OASIS-3 as public data. The results are reported on the corresponding test sets of local data, respectively. From Table 7 we can see that, when compared with the prior art [64], our method achieves very competitive reconstruction performance in terms of SSIM and PSNR with higher communication efficiency while at the same time maintaining the local data privacy by not sharing local model parameters/gradients or any product inferred from local private local data. The counterpart [64] not only iteratively shares local model parameters but also shares the features inferred on each local private data. Besides the superior guarantees w.r.t. data privacy, our method demonstrates higher communication efficiency through lower bandwidth and higher flexibility (offline communication without any synchronization requirements on the local model). Notably, when testing on T2-weighted F and I, we achieved better performance than centralized training (collecting local data together for training), e.g., on T2-weighted F, we achieved 0.9374/32.76 SSIM/PSNR over the 0.9002/30.47 SSIM/PSNR of centralized training. The reason is that we utilize additional unlabeled, non-sensitive, cross-domain public data, which, we assume, are easily acquired in real-world clinical scenarios. Comparisons of qualitative results are shown in Figure 4.

We leverage cross-domain test sets (different domains with the local data) to evaluate the generalizability. Table 8 compares our method with the SOTA FL methods. It shows that our privacy-preserving method owns comparable generalizability with the prior arts [2, 64, 20], which share iterative local model parameters and therefore risk privacy leakage.