Global Knowledge Distillation in Federated Learning

1 Introduction

Federate Learning (FL) is an widely used machine learning framework that learns a global model across multiple decentralized clients. The clients independently train local models using their own data and send the local trained models to a server to aggregate a global model without sharing its own data. Due to the advantages in communication-efficient and privacy-preserving, FL has shown its potential to facilitate real-world applications, such as healthcare, biometrics, computer vision and natural language processing.

Under federated training framework, training methods that allow for flexible local updating and low client participate. FedAvg aistats/McMahanMRHA17 is one classic training algorithm in FL, which directly averages of the participate clients’ model updates and send back the aggregated model to the clients. The client can randomly choose whether to join this round of training. The model convergence can be speed up by this loose aggregation framework. The key challenges in federated learning are system heterogeneity (communication delay and various device) and data heterogeneity (structure different and non-i.i.d data).

For data structure heterogeneous problem, knowledge distillation is one solution in FL, where the heterogeneous models in each clients can be aggregated into one model. The global model as student learns knowledge from multiple client’s model as teacher.

For non-i.i.d data problem, only a few effects have been made, such as FedProx, SCAFFOLD, and FedNova. According to these studies, the non-i.i.d data settings can degrade the effectiveness of all machine learning processing. This is because each client only find an local optimized model based on local non-i.i.d dataset. The aggregated global model can not generate a good solution when the server does not receive enough participate clients’ updates.

To over come the above problem, we propose a novel global knowledge distillation method, named FedGKD, which learns the knowledge from past global models to tackle down the local bias training problem. By learning from global knowledge and consistent with current local models, FedGKD learns a global knowledge model in FL.

Overall, the primary contributions of this paper are summarized as follows:

To the best of knowledge, this paper is the first work that uses knowledge distillation technique to transfer the historical global models information to local model training for preventing over-biased local model training on non-i.i.d data distribution.
We provide a generalization method for federated knowledge distillation based on projection. In some experimental settings, we can see that projection could further help to improve the performance of the federated learning.
We show in extensive experiments and analysis on various dataset (i.e. CIFAR-10/100) and settings (i.e. non-i.i.d data distribution) validate the effectiveness of our proposed FedGKD, when comparing with several state-of-the-art methods.

2 Notations and Preliminaries

Federated Learning

In general, FL aims to learn a global model weight $w$ that minimizes its risk on each of the user tasks $T_{m}$ mcmahan2017communication:

min w E_{T_{m} \in T} [L_{m} (w)],

(1)

where $T = {T_{m}}_{m = 1}^{M}$ is the collection of user tasks. In most cases, FL assumes the same task (i.e., $m = 1$ ) but with different data distribution on each local side, then Eq. 1 is empirically optimized by ${min}_{w} \frac{1}{K} \sum_{k = 1}^{K} L (w_{k})$ , where $L (w_{k})$ is the empirical risk over each local dataset $D_{k}$ . In FL, a global data $D$ is distributed to multiple local clients, with $D = \cup {D_{k}}_{k = 1}^{K}$ .

Knowledge Distillation

KD is referred as teacher-student paradigm that a cumbersome but powerful teacher model transfers its knowledge through distillation to a lightweight student model buci2006model

. One most popular approach in machine learning is KD minimizes the discrepancy between the logits outputs from the teacher model and the student model with a proxy dataset

D^{p}

hinton2015distilling

. The discrepancy could be measured by Kullback-Leibler divergence:

min w^{S} E_{x \sim D^{p}} [f (x; w^{T}) | | f (x; w^{S})],

(2)

where $f (\cdot)$ is the output of an non-linear complex function $f$ , such as logits, softmax outputs, etc.

The core idea of KD has been used for FL to solve the user heterogeneity problem, also well known as non-i.i.d data distribution problem lin2020ensemble; chen2020fedbe. These approaches treat the local models as teachers and transfer their knowledge into a student (global) model to improve its generalization performance. However, former works require a proxy dataset which is not practical in most scenarios. Then, we propose our new learning approach that tackles the problem in more practical settings, such as no proxy dataset, no additional training on the global side, and no additional privacy leakage besides sharing local gradients.

3 Local Training via Global Knowledge Distillation

In this section, we first introduce the proposed Federated Local Training via Global Knowledge Distillation (FedGKD) and its key features.

1:Notations. total number of clients

K

, server

S

, total communication rounds

T

, local epochs

E

, fraction of participating clients

C

, learning rate

η

B

is a set holding client’s data sliced into batches of size

B

3://On the server side

4:procedure ServerExecution

▹

Server Model Aggregation

5: Initial global weight

w^{0}

6: for each communication round

t = 1, \dots, T

7: Send the global weight to all clients

K^{t} \leftarrow

random sample a set of

C \cdot K

clients

9: for each client

k \in K_{t}

in parallel do

10:

Δwt+1k←\textscClientUpdate(k,wt)

11:

w^{t + 1} \leftarrow w_{t} + \frac{1}{C \cdot K} Δ w_{k}^{t + 1}

12://On the data owners side

13:procedure ClientUpdate(

k, w^{t}

)

▹

Local Model Training

14:

w \leftarrow w^{t}

15: for each local epoch

i = 1, \dots, E

16: for each batch

b \in B

17:

w \leftarrow w - η \cdot ▽ L (w, b)

▹

Update Local Weights via Eq.3

18:

Δ w^{t + 1} \leftarrow w - w^{t}

19: Return

Δ w^{t + 1}

back to server

Algorithm 1 FedGKD: Global Knowledge Distillation in Federated Learning

FedGKD is proposed for transferring the global knowledge to local model training, where each client’s data is in non-i.i.d data distribution. Due to the natural non-i.i.d data distribution, the local model training will strongly bias on its own dataset, which cause the complexity and difficulty of global aggregation. In order to mitigate the data and information gaps between each local model due to the data bias, a novel idea stems from the knowledge distillation technique that transfers the past global models information to local model training for preventing over-biased local training. Mathematically, we formulate the loss of FedGKD as a linear combination of original cross-entropy loss and the distillation losses:

	$L\textscFedGKD=λ0LCE(q,pY)$	$+ λ_{1} L_{G K D} (q, q^{G^{1}}) + \dots + λ_{T - 1} L_{G K D} (q, q^{G^{T - 1}}),$		(3)
		$s.t. λ_{0} + λ_{1} + \dots + λ_{T - 1} = 1$

where $q$

is the softmax probability of each local client model,

q^{G^{t}}

is the softmax probability of the former global model at

t

round, and

p^{Y}

is the one-hot label. The distillation loss

L_{G K D}

is defined by Kullback-Leibler divergence same as Eq. 2

. Note that, in our experimental evaluation, we find that using the global model in the last round already can achieve a good performance, which is also easier to set up the hyperparameter. However, we believe that some adaptive global model selection approaches could improve performance after the model converged.

4 Experiment

4.1 Experiment Setup

We evaluate different state-of-the-art FL methods on CIFAR-10 and CIFAR-10 in different FL scenarios. ResNet-8 is used to validate the methods. Considering Batch Normalization(BN) fails on heterogeneous training data, we replace BN by Group Normalization(GN) to produce stabler results and alleviate some of the quality loss brought by BN.

We follow prior work,using the Dirichlet distribution to create disjoint non-i.i.d. client traning data. The value of $α$ controls the degree of non-i.i.d.-ness, the smaller $α$ indicates higher data heterogeneity.

Figure 1: Illustration of # of samples per class allocated to each client

Figure 2: Illustration of convergence rate

Dataset and Evaluation Metrics

We evaluate the learning of different SOTA FL methods on CV task on architectures of ResNet cvpr/HeZRS16. We consider federated learning CIFAR-10/100 krizhevsky2009learning.

Baselines

FedAvg: FedAvg aistats/McMahanMRHA17 is a classic federated learning algorithm, which directly use averaging as aggregation method.
FedProx: FedProx li2018federated improves the local objective based on FedAvg. It directly limits the size of local updates to solve non-i.i.d problem.
MOON: MOON li2021model is a contrastive federated learning method.
FedGKD: The proposed method in Algorithm 1.

Implementation Details

The FL algorithm randomly samples a fraction of clients per communication round for local training. The optimizer used is SGD,the learning rate in our experiments is set as 0.05, and weight decay is 1e-5, momentum is 0.9.

4.2 Evaluation on the Common Federated Learning Settings

Results

As shown in Table 1 and Table 2, the proposed method FedGKD outperforms previous state-of-the-art methods. FedGKT shows superiority when participant ratio is low. This is because when participants are few, the local optimal models will greatly affect the global aggregation.This characteristic helps in the real FL scenarios.

Method	C=0.2		C=0.4		C=0.6
	Best	Final	Best	Final	Best	Final
FedAvg	77.30 $\pm$ 0.72	74.70 $\pm$ 1.09	78.85 $\pm$ 0.32	77.28 $\pm$ 0.42	79.51 $\pm$ 0.34	79.24 $\pm$ 0.34
FedProx	76.39 $\pm$ 0.59	75.45 $\pm$ 0.71	77.45 $\pm 0.53$	76.35 $\pm$ 0.68	77.68 $\pm$ 0.15	77.46 $\pm$ 0.28
MOON	78.60 $\pm$ 0.77	76.27 $\pm$ 1.61	79.03 $\pm$ 0.14	77.10 $\pm$ 0.63	79.49 $\pm$ 0.53	79.45 $\pm$ 0.45
FedGKD	79.45 $\pm 0.42$	78.17 $\pm$ 0.38	79.36 $\pm$ 0.72	78.06 $\pm$ 0.45	79.6 $\pm$ 0.10	79.39 $\pm$ 0.18
FedGKD-PJG	78.58 $\pm$ 0.68	77.95 $\pm$ 0.66	79.38 $\pm$ 0.33	78.73 $\pm$ 0.40	79.8 $\pm$ 0.26	79.72 $\pm$ 0.29

Table 1: Comparison of different FL methods in different sampling ratio, we evaluate on CIFAR-10 with ResNet-8

Dataset	Setting	FedAvg	FedProx	MOON	FedGKD
CIFAR-10	$α$ =1	74.70 $\pm$ 1.09	75.45 $\pm$ 0.71	76.27 $\pm$ 1.61	77.95 $\pm$ 0.66
CIFAR-10	$α$ =0.1	41.94 $\pm$ 4.59	63.43 $\pm$ 0.36	45.67 $\pm$ 5.02	64.00 $\pm$ 2.83
CIFAR-100	$α$ =1	41.70 $\pm$ 0.58	41.23 $\pm$ 0.66	39.83 $\pm$ 0.96	42.33 $\pm$ 0.64
CIFAR-100	$α$ =0.1	30.89 $\pm$ 0.91	31.32 $\pm$ 0.41	27.75 $\pm$ 0.41	33.63 $\pm$ 0.36

Table 2: Performance overview given different data settings for ResNet-8 on CIFAR(20 clients with C=0.2, 100 communication rounds, and 20 local epochs per round)