1. Introduction

2. Related Work

3. Methodology

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.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

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).

5. Conclusion

References

1. Acar, D.A.E.; Zhao, Y.; Navarro, R.M.; Mattina, M.; Whatmough, P.N.; Saligrama, V. Federated Learning Based on Dynamic Regularization. 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021.
2. Ammad-ud-din, M.; Ivannikova, E.; Khan, S.A.; Oyomno, W.; Fu, Q.; Tan, K.E.; Flanagan, A. Federated Collaborative Filtering for Privacy-Preserving Personalized Recommendation System. CoRR, abs/1901.09888, 2019. [CrossRef]
3. Bistritz, I.; Mann, A.; Bambos, N. Distributed Distillation for On-Device Learning. Proceedings of Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems, NeurIPS, 2020.
4. Chen, H.; Chao, W. FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning. 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021.
5. Cho, J.H.; Hariharan, B. On the efficacy of knowledge distillation. Proceedings of 2019 IEEE/CVF International Conference on Computer Vision, ICCV, 2019, pp. 4794–4802.
6. Cohen, G.; Afshar, S.; Tapson, J.; Van Schaik, A. EMNIST: Extending MNIST to handwritten letters. 2017 international joint conference on neural networks (IJCNN). IEEE, 2017, pp. 2921–2926.
7. Goodfellow, I.J.; Pouget-Abadie, J.; Mirza, M.; Xu, B.; Warde-Farley, D.; Ozair, S.; Courville, A.C.; Bengio, Y. Generative Adversarial Nets. Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada, pp. 2672–2680.
8. Guo, P.; Wang, P.; Zhou, J.; Jiang, S.; Patel, V.M. Multi-Institutional Collaborations for Improving Deep Learning-Based Magnetic Resonance Image Reconstruction Using Federated Learning. IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, pp. 2423–2432.
9. Guo, Q.; Wang, X.; Wu, Y.; Yu, Z.; Liang, D.; Hu, X.; Luo, P. Online knowledge distillation via collaborative learning. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR, 2020, pp. 11020–11029.
10. He, K.; Zhang, X.; Ren, S.; Sun, J. Deep residual learning for image recognition. Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 770–778.
11. He, R.; Sun, S.; Yang, J.; Bai, S.; Qi, X. Knowledge Distillation as Efficient Pre-training: Faster Convergence, Higher Data-efficiency, and Better Transferability. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 9151–9161.
12. Hinton, G.; Vinyals, O.; Dean, J. Distilling the knowledge in a neural network. arXiv preprint 2015, arXiv:1503.02531. [Google Scholar]
13. Jeong, E.; Oh, S.; Kim, H.; Park, J.; Bennis, M.; Kim, S. Communication-Efficient On-Device Machine Learning: Federated Distillation and Augmentation under Non-IID Private Data. CoRR, abs/1811.11479, 2018. [CrossRef]
14. Karimireddy, S.P.; Kale, S.; Mohri, M.; Reddi, S.J.; Stich, S.U.; Suresh, A.T. SCAFFOLD: Stochastic Controlled Averaging for Federated Learning. Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, pp. 5132–5143.
15. LeCun, Y.; Cortes, C.; Burges, C. MNIST handwritten digit database, 2010.
16. Li, T.; Sahu, A.K.; Zaheer, M.; Sanjabi, M.; Talwalkar, A.; Smith, V. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems 2020, 2, 429–450. [Google Scholar]
17. Li, X.; Xu, Y.; Song, S.; Li, B.; Li, Y.; Shao, Y.; Zhan, D. Federated Learning with Position-Aware Neurons. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 10072–10081.
18. Lin, T.; Kong, L.; Stich, S.U.; Jaggi, M. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems 2020, 33, 2351–2363. [Google Scholar]
19. Liu, Q.; Chen, C.; Qin, J.; Dou, Q.; Heng, P. FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space. IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, pp. 1013–1023.
20. Lopes, R.G.; Fenu, S.; Starner, T. Data-Free Knowledge Distillation for Deep Neural Networks. CoRR, abs/1710.07535, 2017. [CrossRef]
21. Malinovskiy, G.; Kovalev, D.; Gasanov, E.; Condat, L.; Richtárik, P. From Local SGD to Local Fixed-Point Methods for Federated Learning. Proceedings of the 37th International Conference on Machine Learning, ICML, 13-18 July, Virtual Event, 2020, pp. 6692–6701.
22. McMahan, B.; Moore, E.; Ramage, D.; Hampson, S.; y Arcas, B.A. Communication-efficient learning of deep networks from decentralized data. Artificial intelligence and statistics. PMLR, 2017, pp. 1273–1282.
23. Micaelli, P.; Storkey, A.J. Zero-shot Knowledge Transfer via Adversarial Belief Matching. Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 9547–9557.
24. Netzer, Y.; Wang, T.; Coates, A.; Bissacco, A.; Wu, B.; Ng, A.Y. Reading digits in natural images with unsupervised feature learning 2011.
25. Phuong, M.; Lampert, C.H. Distillation-based training for multi-exit architectures. Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019, pp. 1355–1364.
26. Ramaswamy, S.; Mathews, R.; Rao, K.; Beaufays, F. Federated Learning for Emoji Prediction in a Mobile Keyboard. CoRR, abs/1906.04329, 2019. [CrossRef]
27. Wang, C.; Chen, X.; Wang, J.; Wang, H. ATPFL: Automatic Trajectory Prediction Model Design under Federated Learning Framework. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR, New Orleans, LA, USA, June 18-24, 2022, pp. 6553–6562.
28. Wu, G.; Gong, S. Peer Collaborative Learning for Online Knowledge Distillation. Proceedings of the AAAI Conference on Artificial Intelligence, AAAI, 2021.
29. Xiao, H.; Rasul, K.; Vollgraf, R. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint 2017, arXiv:1708.07747. [Google Scholar]
30. Xie, Q.; Luong, M.T.; Hovy, E.; Le, Q.V. Self-training with noisy student improves imagenet classification. Proceedings of 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR, 2020, pp. 10687–10698.
31. Xin, B.; Yang, W.; Geng, Y.; Chen, S.; Wang, S.; Huang, L. Private FL-GAN: Differential Privacy Synthetic Data Generation Based on Federated Learning. 2020 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2020, Barcelona, Spain, May 4-8, 2020, pp. 2927–2931.
32. Xu, A.; Li, W.; Guo, P.; Yang, D.; Roth, H.; Hatamizadeh, A.; Zhao, C.; Xu, D.; Huang, H.; Xu, Z. Closing the Generalization Gap of Cross-silo Federated Medical Image Segmentation. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 20834–20843.
33. Yang, C.; Xie, L.; Qiao, S.; Yuille, A.L. Training deep neural networks in generations: A more tolerant teacher educates better students. Proceedings of the AAAI Conference on Artificial Intelligence, AAAI, 2019, Vol. 33, pp. 5628–5635.
34. Yang, Z.; Li, Z.; Jiang, X.; Gong, Y.; Yuan, Z.; Zhao, D.; Yuan, C. Focal and Global Knowledge Distillation for Detectors. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 4633–4642.
35. Yoo, J.; Cho, M.; Kim, T.; Kang, U. Knowledge Extraction with No Observable Data. Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 2701–2710.
36. Zhang, L.; Luo, Y.; Bai, Y.; Du, B.; Duan, L. Federated Learning for Non-IID Data via Unified Feature Learning and Optimization Objective Alignment. IEEE/CVF International Conference on Computer Vision, ICCV, Montreal, QC, Canada, October 10-17, 2021, 2021, pp. 4400–4408.
37. Zhang, L.; Shen, L.; Ding, L.; Tao, D.; Duan, L. Fine-tuning Global Model via Data-Free Knowledge Distillation for Non-IID Federated Learning. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 10164–10173.
38. Zhang, Y.; Xiang, T.; Hospedales, T.M.; Lu, H. Deep mutual learning. Proceedings of 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR, 2018, pp. 4320–4328.
39. Zhao, Y.; Li, M.; Lai, L.; Suda, N.; Civin, D.; Chandra, V. Federated Learning with Non-IID Data. CoRR, abs/1806.00582, 2018. [CrossRef]
40. Zhu, Z.; Hong, J.; Zhou, J. Data-free knowledge distillation for heterogeneous federated learning. International Conference on Machine Learning. PMLR, 2021, pp. 12878–12889.