DaFKD: Domain-aware Federated Knowledge Distillation

1. Introduction

Federated learning (FL) has emerged as a prominent distributed machine learning framework to train a global model via the collaboration among users without sharing their original dataset [17,22,27]. Due to the benefits of preserving privacy and economic communication efficiency, FL has been widely adopted in various applications such as medical image processing [8,19,32] and recommendation [2,26].

The classic FL paradigm, FedAvg [22], iteratively optimize the global model by aggregating the parameters of local models trained from data resides on a number of remote devices or servers. However, these methods usually suffer from serious model performance degradation when the data is not independently and identically distributed (Non-IID) across clients, which is a common issue in FL scenarios. This is mainly because the model parameters on different clients are optimized towards diverse directions [14], leading to the overlarge variance of the aggregated model.

To tackle this challenge, federated distillation (FD) [13] proposes to distill the knowledge of multiple local models into the global model by aggregating only the output soft predictions, which recently attracts increasing attention. For instance, Tao et al. [18] leveraged the public dataset as the distillation data samples to obtain the soft predictions from multiple local models and then updated the global model with the average of these soft predictions. Based on [18], Zhu et al. [40] and Zhang et al. [37] improved the distillation by replacing the public dataset with data generated by the generative model. Although these methods achieve significant improvement over existing parameter-averaging methods, they still do not take the model diversity into account and may still limit the model performance. More specifically, only computing the average of soft predictions will inevitably bring errors when some local models make wrong predictions for the distillation sample.

To break the limitations of existing federated distillation methods, we in this paper propose a novel federated distillation method dubbed DaFKD that can discern the importance of each model to the given distillation sample, and thus is able to reduce the impact of wrong soft predictions. More specifically, we consider that the local data in each client constitutes as a specific domain and employ a domain discriminator for each client to identify the correlation factor between the sample and the domain. For a given distillation sample, we endow the local model with high importance when its correlation factor is significant and vice versa. The principle behind this is the fact that a model tends to make the correct prediction when the sample is contained in the domain for training the model. Furthermore, to facilitate the training of domain discriminator, we propose sharing its partial parameters with the target classification model. Through extensive experiments over various datasets (MNIST, EMNIST, FASHION MNIST, SVHN) and different settings (various models and data distributions), we show that the proposed method significantly improves the model accuracy as compared to state-of-the-art algorithms. The contributions of this paper are:

• We propose a new domain aware federated distillation method named DaFKD which endows the model with different importance according to the correlation between the distillation sample and the training domain.
• To adaptively discern the importance of multiple local models, we propose employing the domain discriminator for each client which identifies the correlation factors. To facilitate the training of the discriminator, we further propose sharing partial parameters between the discriminator and the target classification model.
• We establish the theories for the generalization bound of the proposed method. Different from existing methods, we theoretically show that DaFKD efficiently solves the Non-IID problem where the generalization bound of DaFKD does not increases with the growth of data heterogeneity.
• We conduct extensive experiments over various datasets and settings. The results demonstrate the effectiveness of the proposed method which improves the model accuracy by up to $6.02\%$ compared to state-of-the-art methods.

2. Related Work

Federated Learning with Parameters Aggregation The main problem incurred by Non-IID for the model parameters aggregation based methods is the huge diverse of the parameters of local models across clients [21,39], where the local models are optimized to different directions. To tackle this challenge, many prior works seek to to reduce the diverse across local model parameters for efficient model aggregation [1,36]. For instance, Tian et al. [16] proposed adding a regularization item in the local objective function such that the divergence of the local model is constrained by the global model. Sai et al. [14] proposed reducing the variance of local gradient to align the diverse local update. This paper applies the ensemble distillation over the obtained local models which is orthogonal to these methods.

Knowledge Distillation is to transfer the knowledge from one or more networks (teacher) to another (student) [12]. The key step in knowledge distillation is to align the soft prediction of the student model to that of the teacher model [5,25,30,33,38]. For example, some works leverage a proxy dataset to distill knowledge between networks [11,34]. Considering that the proxy dataset may not always exist, some recent works proposed distillation in a data-free manner including reconstructing samples used for training the teacher [20,23] or learning a generator [35]. By following this path forward, we in this paper particularly focus on the distillation in federated learning by tailoring the generative model and ensemble distillation techniques.

Federated Distillation Federated distillation is to distill the knowledge from multiple teacher models trained by different clients to the student model [3,9,28]. Lin et al. [18] first proposed leveraging the knowledge distillation in the server to transfer the knowledge from multiple local models to the global model based on an unlabeled proxy dataset. Chen et al. [4] proposed linearly aggregating multiple local models with weights generated by the Bayesian posterior to produce a series of combined models and then, distilling these models into one global model. Considering that these methods rely on an unlabeled auxiliary dataset in the server which may not exist in real-world settings, Zhu et al. [40] and Zhang et al. [37] proposed replacing the proxy dataset with data generated by the generative models and making ensemble federated distillation in a data-free way. However, most of these methods construct the ensemble knowledge by simply computing the average of soft predictions from multiple local models, which does not take model diversity into account and may limit the model performance. In this paper, we take the domain knowledge for training local models into account and endow the local models with different importance when ensembling these soft predictions.

3. Methodology

In this section, we specify the proposed method DaFKD. As illustrated in Figure 1, the key idea of the DaFKD is to leverage a domain discriminator to discern the importance of each local model to the given distillation sample such that the performance of ensemble distillation can be improved. Specifically, each client in DaFKD trains the local model with its private dataset and the domain discriminator with a global generator in an adversarial way. Then, the server aggregates local models from all participated clients and make ensemble distillation with the generator producing distillation samples and the discriminator producing importance. Besides, to facilitate the training of the discriminator, we further propose sharing partial parameters between the discriminator and the target classification model. The workflow of the proposed algorithm is shown in the Algorithm 1.

3.1. Problem Formulation

We aim to collaboratively train a global model for K total clients in FL. We consider each client k can only access to his local private dataset ${\mathbb{D}}_k:=\{x_i^n,y_i\}$, where $x_i$ is the i-th input data sample and $y_i\in \{1,2,\cdots ,C\}$ is the corresponding label of $x_i$ with C classes. We denote the number of data samples in dataset ${\mathbb{D}}_k$ by $D_k$. The global dataset is considered as the composition of all local datasets $\mathbb{D}=\{{\mathbb{D}}_1,{\mathbb{D}}_2,\cdots ,{\mathbb{D}}_K\}$, $D={\sum }_{k=1}^K{D}_k$. The objective of the FL learning system is to learn a global model w that minimizes the total empirical loss over the entire dataset $\mathbb{D}$:  

$$\begin{array}{cc}\hfill & \underset{w}{min}\mathcal{L}(w):=\sum _{k=1}^K\frac{D_k}{D}{\mathcal{L}}_k(w),\hfill \\ \hfill & \mathrm{where}{\mathcal{L}}_k(w)=\frac{1}{D_k}\sum _{i=1}^{D_k}{\mathcal{L}}_{CE}(w;x_i,y_i),\hfill \end{array}$$

(1)

where ${\mathcal{L}}_k(w)$ is the local loss in the k-th client and ${\mathcal{L}}_{CE}$ is the cross-entropy loss function that measures the difference between the prediction and the ground truth labels.

3.2. Domain Discriminator

To endow the model with suitable importance to the given distillation sample, an intuition is that the model has high probability in making the correct prediction when the sample is contained in the domain for training the model. As a consequence, quantifying the correlation between the domain and any given sample is necessary. Motivated by the adversarial training techniques [7] where the discriminator is trained to distinguish whether the generated data is sampled from the distribution of the target dataset, we in this paper proposes a domain discriminator which views the local dataset as the target and outputs the correlation between the distillation sample and the local domain.

More specifically, we assign a personalized discriminator ${\theta }_k^d$ for each client k and adopt a global generator ${\theta }^g$ shared by all clients. At each round t, each participated client k firstly pulls the generator ${\theta }^g$ from the server to produce pseudo dataset ${\widehat{\mathbb{D}}}_k$ with ${\widehat{D}}_k$ samples by sampling noise from the distribution $p_z$. Then, each client k labels the samples in local private dataset ${\mathbb{D}}_k$ positive and the samples in pseudo dataset ${\widehat{\mathbb{D}}}_k$ negative. Using these data samples with generated labels, each client k trains the domain discriminator ${\theta }_k^d$ with the following loss function:

$$\begin{array}{c}\hfill \begin{array}{cc}\hfill \underset{{\theta }_k^d}{min}{\mathcal{L}}_{adv}^k({\theta }_k^d)& =\frac{-1}{D_k+{\widehat{D}}_k}[\sum _{x_i\in {\mathbb{D}}_k}\log f({\theta }_k^d;x_i)\hfill \\ \hfill & +\sum _{x_i\in {\widehat{\mathbb{D}}}_k}\log (1-f({\theta }_k^d;x_i))],\hfill \end{array}\end{array}$$

(2)

where $f({\theta }_k^d;x_i)$ denotes the probability of $x_i$ being real data. Considering that there have been extensive works for training effective generators [37,40] which can be jointly used with our method, we in this paper simplify the training of generator and adopt the basic FedAvg [22] to train the generator, which still exhibits great effectiveness. Specifically, after obtaining the discriminator ${\theta }_k^d$, the client k in turn leverages the discriminator to train the generator, which maximizes the loss function 2:

$$\begin{array}{c}\hfill \underset{{\theta }_k^g}{max}{\mathcal{L}}_{adv}^k({\theta }_k^g)=-{\mathbb{E}}_{z\sim p_z(z)}\log (1-f({\theta }_k^d;g({\theta }_k^g;z))).\end{array}$$

(3)

It is worthwhile to note that the client can also train the generator ${\theta }_k^g$ and the discriminator ${\theta }_k^d$ in an alternative way like [7]. After obtaining the updated local generator ${\theta }_k^g$, the server receives them from all $K_t$ participated clients and aggregates them to get the new global generator:

$$\begin{array}{c}\hfill \begin{array}{c}\hfill {\theta }^g=\frac{1}{K_t}\sum _{k=1}^{K_t}{\theta }_k^g.\end{array}\end{array}$$

(4)

From a global perspective, the adversarial loss function can be formulated as:

$$\begin{array}{c}\hfill \begin{array}{cc}\hfill & \underset{{\theta }^g}{max}\underset{{\theta }_1^d,\cdots ,{\theta }_K^d}{min}{\mathcal{L}}_{adv}({\theta }_1^d,\cdots ,{\theta }_K^d)\hfill \\ \hfill & =-\frac{1}{K}\sum _{k=1}^K[{\mathbb{E}}_{x\sim p_k(x)}\log f({\theta }_k^d;x)\hfill \\ \hfill & +{\mathbb{E}}_{z\sim p_z(z)}\log (1-f({\theta }_k^d;g({\theta }^g;z)))].\hfill \end{array}\end{array}$$

(5)

Discussion. One concern for this method may be that the local generator will lead to privacy leakage when it is uploaded from client to the server. In fact, there have been many prior works solving this problem. For example, to protect privacy, Zhu et al. [40] proposed generating feature maps instead of the original data and Xin et al. [31] proposed exploiting the differential privacy. Similarly, the privacy concern of the global generator can also be addressed by only allowing the generator to output intermediate features as specified in Zhu et al. [40]. The main method proposed in this paper supports various generators and thus can definitely achieve privacy protection as combined with these privacy-protecting methods together, of which more details can be found in Appendix A.    

Algorithm 1: Workflow of DaFKD Algorithm

3.3. Domain-aware Federated Distillation

To obtain the classification model, in each round t, each participated client k firstly locally trains the model $w_t^k$ and sends it as well as the domain discriminator ${\theta }_k^d$ to the server. After receiving the two models, the server aggregates multiple local models $w_t^k$ by computing their average as:

$$\begin{array}{c}\hfill \begin{array}{c}\hfill {\widehat{w}}_{t+1}=\frac{1}{K_t}\sum _{k=1}^{K_t}{w}_t^k.\end{array}\end{array}$$

(6)

Then, the server uses the global generator ${\theta }^g$ to generate the pseudo dataset ${\widehat{\mathbb{D}}}^g$ as the distillation data. After that, for every distillation sample $x_i\in {\widehat{\mathbb{D}}}^g$, the server computes the importance of each local model $w_t^k$ with the domain discriminator ${\theta }_k^g$ as ${\alpha }_{k,i}=f({\theta }_k^d;x_i)$ and then normalizes it into the probability:

$$\begin{array}{c}\hfill {\widehat{\alpha }}_{k,i}=\frac{f({\theta }_k^d;x_i)}{{\sum }_{k=1}^{K_t}f({\theta }_k^d;x_i)},\end{array}$$

(7)

which guarantees that ${\sum }_{k=1}^{K_t}{\widehat{\alpha }}_{k,i}=1$. Finally, the server inputs the pseudo sample $x_i$ into each local model $w_t^k$ and the average model ${\widehat{w}}_{t+1}$ to obtain the soft predictions $s(w_t^k;x_i)$ and $s({\widehat{w}}_{t+1};x_i)$, and applies ensemble knowledge distillation with the importance ${\widehat{\alpha }}_{k,i}$ to obtain the global model $w_{t+1}$:

$$\begin{array}{cc}\hfill w_{t+1}& ={argmin}_{{\widehat{w}}_{t+1}}{\mathcal{L}}_{KD}({\widehat{w}}_{t+1})\hfill \\ \hfill & =\frac{1}{{\widehat{D}}^g}\sum _{x_i\in {\widehat{\mathbb{D}}}^g}KL\left(\sum _{k=1}^{K_t}{\widehat{\alpha }}_{k,i}·s(w_t^k;x_i),s({\widehat{w}}_{t+1};x_i)\right),\hfill \end{array}$$

(8)

where $KL(·)$ is to compute the Kullback-Leibler divergence (KL-divergence).

Discussion. In our method, the domain discriminator is not necessary to be uploaded to the server. Since the domain discriminator is mainly used to output the importance of each model to the distillation sample, it can definitely be used in each client when the distillation sample is generated locally. Naturally, considering that each client has access to the global generator which is used to train the domain discriminator, the shared distillation dataset can be generated in each local client when the random seed of noise is consistent. Besides, to further protect the privacy of importance, each client can directly upload the soft predictions weighted by importance to the server, which we defer the details to the Appendix B.

3.4. Partial Parameters Sharing

Considering that the domain discriminator is trained over the local dataset, its performance may be deteriorated when the size of local dataset is small. Motivated by the idea in the multi-task learning where sharing the encoder between different tasks can promote each other, we propose sharing partial parameters between the discriminator ${\theta }_k^d$ and the classification model $w^k$ to solve this problem. The intuition behind this is that both the discriminator and the classification model have to distinguish the sample from the extracted features. Besides, another benefit of sharing layers is that the communication cost can be reduced when the discriminator is also uploaded to the server.

To this end, we propose sharing the front model layers between the two models which are used to extract features, as illustrated in Figure 2. Specifically, by denoting the c-layer shared extractor by ${\tilde{w}}^k=[{\tilde{w}}^{k,1},\dots ,{\tilde{w}}^{k,c}]$ where ${\tilde{w}}^{k,i}$ is the i-th layer of the shared extractor, we can re-write the $n_d$-layer discriminator by ${\theta }_k^d=[{\tilde{w}}^k,{\theta }_k^{d,c+1},\dots ,{\theta }_k^{d,n_d}]$, and the n-layer classification model $w^k$ by $w^k=[{\tilde{w}}^k,w^{k,c+1},\dots ,w^{k,n}]$ with $w^{k,i}$ as the i-th layer of the classification model $w^{k,i}$. With the shared layers, the discriminator and the classification model are trained in a joint way:

$$\begin{array}{c}\hfill \underset{{\theta }_k^d,w^k}{min}{\mathcal{L}}_J^k({\theta }_k^d,w^k)={\mathcal{L}}_{adv}^k({\theta }_k^d)+{\mathcal{L}}_k(w^k),\end{array}$$

(9)

where ${\mathcal{L}}_{adv}^k({\theta }_k^d)$ refers to (2) and ${\mathcal{L}}_k(w^k)$ is defined in (1).

3.5. Theoretical Analysis

In this section, we prove that our method can solve the Non-IID problem of FD. To verify this, we first analyze the distribution learned by the generator and discriminator.

Theorem 1.

Denote the data distribution of each client k by $p_k(x)$, the data distribution of all clients by $p(x)$, and the pseudo data distribution of the generator by $p_g(x)$. If the Algorithm 1 trains the discriminator ${\theta }_k^d$ and the global generator ${\theta }^g$ to the optima for the loss function (5), then the pseudo data distribution of the generator is $p_g^{*}(x)=p(x)$, and the discriminator outputs $f^{*}({\theta }_k^d;x)=\frac{p_k(x)}{p_k(x)+p(x)}$ for each client $k=1\dots K$.

The detailed proof is deferred to Appendix C. Theorem 1 exhibits that the global generator still learns the global data distribution even though there are multiple discriminators specialized for different clients. Besides, the discriminator can distinguish the samples either from the global distribution or the local distribution, and thus can generate efficient correlation factors with the collaboration of the generator that produces global distribution. For simplicity of notation, we denote the local model $f(w_k;x)$ trained over the local dataset ${\mathbb{D}}_k$ by $h_{{\widehat{p}}_k}(x)$ and $f(w;x)$ trained over global dataset $\mathbb{D}$ by $h_{\widehat{p}}(x)$. Besides, without losing generality, we consider $D_1=\cdots =D_K=m$. Then, we can further derive the generalization bound of the proposed method.

Theorem 2.

Denote the empirical distribution of activation from each client k by ${\widehat{p}}_k$ and the empirical distribution of global dataset by $\widehat{p}=\frac{1}{K}{\sum }_{k=1}^K{\widehat{p}}_k$. Then, given the constants $0<\delta \le 1$ and $\sigma >0$, with the probability at least $1-\delta$, the expected generalization error ${\mathcal{L}}_p({\sum }_{k=1}^K{\widehat{\alpha }}_k(x)h_{{\widehat{p}}_k})$ of domain-aware ensemble model is:

$$\begin{array}{c}\hfill \begin{array}{cc}\hfill & {\mathcal{L}}_p(\sum _{k=1}^K{\widehat{\alpha }}_k(x)h_{{\widehat{p}}_k})\hfill \\ \hfill & \le (K+1){\mathcal{L}}_{\widehat{p}}(h_{\widehat{p}})+(K+1)\sqrt{\frac{{\sigma }^2\log \frac{2K}{\delta }}{2m}}.\hfill \end{array}\end{array}$$

(10)

The proof can be found in the Appendix D. For ease of comparison, we here present the bounds of state-of-the-art FD methods. Specifically, the bound of FEDFUSION [18] is

$$\begin{array}{cc}\hfill & {\mathcal{L}}_p(\sum _{k=1}^K{\widehat{\alpha }}_k(x)h_{{\widehat{p}}_k})\hfill \\ \hfill & \le {\mathcal{L}}_{\widehat{p}}(h_{\widehat{p}})+\sqrt{\frac{{\sigma }^2\log \frac{2K}{\delta }}{2m}}+\frac{1}{2K}\sum _{i=1}^K{d}_{\mathcal{H}\Delta \mathcal{H}}(p_k,p)+{\lambda }_k,\hfill \end{array}$$

(11)

where ${\lambda }_k={min}_{h\in \mathcal{H}}{\mathcal{L}}_{p_k}(h)+{\mathcal{L}}_p(h)$ and FEDGEN [40] is

$$\begin{array}{cc}\hfill & {\mathcal{L}}_p(\sum _{k=1}^K{\widehat{\alpha }}_k(x)h_{{\widehat{p}}_k})\le {\mathcal{L}}_{\widehat{p}}(h_{\widehat{p}})\hfill \\ \hfill & +\sqrt{\frac{4}{m^{\prime }}(d\log \frac{2em}{d}+\log \frac{4K}{\delta })}+\frac{1}{K}\sum _{i=1}^K{d}_{\mathcal{H}\Delta \mathcal{H}}(p_k^{\prime },p)+{\lambda }_k^{\prime }.\hfill \end{array}$$

(12)

The main differences between DaFKD and baselines are the items $\frac{1}{K}{\sum }_{i=1}^K{d}_{\mathcal{H}\Delta \mathcal{H}}(p_k^{\prime },p)$ and ${\lambda }_k$ which measure the distance between the local data distribution and the global distribution and are incurred by the Non-IID. DaFKD does not include the two items, which indicates that the Non-IID problem is efficiently solved.

Figure 3. Visualized performance w.r.t data heterogeneity.

Figure 3. Visualized performance w.r.t data heterogeneity.

4. Experiments

In this section, we compare the performance of our proposed approach with related works.

4.1. Setup

Baselines: In addition to FEDAVG [22], FEDPROX regularizes the local model training with a proximal term in the model objective [16]. FEDDFUSION is a data-based knowledge distillation method, which applies unlabeled training samples as the proxy dataset [18]. FEDGEN is a data-free knowledge distillation approach that each client can directly regulate the local model updating using the generated unlabeled samples in server [40]. FEDFTG learns a generator to ensemble knowledge of local models in a data-free manner and fine-tunes the global model in server instead of broadcasting the aggregated model back to each client directly [37].

Dataset: We conduct experiments on four image datasets with heterogeneous dataset partition: MNIST [15], EMNIST [6], FASHION MNIST [29] and SVHN [24]. Among them, MNIST, EMNIST and SVHN dataset is for digit and character image classifications, and FASHION MNIST is a fashion-product dataset which is used to learn a multi-class classification task.

Configurations: Unless otherwise mentioned, we set the number of local training epoch E = 20, communication round T = 60, the client number K = 20 with an active ratio r = 0.4. For local training, the batch size is 32 and the weight decay is $1e-3$. The learning rate is $0.01$ for distillation and $0.001$ for training the classifier, generator, and discriminator. Like FEDGEN [40], we use the Dirichlet distribution Dir($\alpha$) on labels to simulate the data heterogeneity. We apply all the training samples and distribute them to user models, and we use all the testing samples for the performance evaluation. For the classifier in all methods, we employ ResNet11 [10] as the basic backbone. For the generator in FEDGEN, FEDFTG and DaFKD, we apply the network composed of two embedding layers(for one-hot label vector and noise vector respectively) and the fully-connected layer with LeakyReLU and BatchNorm layers. For the multi-task learning structure in our DaFKD approach, we treat all previous layers before the last fully-connected layer as share layers, and we use two different fully-connected layers to get outputs as the classifier result and discriminator result.

Figure 4. Fitted learning curve of four image datasets in 60 communication rounds ($\alpha$ = 0.1).

Figure 4. Fitted learning curve of four image datasets in 60 communication rounds ($\alpha$ = 0.1).

4.2. Performance Overview

Test accuracy. Table 1 shows the performance of our DaFKD and other baseline methods on four image datasets. We carry out experiments against different levels of data heterogeneity on each dataset, and DaFKD achieves the best performance in most cases. Among all mentioned approaches, FEDDFUSION, FEDGEN, FEDFTG and DaFKD apply the knowledge distillation with extra data to improve the model training. As shown in Table 1, these KD based methods are notably excellent, and outperform FEDAVG and FEDPROX in most scenarios. Besides, DaFKD is the only KD-algorithm that is robust against different datasets while consistently performs well, especially surpassing the second by 2-6% with different data heterogeneity levels on SVHN dataset. These results verify our motivation that the domain discriminator can identify the correlation and thus solve the heterogeneity problem.

Data heterogeneity and Client participant. Figure 3 displays the test accuracy with different levels of data heterogeneity on four image datasets. As vividly shown in this figure, all methods achieve an improvement in test accuracy with the decreasing degree of data heterogeneity. Most notably, DaFKD gains a significant improvement in test accuracy in data heterogeneity $\alpha$ = 0.1 and steadily outperforms all methods. Figure 5 provides the test accuracy under different ratios r between active clients and total clients. In this figure, DaFKD performs best with different ratios r and the higher accuracy is achieved as applying more active clients in each communication round.

Communication rounds. Table 2 shows evaluation of DaFKD and other baseline methods in terms of the communication rounds to reach the target test accuracy. Here we highlight the best and second-best results in bold. DaFKD reaches most best and the second-best evaluation results on all datasets. What is more, although FEDAVG reach the target accuracy with fewer communication rounds on EMNIST dataset, DaFKD finally can surpass it by significant 5% after all communication rounds. Meanwhile, Figure 4 displays the fitted learning curve of four image datasets in 60 communication rounds, where we adapt the test accuracy per five communication rounds to fit the learning curve. In the figure, each method rapidly increases at the beginning and slows down as training goes. DaFKD always keeps a leading tendency in the test accuracy.

Parameter sensitivity analysis. To figure out whether DaFKD is sensitive to some specific parameters, we select local training epoch E, sample batchsize B and dimension d of noise on SVHN with $\alpha$ = 0.1 to carry out experiments. Figure 6 shows the performance of DaFKD under different configurations. DaFKD achieves the better result when we increase the local training epochs at the beginning. However, DaFKD has a comparable performance when the E is set to 20 and 40. In order to balance communication expense and test accuracy, we give priority to the local training epoch E = 20 here. In addition, DaFKD achieves a similar performance with different sample size B and dimension d of noise. This indicates that DaFKD is only sensitive to few parameters and still robust to most parameters in a large range.

The impact of sharing parameters. We share all front layers (ResNet11) before the last fully-connected layer as share layers between the generator and the discriminator in DaFKD. Table 3 shows the results. As can be seen, parameter sharing facilitates the training of DaFKD in two folds: 1) it enables higher accuracy in most cases; 2) it accelerates convergence where it consistently costs fewer rounds as reaching some given accuracy. Furthermore, parameter sharing can save communication costs by transmitting fewer parameters in each round. Besides, from Table 3.(a), we can find that DaFKD with correlations performs better in all cases (sharing or not). From the last two columns of Table 3.(b), we can find that parameter sharing can accelerate convergence faster than correlations.

5. Conclusion

In this paper, we seek to tackle the data heterogeneity challenge in the federated knowledge distillation. We propose a novel method dubbed domain-aware federated knowledge distillation, namely, DaFKD, which imposes an importance on each local model for some given distillation sample. To quantify the importance, we leverage a domain discriminator to compute the correlation between the distillation sample and the domain for training the local model. Furthermore, to facilitate the training of the domain discriminator, we propose sharing its partial parameters to the classification model. Extensive experiments conducted on various datasets and settings show that our method achieves significant improvement for the model accuracy.