Personalized Multi-tier Federated Learning
Abstract
The challenge of personalized federated learning (pFL) is to capture the heterogeneity properties of data with inexpensive communications and achieving customized performance for devices.
To address that challenge, we introduced personalized multi-tier federated learning using Moreau envelopes (pFedMT) when there are known cluster structures within devices.
Moreau envelopes are used as the devices’ and teams’ regularized loss functions.
Empirically, we verify that the personalized model performs better than vanilla FedAvg, per-FedAvg, and pFedMe.
pFedMT achieves 98 .
30% and 99 .
71% accuracy on MNIST dataset under convex and nonconvex settings, respectively.
introduction
Federated learning (FL) is a distributed on-device learning framework that employs the heterogeneous data privately available at the edge for learning.
In traditional machine learning, edge devices are supposed to send data to the global server for training.
However, federated learning relaxes this restriction by enabling training each model on the end devices and aggregate them on the global server.
Classical federated learning has an objective to compute a single global model from the device’s private data in a communication-efﬁcient and privacy-preserving manner [1].
Virtually all methods in FL follow the following basic steps: (1) The global server sends the initial global model to all participating devices.
(2) A subset of devices is selected to compute local models.
(3) Each selected device learns a local model with their private data.
The global server collects and aggregates all the local models, produces the global model and sends it back to all participating devices.
The steps 2 to 4 continues until the desired accuracy is achieved.
FL suffers from data and system heterogeneity among devices with size and distribution of local data as well as storage and computation capacity, respectively.
These challenges adversely impact the model’s convergence and learning efﬁciency.
Data heterogeneity means the data is disseminated across devices in a non-independent and identical manner (non-IID).
That raises a problem in generalizing the global model for each device.
Unlike conventional FL, personalized FL aims to learn device-speciﬁc models along with a global model to reduce the generalization error and achieve customize learning performance [2].
In FL, it is typically assumed that devices can be grouped into a single cluster [1, 3], and there will be a single global server to which the devices communicate.
Multi-tier FL differs from the vanilla
FL architecture.
Here, devices have no direct communication with the global server.
Another layer called "team" is present between the global server and devices that have direct contact with the global server.
The idea of multi-tier FL is similar to the Cloud-Edge model.
The cloud acts as a "global server", responsible for producing a global model.
Edge servers in different geographical regions acts as "team server" and connect with the cloud and end devices.
End devices compute the local model.
Multi-tier framework is more realistic to solve the real world federated learning problems that includes Cloud-Edge architecture [4].
A non-smooth function can be made smooth using the Moreau envelope.
Motivated by the excellent performance of Moreau envelope [5] in personalized FL [2], we propose multi-tier federated learning where devices simultaneously minimize the Moreau envelopes of the device loss function and build the global model as in the classical FL.
In our proposed architecture, we leverage the hierarchical architecture that has three levels, such as global server, team, and device, and simultaneously learn (1) a personalized model for each device, (2) a personalized model for each team, and (3) a global model.
Geometrically, under this multi-tier system, the global model can be seen as a "central point" where all teams agree.
In contrast, personalized models are the points in various directions that devices take in accordance with their data distributions.
Our algorithm only allows devices to communicate within their team.
Team servers of each team communicates with the global server.
Overall objective of pFedMT is to achieve customized on-device model performance as well as improve performance in global models by in-expensive number of global iterations.
Our key contributions are as follows:
First we formulate the optimization problem for multi-tier personalized FL (pFedMT) by using Moreau envelope as a regularized loss function for both devices and teams.
Second, we designed the algorithm for pFedMT on the basis of multi-tier FL architecture which works for both convex and non-convex settings.
Finally, we empirically evaluated the performance of pFedMT using both benchmark and synthetic datasets in non-iid settings.
pFedMT outperforms the vanilla FedAvg [1], Personalized FedAvg [6], and pFedMe [2] for both convex and non-convex settings.
From the empirical analysis, we have also shown that pFedMT achieves customize on-device performance with in-expensive communication iterations.
The rest of the paper is organized as follows:
Related work is discussed in Section 2.
After that, in
Section 3 we proposed pFedMT.
Experiments and results are given in Section 4.
Related Work
Our work is relevant to hierarchical, multi-task, and personalized federated learning; therefore, here we provide a brief literature survey on these.
Hierarchical federated Learning Most of the recent work in FL assumes a single global parameter server, which may be located either at the edge or in the cloud.
The cloud server has more access to data but is not communication efﬁcient and expected high latency, whereas the edge servers are more communication efﬁcient with devices.
The hierarchical model proposed in [7] takes advantage of both the cloud and edge in federated learning.
In [8], authors introduced the upward and downward divergences in hierarchical FL.
Another work in [9] applied personalization in hierarchical FL to reduce the generalization error at each device.
Multi-task federated learning MOCHA [3] was one of the earliest works on multi-task learning (MTL) in federated settings that considered issues of high communication cost, stragglers, and fault tolerance for distributed multi-task learning.
Multi-task learning helps to handle statistical heterogeneity problems and extended to personalization.
In [10], proposed an MTL approach under a mixture of distributions in federated settings.
Duan et.
al.
[11] introduced a group of adaptive techniques that carefully manage task differences while automatically utilizing potential similarities.
Personalized federated learning Due to its simplicity and low connection cost, FedAvg [1] has become the preferred method for federated learning.
But it suffers from concept-drift when the data is heterogeneous (non-iid), leading to unstable and sluggish convergence.
Concept drift is a problem where no single global model works well for all clients.
To solve this problem and generate a better convergence rate, over the past few years, FL has placed much emphasis on personalization.
SCAFFFOLD [12] is one of the earliest algorithm that requires fewer communication iterations without affected by data heterogeneity.
In [6], a personalized version of model-agnostic meta-learning (MAML) [13] is proposed for federated learning.
The main objective of that work is to ﬁnding an initial shared model that users may quickly adapt to their local dataset by executing few gradient descent steps using their own data.
In [14], authors conducted a systematic learning-theoretic study of personalization and proposed three approaches: user clustering, data interpolation, and model interpolation, that are model-agnostic.
pFedMe
[2], authors formulate personalized FL by using
Moreau envelopes as devices’ regularized loss function.
In [9], authors proposed a hierarchical personalized FL approach where the users come from multiple known clusters.
The algorithm simultaneously learns three models: a global model, a cluster-speciﬁc model for each cluster, and a personalized model for every device.
The architecture is similar to pFedMT, except that we applied Moreau envelope as a regularizer.
Few more recent works on personalized FL are as follows, [15, 16, 17].
The motivation behind pFedMT is two-fold: One, it is reasonable to assume that closer devices contain similar information but data is more heterogeneous for distant devices.
By adding personalization both at team and device levels, we aim to capture a more elaborate personalization model that aligns better with the real-world applications.
Two, communication with the global server is often the most expensive step in FL, but communications within a team is typically much cheaper.
By exploiting the multi-tier architecture and accommodating a large portion of the communication within the teams, pFedMT economizes signiﬁcantly on the communication iterations with the global server.
pFedMT
In a multi-tier FL framework (Figure 1), all devices are divided into M teams.
Each team has a team server ( TS i ) and it is connected with N i devices.
Devices belonging to the same team are supposed to come from similar geographical locations; therefore, devices within a team have more similar models than devices belonging to another team.
Devices within a particular team only communicate with the respective team server( N i ).
Similarly global server ( GS ) only communicates with the team server for global updates.
[Figure]
3.1 pFedMT formulation
We consider a multi-tier FL setup with M teams, with N i devices in each team for i = 1 , . . .
, M .
The standard empirical risk minimization for this problem can be written as minimize
[Equation]
where f i,j ( · ) is the loss function of j th device from i th team.
This formulation involves a single decision variable x that all clients are expected to agree on when the model converges.
This is unsuited for problems in which the clients’ data distributions are non-homogenous because a global model that works for all devices may be nonexistent.
This problem can be addressed effectively by replacing the local objective functions with their Moreau envelopes [2].
We design a multi-tier personalized FL formulation by following this idea.
Let us ﬁrst deﬁne the team-level objective function with device-level personalization:
[Equation]
, where
[Equation]
˜ f
i,j ( · ) is the Moreau envelope of f i,j ( · ) , w i is the model of i th team, θ i,j is the personalized model of j th device in i th team, and λ ≥ 0 is a hyperparameter that controls the device-level personalization impact.
We use Moreau envelope again on the team-level objective to obtain our global loss function,
[Equation]
F i
[Equation]
x ) where
[Equation]
Here, ˜ F i ( · ) is the Moreau envelope of F i ( · ) , x is the global (server) model, w i is the personalized model of i th team, and γ ≥ 0 is a hyperparameter that controls the team-level pesonalization.
The goal in pFedMT is to minimize φ ( x ) .
We use the gradient method: Starting from an initial estimate x 0 ∈ R d , and for a given step-size β > 0 , for t = 0 , 1 , . . .
, T x t + [Equation] =
[Equation]
where the last equality follows from the deﬁnition of the gradient of Moreau envelope:
[Equation]
In general, prox F i /γ ( x t ) can be difﬁcult to compute, but we can compute it by considering the
Moreau envelope as a subproblem,
[Equation]
At every global round t , each team approximates the solution of Problem 6 by using the gradient method.
Starting from w t, 0 i = x t , and for given step-size η i > 0 , we compute for k = 0 , 1 , . . .
, K [-,w t,k +
 i-],
[Equation]
where the last equality follows from the deﬁnition of the gradient of Moreau envelope:
[Equation]
Once again, prox f i,j /λ ( w t,ki ) can be difﬁcult to compute.
Therefore we handle this operation as a subproblem solved at the device-level:
[Equation]
In every team-level update, each device approximate prox f i,j /λ ( w t,k i ) by applying gradient method to problem 9.
Starting from θ t,k, 0 i,j = w t,ki , and using a step-size α > 0 , we compute for l = 0 , 1 , . . . , L ,
[Equation]
t,k,l [Equation] i [Equation]
[Equation]
This framework gives us a triple-loop mechanism for personalized multi-tier FL, described below.
3.2 Algorithm
We propose an algorithm for personalized multi-tier federated learning using Moreau envelope (pFedMT).
Global server initializes the global model ( x 0 ).
Every team connected to the global server, copies the global model to their respective team model ( ∀ Mi =1 w 0 ,
i = x 0 ).
Each device within a team copies the initial team model ( w 0 , 0 i ) as their local model( θ i,j ).
Global server initializes total number of global iterations, team iterations, and local iterations as T, K, and L respectively.
At each global iteration t (step 1), global server broadcasts global model x t to every teams (step 2).
Similarly, at each team iteration k , each team boradcasts w t,k i to all the devices within the team (step 4).
For each local iteration ( l ), each device D i,j solves Equation 10 separately but in parallel to obtain the personalized model θ t,k,l i,j (steps 5 and 6).
The team server ( N i ) of each team collects the device updates from the respective devices after L local iterations and perform aggregation ( ¯ θ t,ki ) on the device updates (step 8) which is similar to weighted FedAvg [18].
δ j is a weighting factor based on the ratio between the device’s number of samples and the total number of samples.
And also, (cid:80) N i j =1 δ j = 1 .
Each N i produce personalized model ( w t,ki ) using Equation 7(step 9).
Each team broadcast the updated team level model to the devices registered with that team and continue stepsto 10 for next K − 1 team iterations.
After all teams ﬁnished K team iterations, the global server collects team updates ( w t,K i ) from each team and performs weighted aggregation ( w t ) on team updates (step 11).
τ i is a weighting factor that depends on the ratio between the number of samples of the team and the total amount of data.
Similar to δ j , (cid:80) Mi
=1 τ i = 1 .
The global server produces a global update( x t ) by solving Equation 4 (step 12).
The global aggregation and global update both would be similar if value of γ and β = 1.
The global server broadcasts the updated global model to teams and continues the steps 1 to 12 for the upcoming
T − 1 global iterations.
Input : x
Initialize : ∀ Mi =1 w 0 ,
i = x 0 , ∀
Mi =1 ∀ N i j
=1 θ i,j = w 0 ,
i , T , K , L , α i,j , β , γ , λ , η i
Output : ∀ Mi =1 ∀ N i j
=1 θ T,K,Li,j , x T 1: for t = 0 , 1 ,...,T do (cid:46)
Global iterations 2: global server sends x t to the teams.
(cid:46)
Global model 3: for k = 0 , 1 ,...,K do (cid:46)
Team iterations 4:
Teams send w t,k to the devices.
(cid:46) Team-level personalized model 5: for l = 0 , 1 ,...,L do .
(cid:46)
Local iterations 6: θ t,k,l +1 i,j = θ t,k,li,j − α i,j ∇ f i,j ( θ t,k,li,j ) − α i,j λ ( θ t,k,li,j − w t,ki )
(cid:46) Personalized local models 7: end for 8: ¯ θ t,ki = (cid:80) N i j =1 δ j θ t,k,Li,j (cid:46) Aggregation within a team 9: w t,k
+1 i =
(1 − η i λ − η i γ ) w t,ki + η i γx t + λη i ¯ θ t,ki (cid:46) Personalized team update 10: end for 11: ¯ w t = (cid:80) Mi
=1 τ i w t,Ki (cid:46)
Global aggregation 12: x t +1 =
(1 − βγ ) x t + β ¯
[Table]
Experiments
We study a classiﬁcation problem to validate pFedMT using both benchmark (MNIST [19]) and synthetic datasets.
Both datasets contain ten class labels.
We divided MNIST in a non-iid manner into two teams, each containing 10 devices.
We adopt the non-iid data generation and distribution procedure from [2] for synthetic data preparation.
Similarly, the generated synthetic data is distributed into two teams, each containing 25 devices.
We took into account an l 2 -regularized multinomial logistic regression (MLR) with softmax activation loss function for µ -strongly convex settings.
A two-layered deep neural network for synthetic dataset and a convolutional neural network is built for
MNIST dataset as non-convex settings.
All datasets are split into train and test sets with a ratio of 3:1.
For the entire experiments we used the learning rate α = 0 .
For pFedMT, we set η = 0 .
03 , and γ = 0 .
1 for convex settings, and η = 0 .
03 and γ = 1 for non-convex settiings.
The value of λ , γ , β is given in Table 1 for the respective experiments.
We consider the value of T = 100 and 600, K =and L = 20 for pFedMT.
The personalized models and global model of pFedMT are represented by pFedMT(PM) and pFedMT(GM) respectively.
Similarly, pFedMe(PM) and pFedMe(GM) denotes personalized and global model of pFedMe, respectively.
4.1 Results and Analysis
We analyzed the communication efﬁciency and convergence of pFedMT with the state-of-the-art (SOTA) models.
From Figure 2, we observed convergence of the personalized model pFedMT(PM) and global model pFedMT(GM) are quicker than the personalized pFedMe(PM) and global model pFedMe(GM) of pFedMe, and Per-FedAvg, respectively.
pFedMT(PM) converges within fewer than 100 global iterations.
The global model needs more iterations to converge.
[Figure]
From Table 1, we observed both convex and non-convex cases.
The pFedMT(PM) outperforms the SOTA.
pFedMT(GM) is better than the FedAvg on both MNIST and synthetic datasets.
For synthetic datasets, pFedMT(GM) performs better than per-FedAvg and pFedMe.
In convex settings, pFedMT(PM) reached 98.30% test accuracy within 100 global iterations.
The closest one is PerFedAvg which reaches 96.26% after 600 global iterations.
On synthetic datasets after training ofglobal iterations, pFedMT(GM) produces slightly better accuracy than pFedMT(PM).
In non-convex settings, the test accuracy of pFedMT(PM) on MNIST datasets reaches 99 .
71% , which is 0 .
38% higher than the pFedMe for 600 global iterations.
On synthetic datasets using deep neural network (DNN), pFedMT(PM) and pFedMT(GM) produce, 95.42% and 94.05% test accuracy, respectively, after training for 600 iterations.
Whereas pFedMe(PM) can produce, 89.53% accuracy on the similar settings.
From the results, we can say that empirically pFedMT(PM) outperforms the SOTA.
[Table]
4.2 Conclusions
We proposed pFedMT, a personalized multi-tier federated learning system that consists of a global server, teams and devices.
We applied the Moreau envelope as a regularized loss function on devices and teams.
pFedMT produces device-speciﬁc personalized models and a global model simultaneously.
From the empirical results, we observed that pFedMT is communication efﬁcient and well generalized model across devices.
This work is still in its early stage.
Therefore, in the near future, we will provide theoretical guarantees of pFedMT, empirical analysis of the effect of hyperparameters λ and γ , and extensive experiments with state-of-the-art models
Acknowledgments and Disclosure of Funding
This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by Knut and Alice Wallenberg Foundation.
The computations were enabled by the supercomputing resource Berzelius provided by National Supercomputer Centre at Linköping
University and the Knut and Alice Wallenberg foundation.