Preservation of the Global Knowledge by Not-True Distillation in Federated Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
In federated learning, a strong global model is collaboratively learned by aggregating clients’ locally trained models.
Although this precludes the need to accessclients’ data directly, the global model’s convergence often suffers from data heterogeneity.
This study starts from an analogy to continual learning and suggests thatforgetting could be the bottleneck of federated learning.
We observe that the globalmodel forgets the knowledge from previous rounds, and the local training inducesforgetting the knowledge outside of the local distribution.
Based on our findings,we hypothesize that tackling down forgetting will relieve the data heterogeneityproblem.
To this end, we propose a novel and effective algorithm, Federated
Not-True Distillation (FedNTD), which preserves the global perspective on locallyavailable data only for the not-true classes.
In the experiments, FedNTD showsstate-of-the-art performance on various setups without compromising data privacyor incurring additional communication costs.
Introduction
At present, massive data is being collected from edge devices such as mobile phones, vehicles andfacilities.
As the data may be distributed on numerous devices, decentralized training is often requiredto train deep network models.
Federated learning [23, 24] is a distributed learning paradigm thatenables the learning of a global model while preserving clients’ data privacy.
In federated learning,clients independently train local models using their private data, and the server aggregates them intoa single global model.
In this process, most of the computation is performed by client devices, whilethe global server only aggregates the model parameters and distributes them to clients [2, 49].
Most federated learning algorithms are based on FedAvg [37], which aggregates the locally trainedmodel parameters by weighted averaging proportional to the amount of local data that each clienthad.
While various federated learning algorithms have been proposed thus far, they each conductparameter averaging in a certain manner [1, 9, 20, 30, 46, 51].
Although this aggregation schemeempirically works well and provides a conceptually ideal framework when all client devices areactive and i.i.d.
distributed (a.k.a. LocalSGD), the data heterogeneity problem [31, 55] is a notoriouschallenge for federated learning applications and prevents their widespread applicability [19, 29].
As the client generates its own data, the data is not identically distributed.
More precisely, thelocal data across clients are drawn from heterogeneous underlying distributions; thereby, locallyavailable data fail to represent the overall global distribution, which is referred to as data heterogeneity.
Despite its inevitable occurrence in many real-world scenarios, data heterogeneity not only makestheoretical analysis difficult [31, 55] also degrades many federated learning algorithms’ performances[18, 27].
By resolving the data heterogeneity problem, learning becomes more robust against partialparticipation [31, 37], and the communication cost is also reduced by faster convergence [45, 51].
Interestingly, continual learning [41, 44] faces a similar challenge.
In continual learning, a learner model is continuously updated on a sequence of tasks, with an objective to perform well on whole
[Figure]
tasks.
Unfortunately, owing to the heterogeneous data distribution of each task, learning on the tasksequence often results in catastrophic forgetting [36, 39], whereby fitting on a new task interfereswith the important parameters for previous tasks.
As a consequence, the model parameters drift awayfrom the area where the previous knowledge is desirably preserved (Figure 1a).
Our first conjecture is such forgetting also exists in federated learning.
While the server aggregateslocal models, the distribution where they are trained may be largely different from those of previousrounds.
As a result, the global model faces the distributional shifts at each round, which may causethe forgetting as in continual learning (Figure 1b).
To empirically verify this analogy, we examine theglobal model’s prediction consistency.
More specifically, we measure the global model’s class-wiseaccuracy while the communication rounds proceed.
The observations verify our conjecture: the global model’s prediction is highly inconsistent acrosscommunication rounds, significantly reducing the performance of predicting some classes that theprevious model originally predicted well.
We dig deeper to analyze how averaging the locally updatedparameters induces such forgetting and confirm that it occurs in local training: the global knowledge,corresponding to the outside of local distribution, is prone to be forgotten.
As merely averaging thelocal models cannot recover it, the global model struggles to preserve previous knowledge.
Based on our findings, we hypothesize that mitigating the issue of forgetting can relieve dataheterogeneity (Figure 1c).
To this end, we propose a novel algorithm Federated Not-True Distillation(FedNTD).
FedNTD utilizes the global model’s prediction on locally available data, but only for thenot-true classes.
We demonstrate the effect of FedNTD on preserving global knowledge outside ofa local distribution and its benefits on federated learning.
Experimental results show that FedNTDachieves state-of-the-art performance in various setups.
To summarize, our contributions are:
• We present a systematic study on forgetting in federated learning.
The global knowledge
outside of the local distribution is prone to be forgotten and is closely related to the dataheterogeneity issue (Section 2) .
• We propose a simple yet effective algorithm, FedNTD, to prevent forgetting.
Unlike priorworks, FedNTD neither compromises data privacy nor incurs additional communicationburden.
We validate the efficacy of FedNTD on various setups and show consistentlyachieves state-of-the-art performance (Section 3, Section 4) .
• We analyze the FedNTD benefits on the federated learning.
The knowledge preservation by
FedNTD improves weight alignment and weight divergence after local training (Section 5) .
1.1 Preliminaries
Federated Learning We aim to train an image classification model in federated learning systemthat consists of K clients and a central server.
Each client k has a local dataset D k , where the wholedataset D “ Ť k Pr K s D k .
At each communication round t , the server distributes the current globalmodel parameters w p t ´ 1 q to sampled local clients K p t q .
Starting from w p t ´ 1 q , each client k P K p t qupdates the model parameters w p t q k using its local datasets D k with the following objective:
[Equation]
where L is the loss function.
At the end of round t , the sampled clients upload the locally updatedparameters back to the server and aggregate by parameter averaging as w p t q as follows:
[Equation]
Knowledge Distillation Given a teacher model T and a student model S , knowledge distillation[17] matches their softened probability q Tτ and q
Sτ using temperature τ .
The c -th value of the q τ canbe described as q τ p c q “ exp p
z c { τ q ř
i exp p
z i { τ q , where z c is the c -th value of logits vector z and C is the numberof classes.
Given a sample x , the student model S is learned by a linear combination of cross-entropyloss L CE for one-hot label 1 y and Kullback-Leibler divergence loss L KL using a hyperparameter β : [Equation]
[Equation]
Forgetting in Federated Learning
To understand how the non-IID data affects federated learning, we perform an experimental studyon heterogeneous locals.
We choose CIFAR-10 [25] and a convolutional neural network with fourlayers as in [37].
We split the data to 100 clients using Latent Dirichlet Allocation (LDA), assigningthe partition of class c samples to clients by p „ Dir p α q .
The heterogeneity level increases as theα decreases.
We train the model with FedAvg for 200 communication rounds, and 10 randomlysampled clients are optimized for 5 local epochs at each round.
More details are in Appendix C.
2.1 Global Model Prediction Consistency
To confirm our conjecture on forgetting, we first consider how the global model’s prediction variesas the communication rounds proceed.
If the data heterogeneity induces forgetting, the predictionafter update (i.e., parameter averaging) may be less consistent compared to the previous round.
Toexamine it, we observe the model’s class-wise test accuracy at each round, and measure its similarityto the previous round.
The results are provided in Figure 2a and Figure 2b.
As expected, while the server model learned from i.i.d.
locals (IID Server) predicts each class evenlywell at each round, prediction is highly inconsistent in the non-IID case (NIID Server).
In the non-IIDcase, the test accuracy on some classes which originally predicted well by the previous global modeloften drops significantly.
This implies that forgetting occurs in federated learning.
To measure how the forgetting is related to data heterogeneity, we borrow the idea of Backward
Transfer (BwT) [5], a prevalent forgetting measure in continual learning [4, 7, 8, 14], as follows:
[Equation]
(5) where A p t q c is the accuracy on class c at round t .
Note that the forgetting measure F captures theaveraged gap between peak accuracy and the final accuracy for each class at the end of learning.
The result on the varying heterogeneity levels is plotted in Figure 2c, showing that the global modelsuffers from forgetting more severely as the heterogeneity level increases.
2.2 Knowledge Outside of Local Distribution
We take a closer look at local training to investigate why aggregating the local models inducesforgetting.
In the continual learning view, a straightforward approach is to observe how fitting on thenew distribution degrades the performance of the old distribution.
However, in our problem setting,the local clients can have any class.
Given that only their portions in the local distribution differacross clients, such strict comparison is intractable .
Hence, we formulate in-local distribution p p D qand its out-local distribution ˜ p
p D q to systematically analyze forgetting in local training.
Definition 1.
Consider a local dataset D consists of N data points x i and its label y i in C -classclassification problem.
The in-local distribution vector p k “ p p D k q and its out-local distributionvector ˜ p
k “ ˜ p
p D k q are
[Equation]
[Equation]
. . , p C s , where p c
[Equation]
˜ p C s , where ˜ p
[Equation]
[Equation]
The underlying idea of out-local distribution ˜ p
p D q is to assign a higher proportion to the classes withfewer samples in local datasets.
Accordingly, it corresponds to the region in the global distributionwhere the in-local distribution p p
D q cannot represent.
Note that if p p D q is uniform, ˜ p
p D q alsocollapses to uniform, which aligns well intuitively.
An example of label distribution for 10 clientsand their p p D q s and ˜ p
p D q s are provided in Figure 3.
We measure the global and local accuracy varies on p p D q and ˜ p
p D q , as in Figure 4.
After localtraining, the local models are well-fitted towards p p D q (Figure 4a), and the aggregated global modelalso performs well on it.
On the other hand, the accuracy on ˜ p
p D q significantly drops, and the globalmodel accuracy on it also degrades (Figure 4b).
To summarize, the knowledge on the out-local distribution ˜
p
p D q is prone to be forgotten in localtraining, resulting in the global model’s forgetting.
Based on our findings, we hypothesize thatforgetting could be the stumbling block in federated learning.
2.3 Forgetting and Local Drift
First empirically observed by [55], the deviation of local updates from the desirable global directionhas been widely discussed as a major cause of slow and unstable convergence in heterogeneousfederated learning [21, 30, 31].
Unfortunately, given that it is very difficult to analyze such drift, acommon approach is to assume bounded dissimilarity between the local function gradients [20, 31].
One intriguing property of knowledge preservation on out-local distribution is that it corrects the localgradients towards the global direction.
We define the gradient diversity Λ to measure the dissimilarityof local gradients and state the effect of knowledge preservation as follows:
Definition 2.
For the uniformly weighted K clients, the gradient diversity Λ of local functions f ktowards the global function f “ 1 K ř Kk “ 1 f k is defined as:
[Equation]
Here, Λ ě 1 measures the alignment of gradient direction of the local function f k s w.r.t.
the globalfunction f .
Note that the Λ becomes smaller as the directions of local function gradients ∇ f k sbecome similar, e.g., if the magnitudes of the ∥∇ f k ∥ 2 s are fixed, the smallest Λ is obtained when thedirection of ∇ f k s are identical.
Proposition 1.
Suppose uniformly weighted K clients with in-local distribution p k “ r p k 1 , . .
, p k C s .
Ifwe assume the class-wise gradients g c are orthogonal with uniform magnitude, increasing β ď C { 2 ´ 1reduces the gradient diversity Λ from the local gradient ∇ f k “ p p k ` β ˜ p
k
q ¨ g with the ratio:
[Equation]
where β stands for the effect of knowledge preservation on the out-local distribution ˜ p k .
The 147 M
K, C ,p ą 0 is a constant term consists of K , C , and
p p k q Kk “ 1 ,
The proof is given in Appendix O. Note that here we treat f k as a sum of class-wise losses ř c p kc L c ,where L c “ E x | y “ c r L p x ; w qs is the loss on the specific class c .
When β “ 0 , there is no regularizationto preserve out-distribution knowledge, so the local model only needs to fit on the in-local distributionp k .
The above proposition suggests that the preserved knowledge on out-local distribution ˜ p k (i.e.,as β increases) guides the local gradient directions to be more aligned towards the global gradient,reducing the gradient diversity Λ .
Such forgetting perspective provides the opportunity to handle thedata heterogeneity at the model’s prediction level.
FedNTD
In this section, we propose Federated Not-True
Distillation (FedNTD) and its key features.
Thecore idea of FedNTD is to preserve the globalview only for the not-true classes.
More specifically, FedNTD conducts local-side distillationby the linearly combined loss function L between the cross-entropy loss L CE and the nottrue distillation loss L NTD :
[Equation]
Here, the hyperparameter β stands for the strength of knowledge preservation on the out-localdistribution.
Then, the not-true distillation loss L NTD is defined as the KL-Divergence loss betweenthe not-true softmax prediction vector ˜ q lτ and ˜
q gτ as follows:
[Equation]
, where
[Equation]
[Equation]
which take softmax with temperature τ only for the not-true class logits.
Figure 5 illustrates how thenot-true distillation works given a sample x .
Note that ignoring the true-class logits makes gradientsignal of L NTD to the true-class as 0 .
The detailed algorithm is provided in Appendix B.
We now explain how learning to minimize L NTD preserves global knowledge on out-local distribution˜ p p D q .
Suppose there are N number of data points in the local dataset D .
The accumulated KullbackLeibler divergence loss L KL between q l,iτ , the probability vector for the data x i , and its reference q g,iτto be matched for is:
[Equation]
By splitting the summands for the true and not-true classes, the term becomes:
L trueKL “ ´ 1 N N ÿ i “q g,iτ p
y
i q log « q l,iτ p
y i q q g,iτ p
y
i q ff , L not-trueKL
“ ´ 1 N
N ÿ i “
C ÿ c 1 ‰ y i q g,iτ p c 1 q log « q l,iτ p c 1 q q g,iτ p c 1 q ff .
Proposition 2.
Consider the in-local distribution p p D q “ r p 1 . . .
p C s such that p c “ | S c | N and itsout-local distribution ˜ p
p D q “ r ˜ p 1 , . . .
˜
p C s , where S c is the set of indices satisfying y i
“ c .
Then
the
L trueKL and L not-trueKL each are equivalent to the weighted sum on p p D q and ˜ p
p D q as
L trueKL “
C ÿ c “p c E i P S c » – ´ q g,iτ p c q log « q l,iτ p c
q q g,iτ p
c
q fffi fl (14)
L not-trueKL
C ´ 1 “
C ÿ c “˜ p
c E i R S c » – ´ q g,iτ p c q log « q l,iτ p c
q q g,iτ p
c q fffi fl
With a minor amount of calculation from Equation 13, we derive the above proposition.
The derivationis provided in Appendix P.
The proposition suggests that matching the true-class and the not-trueclass logits collapses to the loss on the in-local distribution p p D q and the out-local distribution ˜ p
p D q .
In the loss function of our FedNTD (Equation 10), we attain the new knowledge on the in-localdistribution by following the true-class signals from the labeled data in local datasets using the L CE .
In the meanwhile, we preserve the previous knowledge on the out-local distribution by followingthe global model’s perspective, corresponding to the not-true class signals, using the L NTD .
Here,the hyperparameter β controls the trade-off between learning on the new knowledge and preservingprevious knowledge.
This resembles to stability-plasticity dilemma [38] in continual learning, wherethe learning methods must balance retaining knowledge from previous tasks while learning newknowledge for the current task [35].
Experiment
4.1 Experimental Setup
We test our algorithm on MNIST [11], CIFAR-10 [25], CIFAR-100 [25], and CINIC-10 [10].
Wedistribute the data to 100 clients and randomly sample clients with a ratio of 0.1.
For CINIC-10, weuse 200 clients with a sampling ratio of 0.05.
We use a momentum SGD with an initial learning rateof 0.1, and the momentum is set as 0.9.
The learning rate is decayed with a factor of 0.99 at eachround, and a weight decay of 1e-5 is applied.
We adopt two different NIID partition strategies:
• (i) Sharding [37]: sort the data by label and divide the data into same-sized shards, and control the heterogeneity by s , the number of shards per user.
In this strategy only considers
the statistical heterogeneity as the dataset size is identical for
each client.
We set s as MNIST
( s “ 2 ), CIFAR-10 ( s P t 2 , 3 , 5 , 10 u ), CIFAR-100 ( s “ 10 ), and CINIC-10 ( s “ 2 ).
• (ii) Latent Dirichlet Allocation
(LDA) [34, 45]: assigns partition of class c by samplingp c
« Dir p α q .
In this strategy, both the distribution and dataset size are different for each
client.
We set α as MNIST ( α “ 0 . 1 ), CIFAR-10 ( α P t 0 .
5 u ), CIFAR- ( α “ 0 . 1 ), and CINIC-10 ( α “ 0 . 1 )
More details on model, datasets, hyperparameters
4.2 Performance on Data Heterogeneity
We compare our FedNTD with various existing works. The results are provided in Table 1.
Asreported in [27], even the state-of-the-art methods perform well only in specific setups, and theirperformance often deteriorates below FedAvg.
However, our FedNTD consistently outperforms thebaselines on all setups, achieving state-of-the-art results in most cases.
For each experiment in Table 1, we also report the forgetting F in the parenthesis along withthe accuracy.
Note that the smaller F value indicates the global model less forgets the previousknowledge.
We find that the performance in federated learning is closely related to forgetting, withthe performance improving as the forgetting reduces.
We believe that the gain from the prior works isactually from the forgetting prevention in their own ways.
[Table]
We emphasize that the prior works to learn from heterogeneous locals often require statefulness (i.e.,clients should be repeatedly sampled with identification) [28, 46], additional communication cost[20], or auxiliary data [33].
However, our FedNTD neither compromise any potential privacy issuenor requires additional communication burden.
A brief comparison is provided in Appendix E.
We further conduct experiments on the effect of local epochs, client sampling ratio, and modelarchitecture is in Appendix H, the advantage of not-true distillation over knowledge distillation in
Appendix I, and the effect of hyperparameters of FedNTD in Appendix J.
In the next section, weanalyze how the knowledge preservation of FedNTD benefits on the federated learning.
Knowledge preservation of FedNTD
[Figure]
In Figure 7, we present the test accuracy ondiverse heterogeneity levels.
Although both
FedAvg and FedNTD show little change inlocal accuracy on in-local distribution p p x q ,
FedNTD significantly improves the local accuracy on out-local distribution ˜ p
p x q , whichimplies it prevents forgetting.
Along with it,the test accuracy of the global model also substantially improves.
These gaps are enlargedwhen the number of local epochs becomesgreater, where the local models much deviatefrom the global model.
The accuracy curvesduring local training are in Figure 6.
It showsthat fitting on the local distribution rapidlyleads to forgetting on out-local distribution, But FedNTD effectively relieves this tendency withouthurting the learning ability towards the in-local distribution.
Our interest is how the FedNTD’s knowledge preservation on out-local distribution benefits on thefederated learning, despite of little change of its performance on in-local distribution.
To figure it out,we analyze the local models in FedNTD after local training, and suggest two main reasons for it as:
• Weight Alignment : How much the semantics of each weight is preserved?
• Weight Divergence : How far the local weights drift from the global model?
5.1 Weight Alignment
In recent studies
[Table]
it has been suggested that there is a mismatch of semantic information encodedin the weight parameters across local models, even for the same coordinates (i
[Table]
e., same position)[45, 52, 53].
As the current aggregation scheme averages weights of the identical coordinates,matching the semantic alignment across local models plays an important role in global convergence.
To analyze the semantic alignment of each parameter, we identify the individual neuron’s classpreference by which class has the largest activation output on average.
We then measure the alignmentfor a layer between two different models as the proportion of neurons that the class preference ismatched for each other.
The result is provided in Table 2.
While FedAvg and FedNTD show little difference in the IID case, FedNTD significantly enhancesthe alignment in the NIID cases.
The visualized results are in Appendix L, with more details on howthe alignment is measured.
We further analyze the change of feature after local training using a unithypersphere in Appendix M and T-SNE in Appendix N.
[Table]
W p t ´
[Table]
G vs
[Table]
W
[Table]
t
[Table]
G [Table] [Table]
[Table]
G vs
[Table]
W
[Table]
L
[Table]
W
[Table]
t ´
[Table]
q G vs
[Table]
W
[Table]
t
[Table]
G
[Table]
5.2 Weight Divergence
The knowledge preservation by FedNTD leads the global model to predict each class more evenly.
Here we describe how the global model with even prediction performance stabilizes the weightdivergence.
Consider a model fitted on a specific original distribution, and now it is trained on a newdistribution.
Then the weight distance between the original model and fitted model increases as thedistance between the original distribution and new distribution grows.
We argue that if the distance between the global model’s underlying distribution and local distributionsis close, the moved distance between the global model and local models also becomes close.
If weassume the local distributions are generated arbitrarily, the most robust choice for the global model’sunderlying distribution is a uniform distribution.
We formally rephrase our argument as the follows:
Proposition 3.
Let P : ∆ C Ñ R ě 0 be the probability measure for the client’s local distribution andΠ be the set of measure in the hypothesis.
Assume that the class-wise loss L c p w q “ E x | y “ c r L p x ; w qsis λ -smooth and w c be a optimum of L c .
Then the loss L of the client with distribution p becomes:
[Equation]
Then for arbitrary P P Π , the uniform distribution attains the minimax value:
[Equation] [Equation]
The proof is provided in Appendix R.
Although the global model’s underlying distribution is unknown,the normalized class-wise accuracy vector is a handy approximation for it as: A G “ 1 A ¨ r a 1 , . . .
a C s ,where A is the global model’s test accuracy and a c is its class-wise accuracy on the class c .
The results in Figure 8 empirically validate our argument.
There is a strong correlation betweenweight divergence } w ´ w k } (for global model w and client k ’s local model w k ) and distributiondistance } A G ´ p k } (for client k ’s distribution p k ).
By providing a better starting point for localtraining, FedNTD effectively stabilizes the weight divergence.
Related Work
Federated Learning (FL) is proposed to update a global model while the local data is kept inclients’ devices [23, 24].
The standard algorithm is FedAvg [37], which aggregates trained localmodels by averaging their parameters.
Although its effectiveness has been largely discussed in i.i.d.settings [43, 47], many algorithms obtain the sub-optimal when the distributed data is heterogeneous[31, 55].
Until recently, a wide range of variants of FedAvg has been proposed to overcome sucha problem.
One line of work focuses on local-side modification by regularizing the deviation oflocal models from the global model [1, 20, 30, 46].
Another is the server-side modification, whichimproves the efficacy of aggregation of local models in the server [9, 33, 45, 54].
Our work aims topreserve global knowledge during local training, which belongs to the local-side approach.
Forgetting View in FL A pioneer work that considers forgetting in FL is FedCurv [42].
It regardseach local client as a task, and [42] regulates the change of local parameters to prevent accuracydrop on all other clients.
However, it needs to compute and communicate parameter-wise importanceacross clients, which severely burdens the learning process.
On the other hand, we focus on theclass-wise forgetting and suggest that not-true logits from local data contain enough knowledge toprevent it.
A concurrent work of our study is [48], which also reports the forgetting issue in localclients by empirically showing the increasing loss of previously learned data after the local training.
To prevent forgetting, [48] exploits generated pseudo data.
Instead, we focus on the class-wiseforgetting and suggest that not-true logits from local data contain enough knowledge to prevent it.
The continual learning literature is further discussed in Appendix D.
Knowledge Distillation (KD) in FL
In FL, a typical approach is using KD to make the globalmodel learn from the ensemble of local models [9, 26, 33, 56].
By leveraging the unlabeled auxiliarydata, KD effectively tackles the local drifts by enriching the aggregation.
However, such carefullyengineered proxy data and may not always be available [30, 54, 57].
Although more recent worksgenerate pseudo-data to extract knowledge by data-free KD [54, 57], they require additional heavycomputation, and the quality of samples is sensitive to the many hyperparameters involved in theprocess.
On the other hand, as a simple variant of KD, our proposed method surprisingly performswell on heterogeneity scenarios without any additional resource requirements.
Conclusion
This study begins from an analogy to continual learning and suggests that forgetting could be amajor concern in federated learning.
In our observations, the knowledge outside of local distributionis prone to be forgotten in local training and is closely related to the unstable global convergence.
To overcome this issue, we propose a simple yet effective algorithm, FedNTD, which conductslocal-side distillation only for the not-true classes to prevent forgetting.
FedNTD does not have anyadditional requirements, unlike previous approaches.
We analyze the effect of FedNTD from variousperspectives and demonstrate its benefits in federated learning.
Broader Impact We believe that Federated Learning is an important learning paradigm that enablesprivacy-preserving ML.
This paper suggests the forgetting issue and introduces the methods to reliveit without compromising data privacy.
The insight behind this work may inspire new researches.
However, the proposed method maintains the knowledge outside of local distribution in the globalmodel.
This implies that if the global model is biased, the trained local model is more prone to have asimilar tendency.
This should be considered for ML participators.