Tackling Personalized Federated Learning with Label Concept Drift via Hierarchical Bayesian Modeling
Abstract
Federated Learning (FL) is a distributed learning scheme to train a shared model across clients.
One fundamental challenge in FL is that the sets of data across clients could be non-identically distributed.
Personalized Federated Learning (PFL) attempts to solve this challenge.
Most methods in the literature of PFL focus on the data heterogeneity that clients differ in their label distributions.
In this work, we focus on label concept drift which is a broad but relatively unexplored area.
We present a general framework for PFL based on hierarchical Bayesian inference and propose a variational inference algorithm based on this framework.
We demonstrate our methods through empirical studies on CIFAR100 and SUN397.
Experimental results show our approach significantly outperforms the baselines when tackling the label concept drift across clients.
Introduction
Since the introduction of FL in [12] in 2016, there has been increasing attention for this distributed learning setting, as it turns out to be a general framework to address the privacy concerns arise from many different areas.
Despite its strength, there are also many challenges in the application of FL.
One of them is the statistical heterogeneity of client data sets.
Since clients sit in various environments, their data exhibit certain concepts that correlate with the local environments and deviate from each other.
As summarized in [8] there exist different kinds of data heterogeneity.
In this work, we focus on the so-called label concept drift .
That is, the same label can have different features for different clients.
As a result the aggregated global model is not optimal for every individual client.
Personalized federated learning (PFL) methods target data heterogeneity issue and intend to achieve a better local utility via personalized model at the clients’ side.
Many of them view the learning process of PFL through optimization [5; 6; 7; 13], while a few [1; 4; 2] try to understand PFL as a posterior Bayesian inference problem.
However, in the current literature on PFL, the most methods [1; 3; 7; 15; 13; 5; 6] focus on label distribution skew , i.e. clients have data corresponding to different labels.
In this work we also take a Bayesian perspective, but aim at label concept drift and manage to capture the common trend while allowing each client to specialize in individual concepts.
Our contributions In this work we provide a variational Bayes framework for PFL to tackle label concept drift.
We demonstrate the effectiveness of our proposed approach through extensive empirical experiments.
Method
In this section, firstly we state the problem studied in this work, then we describe the proposed variational Bayesian framework for PFL under the label concept drift.
2.1 Problem Formulation
Let the data distribution of client j be P j ( x, y ) where x, y denote data and label, we present the definition of label concept drift following [8].
Definition 2.1 (Label Concept Drift [8]) .
Let J be the number of clients, label concept drift indicates that the conditional generating distributions { P j ( x | y ) }
Jj =1 are different across different clients, but marginal distributions { P j ( y ) }
Jj =1 are the same.
In FL a single shared model is used to fit on all clients’ data, instead of doing that in PFL we aim to solve the following minimization problem:
[Equation]
where J is the number of clients, D j is the set of data available in client j , w 1: J is a shorthand for the set of parameters { w 1 , · · · , w J } , w j is the personalized parameter for the j -th client and f j ( w j ; D j ) is the empirical average loss function of the j -th client:
[Equation]
where ( x ( j ) i , y ( j ) i )
∈
D j is one data point of client j , l ( · , ·| w ) is the loss function using weight parameter w and n j := |D
j | is the number of data points on the j -th client.
2.2 The Augmented Joint Distribution
To develop a Bayesian framework, we need to obtain a posterior distribution for parameters which we are interested in.
One way to obtain the posterior distribution p ( w j |D ) is by performing Bayesian inference on the j -th client locally, i.e. p ( w j |D ) :
= p ( w j |D j ) .
Given a vague prior p ( w j ) , one disadvantage of this approach is the variance of p ( w j |D j ) could be high if the number of data points on client j is small.
In another way, since all clients are running similar tasks, D { 1 , ··· ,J }\ j should be able to provide information to form the posterior of w j .
Therefore we introduce a global variable w such that all w 1: J depend on w and w captures the correlations between different clients, namely the common trend of the task.
This also implies the conditional independence between w i and w j :
[Equation]
Using Bayes’ theorem, the augmented posterior distribution of { w , w 1: J } is proportional to the product of the prior and the likelihood function, thus we have:
[Equation]
where p ( w ) is the prior of the introduced global variable, f j ( w j |D j ) is defined in Equation (2) and is proportional to the negative of the data log-likelihood on client j .
Based on the above augmented model, two inference algorithms could be used.
Maximum a Posteriori Probability (MAP) seeks a maximizer to the unnormalized posterior and avoids the intractable integration of the model evidence.
However it can be difficult to select a set of hyper-parameters for the conditional priors.
To mitigate this limitation, we present an algorithm based on the principle of maximizing the marginal likelihood in Section 2.3.
The optimization will be conducted using variational expectation maximization.
2.3 Maximize the marginal likelihood
We assume an isotropic Gaussian as the form of conditional prior p ( w j | w ) and use factorized variational approximation q ( w 1: J ) :
= (cid:81)
Jj =
1 q j ( w j ) to the true posterior distribution p ( w 1: J |D ) .
Additionally, the axis-aligned multivariate Gaussian is used as the variational family, that is, q j ( w j ) =
N ( w j | µ j , Σ j ) and Σ is a diagonal matrix.
To optimize these approximations { q j ( w j ) } J j =1 , we maximize the evidence lower bound (ELBO) of the marginal likelihood:
[Equation]
The above ELBO can be optimized using variational expectation maximization through blockwise coordinate descent.
To obtain the variational approximation q ( w j ) , we only need to maximize E q ( w j ) [log p ( D j | w j )] − KL[ q ( w j )
∥ p ( w j | w , ρ 2 j )] using the local set of data on client j .
After these local approximations have been formed, clients upload their own variational parameters to the server.
Given these parameters, the server tries to optimize the ELBO in Equation (5) by updating the global variable w .
Simplifying Equation (5), we obtain the objective function on the server:
[Equation]
Setting the gradient of Equation (6) w.r.t.
w and ρ 21: J to be zero, we obtain closed form for the optimization of the global variable:
[Equation]
where D is the dimension of w and Tr( Σ j ) is the trace of the variational variance-covariance parameter.
We summarize the update rules in the following Algorithm 1.
Input: T rounds, µ 01: J , Σ 01: J , w 0 , ρ 01: J , sampling ratio r
Output w T 1: J 1: for t = 0 to T − 1 do 2:
Server sends w t to all clients 3: ▷ clients update their personalized variational distributions ◁ 4: for j = 1 , . .
, J do 5:
Update ρ t +1 j by [Equation] D 6: ( µ t +1 j , Σ t +1 j ) ∈ arg min ( µ j , Σ j ) E q ( w j ) [log p ( D j | w j )] − [Equation]
[ q ( w j ) ∥ p ( w j | w , ρ 2 j )] 7:
Server select a random subset of clients I t of size ⌊ J × r ⌋ 8:
Each client j ∈ I t sends its updated variational parameters µ t +1 j and ρ t +1 j to the server 9: ▷ server optimizes the global variable [Equation]
[Equation]
Experiments
To evaluate our methods, we target image classification tasks and compare our methods against baselines under various settings.
We present results of the following frameworks: 1) Local: all clients train locally; 2) FedAvg [11]: all clients rely on an aggregated global model; 3) pFedPer [3]: the network consists of representation model and linear classifier, the representation model participates in collaboration while every client trains its own linear classifier; 4) pFedHN
[14]: the server uses a hypernetwork and client-specific latent variable to generate a personal network for each client.
5) pFedME [16]: bi-level optimization of local models and the global model using Moreau envelops.
When applying our method, we follow the treatment in pFedPer, which considers the entire network comprising a representation model and a linear classifier, we only personalize the linear classifier layer while let the representation model to be trained via FedAvg to generate a general representation.
Unifying the representation model could have additional profit for the representation model as it has many parameters while individual client has possibly limited data points.
We conduct experiments on two variants of our method: 1) pFedVEM: our approach of variational inference.
Every client has an estimation of the conditional distribution p ( w j | w ) over its personalized linear classifier and the global variable w of linear classifiers w 1: J are aggregated w.r.t.
the estimated distributions { p ( w j | w ) } Jj =1 . 2) pFedMAP:
We also test a maximum posterior method based on our hierarchical Bayesian framework under non-informative prior of w and normal distributed p ( w j | w ) with predefined variance ρ j .
We set ρ 1: J to be the same to avoid a prohibitive tuning workload.
The resulting algorithm looks similar to Ditto [10].
[Table]
3.1 Training protocol
To better present label concept drift, i.e. varying P j
( x | y )
[Table]
we use hierarchical datasets CIFAR100 [9] and SUN397 [17], which contain superclasses and subclasses.
We set the classification task to be superclass prediction, while for every client the data of each superclass is sampled from a random subclass and hence label concept drift is induced.
To get closer to reality, we allow the number of observations by each client to be different by random sampling without replacement.
When running the experiments the number of communication rounds is set to be 100.
We adopt the communication protocol proposed in pFedMe [16], i.e. at the beginning of each round, the server broadcasts the aggregated model to all clients, while only a subset of clients send their parameters back to the server.
Each client has a probability of 0 .
1 to be sampled, we evaluate all frameworks with the number of clients C ∈ { 50 , 100 , 200 } .
We use a CNN network with five convolutional layers followed by a fully connected layer.
In case the framework consists of two speparate components, the convolutional layers block is used as the representation model and the fully connected layer as the linear classifer.
We run all experiments five times on a cluster within the same container environment.
The the same group of five random seeds has been used for local data sampling, clients sampling and parameters initialization.
We then present the mean and standard error.
3.2 Results
We evaluate all the frameworks over various settings, the results are given in Table 1.
We observe that: 1) Although the label concept drifts among clients, FedAvg still outperforms Local.
When the expected # Sample/Clients is large, e.g. on CIFAR100, FedAvg even outperforms some pFed frameworks, implying that although label concept drift across clients using a single global model is still adequate when all participants have sufficient data; 2) Our pFedMAP, which is lightweight and easy to implement, is competitive compared with previous PFL methods; 3) Our pFedVEM always achieves the best accuracy and performs significantly better than baselines over all settings.
Conclusions
In this paper, we addressed the problem of personalized federated learning when label concept distributions differ between clients.
We developed an evaluation framework for PFL with concept drift, and proposed a variational Bayes framework for PFL.
Experimental results show our method yields consistently superior performance to main competing frameworks.
1) When conducting the local framework, we train the model for 100 epochs. 2)
pFedHN limits the to connect a single client at each iteration, we let it connect with sampled clients in series at each round.
Note that in reality this leads to a significant higher time complexity than other frameworks.