D ISENTANGLED
K NOWLEDGE
T RANSFER :
A
N EW
P ERSPECTIVE
FOR
P ERSONALIZED
F EDERATED
L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Personalized federated learning (pFL) is to collaboratively train non-identical machine learning models for different clients to adapt to their heterogeneously distributed datasets.
State-of-the-art pFL approaches pay much attention on exploiting clients’ inter-similarities to facilitate the collaborative learning process, meanwhile, can barely escape from the irrelevant knowledge pooling that is inevitable during the aggregation phase (e.g., inconsistent classes among clients), and thus hindering the optimization convergence and degrading the personalization performance.
To tackle such conflicts between facilitating collaboration and promoting personalization, we propose a novel pFL framework, dubbed pFedC, to disentangle the global aggregated knowledge into several compositional branches and only aggregate relevant branches for supporting conflicts-aware collaboration among contradictory clients.
Specifically, by reconstructing each local model into a shared feature extractor and multiple disentangled task-specific classifiers, the training on each client transforms into a mutually reinforced and relatively independent multi-task learning process, which provides a new perspective for pFL.
Besides, we conduct a personalized aggregation mechanism on disentangled classifiers via quantifying the combination weights for each client to capture clients’ common prior, as well as mitigate potential conflicts from the divergent knowledge caused by the heterogeneous data.
Extensive empirical experiments are conducted over various models and datasets to verify the effectiveness and superior performance of the proposed algorithm.
I NTRODUCTION
Recently, Federated Learning (FL) has gained growing attention for its capability of collaboratively training machine learning models among distributed clients without accumulating their private data in a central repository McMahan et al.
(2017b)
; Yang et al.
FL has successfully contributed to a wide range of smart applications, such as the next-word prediction Hard et al.
(2018), voice recognition Sattler et al.
(2020), and health care applications
Xu et al.
However, there exists a fundamental challenge that prevents its further application to users with data that having statistical defects, i.e., the data could be non-independent and identically distributed (non-IID).
In such circumstance, the knowledge learnt from local training may have conflicts in the aggregation process, and thus sharing a global model for all clients may significantly slow down the training convergence process and seriously degrade the inference performance Jiang et al.
To deal with such data heterogeneity issue, personalized federated learning (pFL) has emerged that allows FL clients to train personalized models rather than an identical sharing model with others
Hanzely et al.
Tan et al.
; Huang et al.
The principle is to explore the interuser similarities from the joint training process meanwhile producing unique models that can adapt to their diverse data distributions Li et al.
(2021a);
T Dinh et al.
Fallah et al.
; Collins et al.
; Liang et al.
; Li et al.
(2021b)
; Zhang et al.
(2021b)
; Huang et al.
; Zhang et al.
(2021a).
Among the statue quo pFL methods, two of the best popular solutions are similaritybased multiple-model scheme Zhang et al.
(2021b)
; Huang et al.
; Zhang et al.
(2021a)
and parameter decomposition way Collins et al.
; Liang et al.
; Li et al.
(2021b).
The former one is to produce a personalized model for each client by dissecting the similarity among different
[Figure] [Figure]
clients, while the latter one focuses on decoupling the whole model into two parts, top layers (e
[Figure]
g., feature extractor) and bottom layers (e
[Figure]
g., classifier), and only aggregating specific parameters (e.g., feature extractor) in the server.
Existing pFL methods have demonstrated improvements on FL performance via facilitating the collaboration among clients with similar data distributions, while very few literature have paid attention to the contradictory knowledge from users with irrelevant data (i.e., users have different classes or tasks), which may potentially degrade the aggregation efficiency if their cooperation is reinforced during the pFL process.
For example, considering the aggregation phase on two clients with heterogeneous data distributions (i.e., Figure 1), the knowledge transferred from client 1 to 2 can be simply divided into two types: relevant knowledge (i.e., information on class 0) and irrelevant knowledge (i.e., information on class 2, 3).
Traditional client-level aggregation mechanisms in pFL take knowledge from all other clients as a whole, and apply a unified combination weight for both the relevant and irrelevant parts.
Worse still, the intertwined irrelevant knowledge transfer obtained from other clients would be aggravated among users with higher data similarity.
Namely, facilitating personalized collaboration cannot avoid irrelevant knowledge pooling.
As shown in Figure 1, by using different classifier aggregation mechanisms, the average model accuracy on customized classifiers outperforms that of a shared classifier, while the disentangled classifier has the best performance.
This inspires us to investigate a more efficient approach to disentangle the coupled relationship between relevant and irrelevant knowledge transfer during the pFL aggregation process.
Motivated by the multi-task learning (MTL) paradigm, commonalities and differences across multiple related tasks can be exploited simultaneously and further reused to improve the learning performance for task-specific models Hashimoto et al.
; Ruder (2017); Sener & Koltun (2018).
We seek to develop a novel pFL training framework that can circumvent the aggregation conflicts of the irrelevant knowledge derived from the distributed non-IID data during pFL training.
To this end, we formulate the learning phase in pFL to an MTL-based conflicts-aware knowledge transfer training algorithm, dubbed pFedC, whose main idea is to disentangle the global aggregated knowledge into several compositional branches that can be trained in a mutually reinforced or relatively independent way according to their knowledge conflicts.
Specifically, we disentangle the original model in a brand-new perspective and reconstruct the parameters into two types, i.e., the global representation parameters shared by all clients and multiple task-specific parameters owned by each client.
The shared representation is to capture the generic knowledge across all clients, while the task-specific parameters are to optimize the personalization.
Different from previous pFL methods, e.g., FedFomo Zhang et al.
(2021b)
and FedRep Collins et al.
(2021), pFedC disentangles the classifier at the semantic dimension to further improve the FL personalization.
We demonstrate that pFedC can let clients selectively acquire relevant knowledge from others while excluding irrelevant knowledge transfer to avoid potential conflicts.
The pro- posed pFedC can even deal with the extreme case of user heterogeneity, e.g., each client needs to solve different classification tasks.
Extensive experiments are conducted and we show that pFedC can outperform the state-of-the-art baselines over widely used models and datasets, i.e., EMNIST,
FashionMNIST, CIFAR10 and CIFAR100.
The contributions of the paper are summarized as follows:
• To the best of our knowledge, we are the first to consider the conflicts between collaboration and personalization among different clients during the pFL training, and explicitly reveal the benefits of classifier disentanglement in pFL with non-IID data distributions.
• A flexible classifier disengagement mechanism based on semantic knowledge from different classes is proposed to avoid conflicts in knowledge transfer for each client.
• We further design a conflict-aware training framework, pFedC, that can efficiently exploit the inter-similarities among all clients while producing accurate personalized models for each client.
• We conduct extensive experiments on four typical image classification tasks and two nonIID data settings.
The empirical evaluation shows the superior performance of pFedC over the state-of-the-art approaches.
R ELATED W ORK
Personalized Federated Learning.
Existing works on pFL can be categorized into data-based and model-based methods.
The former ones focus on reducing the statistical heterogeneity among clients’ datasets to boost the model convergence, i.e., data-sharing Zhao et al.
(2018), dataaugmentation Jeong et al.
; Duan et al.
Model-based methods emphasize on producing customized model structures or parameters for different clients, which can also be divided into two types: single-model and multiple-model methods.
Single-model based methods extended from the conventional FL algorithms like FedAvg McMahan et al.
(2017a) combine the optimization of the local models and global model, including local fine-tuning Wang et al.
; Schneider & Vlachos (2019); Arivazhagan et al. (2019), regularization T Dinh et al. Hanzely & Richt´arik (2020); Hanzely et al.
model mixture Mansour et al. ; Deng et al. (2020), meta-learning Fallah et al. ; Jiang et al.
and parameter decomposition Bui et al.
; Collins et al.
Arivazhagan et al.
Oh et al. (2022).
Considering the diversity and inherent relationship of local data, a multi-model-based approach is applied to train multiple global models for heterogeneous clients Huang et al.
; Smith et al.
(2017a); Ghosh et al.
; Mansour et al.
Huang et al.
; Zhang et al.
(2021b).
The main idea is to facilitate the collaboration among different clients with similar data distributions.
Existing pFL researches focus on designing efficient methods to facilitate collaboration among clients by measuring inter-similarities, yet none of them have considered the conflict problem caused by irrelevant knowledge aggregated from other clients, especially for those who have higher similarities among their local data distributions.
To this end, we need to explore a novel pFL method to efficiently and selectively capture the relevant knowledge and circumvent the irrelevant knowledge transfer from other clients.
Multi-task Learning.
Multi-task learning (MTL) refers to learning a single model that can solve multiple different tasks Hashimoto et al.
; Ruder (2017); Sener & Koltun (2018).
By sharing parameters between all tasks, MTL can learn more efficiently with an overall model with a smaller size compared with traditional separate model learning methods Zhang & Yang (2021); Vandenhende et al. ; Chen et al. Liu et al. ; Kendall et al.
Multi-class learning can be seen as a special case of MTL, in which the classification on each class can be considered as one task Li et al. ; Huang et al. Smith et al.
Smith et al. (2017a) are the first to apply multi-task learning to the FL settings. Unfortunately, the proposed MOCHA algorithm can only train linear models. FedEM Marfoq et al.
(2021) learns personalized models by implicitly clustering all clients into M groups. However, all above federated MTL approaches treat each client as a single task, neglecting the inherent MTL attribute (i.e., multiple classes) of each client.
Summary.
Aiming to eliminate the conflict between collaboration and personalization during pFL training, we are motivated to reformulate the original pFL task on each client into an MTL problem to improve the learning efficiency meanwhile mitigating the aggregation conflicts from heterogeneous clients.
By such means, the global aggregated knowledge can be disentangled into multiple complementary branches, so as to help optimize both the local training and global aggregation.
P ROBLEM F ORMULATION 3. N OTATIONS AND P RELIMINARY
In pFL, the goal is to collaboratively train personalized models among multiple clients to adapt to the distributed heterogeneous data.
Considering N clients with non-IID datasets, let D i
=
{ x ij , y ij } D
i i = ( 1 ≤ i ≤ N ) be the dataset on the i -th client, where x j ∈ X is the j -th input data sample, y j ∈ Y is the corresponding label.
X ∈ R d and Y = { 1 , 2 , . .
. , C } ∈ R represent the input and output space, respectively.
We define Π i =
[ π 1 i , π 2 i , . . .
, π Ci ] as the label distribution on i -th client, in which each element π ci is a binary indicator that represents whether label c exits in i -th client (i.e., π ci = 1 ) or not (i.e., π ci = 0 ).
Let Π
= { Π 1 , Π 2 , . .
, Π N } denotes the whole label distributions on all clients.
The size of the datasets on the i -th client is denoted by D i .
The size of all clients’ datasets is D = (cid:80) Ni
=
1 D i .
The model on i -th client is denoted by q i : X → Y , which maps input to predicted labels, and is parameterized by θ i .
The objective of pFL can be formulated as
[Equation]
where L i ( θ i ) =
1 D i (cid:80) D i j =1 ℓ ( θ i ; x ij , y ij ) and Θ
= { θ i , . . .
, θ N } is the set of personalized parameters for all clients.
L i is loss function of i -th client associated with dataset D i .
The difference between the predicted value and the true label of data samples is measured by ℓ , e.g., cross-entropy loss
. 3. T ASKWISE M ODEL D ISENTANGLEMENT
We reformulate the original pFL problem into an MTL process to avoid potential aggregation conflicts.
Inspired by recent works on disentangled representations Luo et al. ; Sarhan et al. (2020), we first re-construct the model parameters into two different types of parameters: shared representation ( ϕ ) which denote the global and common knowledge for all tasks, and disentangled task-specific parameters ( ω ) that capture the local conflicts-aware knowledge for each client.
The latter one can be further disentangled into T parts: ω ≜ { ω 1 , . .
. , ω T } .
Note that the number of virtual tasks T can be set flexibly, i.e., T ≤ C .
When T < C , it means a C -class classification task is grouped into T virtual sub-tasks, in which each virtual task becomes a new classification task.
For simplicity, in the formulation part, we disentangle the task-specific parameters according to the number of classes, which is reasonable in practical applications.
In other words, we disentangle the original task into C (i.e., T = C ) binary classification tasks, in which the output space for task c covert to Y c = { 0 , 1 } .
Then the whole local model at i -th client can be represented as
[Equation]
Given the MTL-based version of original model, we then reformulate the whole optimization problem.
Referring to existing hypothesis and loss functions definitions in existing MTL literatureSener & Koltun (2018), we apply the typical minimization formulation at each client i as follows:
[Equation]
where L c ( ϕ i , ω ci ) =
1 D ci (cid:80) D ci j =1 ℓ ( ϕ i , ω ci ; x ij , ˆ y cj ) .
λ c is the weight for task c , which can be static or dynamically computed during the training process.
D ci denotes the number of data samples belonging to the class c .
Since we reformulate the local objective in an MTL version, the j -th data sample on i -th client { x ij , y ij } can be re-represented as { x ij , ˆ y i, 1 j , . . .
, ˆ y i,C j } . We omit the symbol i in ˆ y i,cj in the later content for simplicity. For each client i , the local training becomes a multi-task learning process. To achieve conflictsaware pFL, we should design an efficient MTL-based training and aggregation mechanism that can facilitate collaboration among clients with similar data distribution meanwhile reducing the negative/irrelevant knowledge transfer from other clients.
Require: {D 1 , D 2 , . . .
, D N } , η , E , { Π i } Ni =1 , T .
Ensure: Trained personalized models { θ 1 , θ 2 , . . .
, θ N } .
1: Initialize the clients’ model parameters 2: procedure S ERVER EXECUTES 3: for each communication round t ∈ { 1 , . .
, T } do 4:
Receive { θ i } Ni =1 from all the clients 5:
Calculate ϕ via Eq.
(6) 6: for each client i in parallel do 7:
Calculate ω i via Eq.
Distributed ϕ and { ω i } Ni =1 accordingly 9: for each client i in parallel do 10: θ i ← ϕ ◦ ω i 11: θ i ← CL IENT U PDATE ( θ i ) return Trained personalized models { θ 1 , . .
, θ N } 12: procedure C LIENT U PDATE ( θ i ) 13:
Client i receives θ i from the server.
14: for each local epoch e ∈ { 1 , . . .
, E } do 15: for mini-batch ξ t ⊆
D i do 16:
Update model parameters θ i via Eq.
(5) 17: return θ i
P F ED C A LGORITHM
To conduct efficient pFL via the MTL training, we update the shared representation and task-specific parameters at the server side.
We apply a conflict-aware personalized federated learning framework to select only the relevant knowledge from other clients during the aggregation phase.
4. C ONFLICTS - AWARE P ERSONALIZED F EDERATED L EARNING
In this section, we elaborate the conflicts-aware model training and aggregation mechanism in pFedC step by step.
Since the shared representation should capture the generic and common knowledge from all data samples, we let each client to maintain a global version of shared representation at the server side, and download it to local side to form a new local model before each communication round, i.e., ϕ ← (cid:80) Ni
=1 D i D ϕ i .
As for task-specific parameters, we introduce a mask vector m i at the server side to guide the aggregation phase of i -th client:
[Equation]
where m ci represents the aggregation weight vector of c -th task in i -th client, while each element m ci,n denotes the aggregation decision of task c on n -th client for i -th client, i.e., m ci,n = 1 if and only if client i and n both have class c ; otherwise m ci,n = 0 .
Therefore, the value of m ci,n can be easily calculated by m ci,n = π ci · π cn ( 1 ≤ i, n ≤ N ).
In aggregation phase, each client i calculates the personalized task-specific parameters
ω
i by an element-wise multiplication of m i and all local { ω n } Nn =1 .
Then, the server can update the personalized model consisting of the shared representation and the task-specific parameters for each client.
4. W ORKFLOW
As discussed earlier, pFedC enables each client to perform an MTL locally, and the server aggregates the local model parameters in a conflict-aware manner.
To do so, pFedC alternates between client updates and a server update during each communication round.
The overall learning scheme is summarized in Algorithm 1.
Client Update.
In each communication round, the clients download the shared representation ϕ and personalized task-specific parameters
ω i from the server.
In the client update phase, each of them makes several local gradient-based updates (e.g., E epochs) to optimize the objective in Eq.
(3) by adapting the MGDA-UB algorithm Sener & Koltun (2018) from MTL.
Namely, for each iteration, client i updates its model as follows,
[Equation]
) where θ i = { ϕ ◦ ω i } , and η denotes the learning rate.
∇ θ i L i ( θ i ) is the gradient of local objective function L i
( θ i ) with respect to θ i .
Server Update.
After training the local models at the client-side, corresponding models will be sent to the server for further aggregation.
Specifically, the server first aggregates the shared representation by taking a weighted averaged of them:
[Equation]
The personalized task-specific parameters for each task c of client i can be aggregated by using mask vector m i , i.e.,
[Equation]
Updated shared representation and personalized task-specific parameters are then distributed to the local clients for next round of training.
4. A NALYSIS ON P F ED C
Convergence Analysis.
We prove that pFedC can assist clients converge to their respective local optimums under the following assumptions: 1) L 1 , · · · , L N are all µ -strongly-convex, 2) L 1 , · · · , L N are all L -smooth, 3) the variance of stochastic gradients in each user is bounded by σ 2 i and 4) the expected squared norm of stochastic gradients is uniformly bounded by G 2 .
Theorem 1.
Let all above assumptions hold and µ, L, σ i , G are defined therein.
Choose κ = Lµ , γ = max { 8 κ, E } and the learning rate η t = [Equation] ( γ + t ) .
Suppose task c exists in client
i (i.e., π ci = 1 ), then client i in pFedC satisfies
[Equation]
where
B =
[Equation]
=
[Equation]
The full convergence analysis proof is elaborated in Appendix 7.3.
From Theorem 1, we show that pFedC converges to the respective local optimums at a rate of O ( 1 T ) for strongly convex and smooth functions.
Privacy Preserving. In pFedC, the server requires to access { Π i } Ni =1 from local clients, which may lead to privacy concerns on user data profiles.
However, there are already some literature Tan et al.
; Zhu et al.
(2021) show that transmitting local category-level information is considered reasonable and will not cause privacy leakage in image classification tasks, where different clients have different subsets of the whole classes.
Thus, it is feasible to upload data distribution information to boost personalization performance in a federated framework.
Moreover, pFedC can be integrated with various privacy-preserving techniques, such as Homomorphic Encryption (HE) Phong et al.
(2018), differential privacy Wei et al.
and
even trusted execution environment Mondal et al.
(2021) to enhance the reliability of the training framework.
Taking Homomorphic Encryption as an example, we design a privacy-preserved version of pFedC.
As space is limited, we omit it in the main content and explain the details in Appendix 7.4.
Communication Efficiency.
Compared with conventional pFL methods that only transmit local model parameters or partial parameters, e.g., top layers in FedRep Collins et al.
(2021), our proposed pFedC requires the exchange of multi-branches model parameters, which may lead to additional communication overhead.
However, we clarify that the additional communication cost on pFedC [Table] can be negligible, which only has about 0.011% increments.
We will conduct detailed explanations on communication efficiency of pFedC in the experimental part
[Table]  E XPERIMENTS 5. E XPERIMENTAL S ETUP
Datasets.
We evaluate the pFedC framework over four datasets, EMNIST, FashionMNIST, CIFAR10 and CIFAR100.
We consider two non-IID scenarios: 1) pathological non-IID setting — each client is randomly assigned two classes with the same amount of data on each class McMahan et al.
(2017b)
; 2) real-world non-IID setting — each client contains all classes, while the data on each class is not uniformly distributed, i.e., using a Dirichlet distribution Dir ( α ) Hsu et al.
Lin et al.
(2020), where α indicates the non-IID level, and a smaller value of α indicates higher data heterogeneity.
All data are divided into 75% training set, and 25% test set.
The test set and the training set have the same data distribution for all clients.
Baselines.
We compared the performance of pFedC with the state-of-the-art methods.
In addition to
FedAvg McMahan et al.
(2017b)
, we also include Per-Fedavg Fallah et al.
(2020), a pFL algorithm based on meta-learning; pFedMe T Dinh et al.
(2020), a pFL algorithm with regularization term added in the objective function;
FedBN Li et al.
(2021b), keeps each client’s BN layer updating locally, while other layers are aggregated according to the FedAvg algorithm; FedRep Collins et al.
(2021), a pFL algorithm that keeps each client’s classifier updating locally, while the other parts are aggregated at the server; FedFomo Zhang et al.
(2021b), a pFL algorithm that uses distance to calculate the aggregation weights based on the model and loss differences.
Training Details.
For MNIST and EMNIST dataset, we use a CNN architecture based on LeNet
LeCun et al.
We treat all layers except the last as the shared representation function and put ten fully-connected layers as task-specific functions.
For CIFAR10 and CIFAR100 dataset, we use an architecture based on ResNet-18 He et al.
and make similar modifications to it as LeNet.
The detailed architectures used for experiments is visualized in Appendix 7.2.
It’s worth noting that we disentangle the model of CIFAR100 based on the number of superclasses (i.e., 20).
In all baselines, we use the original LeNet and RestNet-18 architectures.
All the models have the same structure between different clients under the same setting.
We evaluate the performance of pFedC on 20 clients with 100% participation, and 100 clients with 10% participation, respectively.
The average model accuracy of all clients is obtained after 1000 communication rounds training for all baselines and 40 communication rounds training for pFedC, respectively.
Implementation.
We simulate all clients and the server on a workstation with an RTX 2080Ti GPU, a 3 . 6 -GHZ Intel Core i9-9900KF CPU and 64 GB of RAM.
5. E XPERIMENTAL R ESULTS
Performance Comparison.
We compare the performance of pFedC with all baselines under two different non-IID cases: pathological non-IID setting and real-world non-IID setting.
For all experiments, we use SGD optimizer with a batch size of 32 for MNIST and EMNIST, 16 for CIFAR and CIFAR100.
The learning rate is 0 .
01 for CIFAR10 and CIFAR100, and 0 .
005 for MNIST and
EMNIST.
The performance of all baselines is evaluated by the mean testing accuracy, which is the [Table] [Table] average of the testing accuracy on all clients and the ± indicates the error range of the mean testing accuracy after 3 repeated experiments.
For pFedC, we evaluate the performance by calculating the mean testing accuracy on all tasks.
During the training process, MGDA-UB Sener & Koltun (2018) is used to solve the local optimization problem on each client.
The experimental results are listed in
Table 1 and Table 2, respectively.
Table 1 demonstrates the mean testing accuracy of all methods on four datatsets with the pathological non-IID data setting.
It can be observed that pFedC constantly outperforms other competitors by large margins under a variety of settings.
The reason mainly lies in the MTL-based local training and personalized aggregation of the pFL models, which encourages the clients to optimize the objective in a conflict-aware manner.
For real-world non-IID setting, Table 2 shows the results of FedAvg, PerFedAvg, pFedMe, FedBN, FedRep, FedFomo and our pFedC on four datasets under three different level of data heterogeneity (i.e., α = 0 .
From Table 2, we can notice that pFedC provides superior performance than the baselines over different datatsets and data distributions in most cases.
Besides, we emphasize that even if all clients own all types of classes (e.g., α = 1 ), the performance of pFedC is still better than that of FedAvg.
This is enabled by the disentangled architecture and task-specific aggregation process.
Effect of Data Heterogeneity.
To further verify the efficiency of our proposed pFedC, we investigate the effect of data heterogeneity on training performance.
The final mean test accuracies of each algorithm under different value of α = 0 .
1 , 1 are recorded to plot the trends between performance and data heterogeneity.
As shown in Figure 2, pFedC can provide stable gain over other methods when the data distributions are highly heterogeneous (i.e., with smaller α ).
In most cases, pFedC shows an obvious advantage in training performance, while the advantage of Per-FedAvg,
FedRep, FedFomo vanishes as data heterogeneity alleviates.
These results show the robustness of pFedC under different levels of data heterogeneity.
Scalability.
To verify the scalability of pFedC, we further conduct experiments on large-scale FL systems, where the number of clients reaches 100.
During the training process of pFedC, partial clients (sampling rate: 0.1) will be selected and related partial elements in the mask matrix will be [-= 0. = 0. = 
-]used for aggregation in each communication round.
[-= 
-]
In Table 3, the experimental results demonstrate the efficiency and scalability of pFedC. Besides, we notice that disentangling the original bottom layer (i.e., classifier) into T parts in pFedC only expands the width in the last FC layer while the depth has not changed at all.
[-= 0. = 
 A v e r a g e T e s t A cc u r ac y ( % ) pFedC, E =3 pFedC, E =5 pFedC, E = pFedC, E = (c) CIFAR10 Dataset.-]Taking ResNet-18 on CIFAR-100 for example, one FC layer (50 × 1) is disentangled and reconstructed into twenty FC layers (50 × 1) for pFedC (in our experimental setting).
[-= 0. = 0. = 
 A v e r a g e T e s t-]The parameter size on the original model and our proposed model are 8,869,427 and
[Figure]
[Table]
[-= [Table] [Table] = 
-]8,870,396, respectively, which has only 0.011% increments.
Compared with the original model, the additional cost on memory, computation and communication are negligible.
Comparison with Federated MTL methods.
To further demonstrate the efficiency of MTL process in pFedC, we compare it with MOCHA Smith et al. (2017b) and FedEM Marfoq et al.
Since MOCHA only works for linear models, we conduct experiments on two different models: linear models and CNN models (i.e., ReNet-18). We tuned the parameter λ of MOCHA on a holdout validation set via grid search in { 10 1 , 10 0 , 10 1 , 10 2 , 10 3 } , and we found that the optimal value of λ is 10 0 . Unless otherwise stated, the number of components considered by FedEM is M = 3 .
We show that pFedC can obtain the best performance when compared with other federated MTL methods.
The detailed results are shown in Table 6.
Effect of Different Disentanglement Strategies.
For large-scale classification datasets, disentangling the original model at a class level may be infeasible due to the huge number of classes in the dataset.
In this case, how to disentangle the model becomes important and should be carefully investigated.
Therefore, We conduct analysis on the effect of disentanglement strategies on CIFAR100.
Three different disentanglement strategies are proposed in the experiments: 1) T = 10 , splitting the 100 classes into 10 groups and considering each 10 consecutive classes as a task von Oswald et al.
(2019); 2) T = 20 , splitting the dataset according to the superclasses (i.e., 20 superclasses) Yoon et al.
(2019); 3) T = 100 , splitting the dataset according to the original classes (i.e., 100 classes).
T = 1 represents no model disentanglement in the training process, which is equivalent to FedAvg.
As shown in Table 4, the training performance on CIFAR100 with different data heterogeneity varies greatly with the number of tasks (i.e., T ).
Specifically, comparing with FedAvg (i.e., T = 1 ), pFedC achieves a significant improvement on model accuracy, e.g., a 32.13% - 59.31% improvement over
FedAvg when T = 10 , and 4.01% - 24.49% improvement over FedAvg when T = 100 .
These results verify the efficiency of our proposed model disentanglement-based pFL training framework.
However, larger values of T cannot always benefit the convergence of the personalized models.
The reason is that a large number of tasks makes it extremely difficult to learn the relationship between these different tasks.
Moreover, larger values of T may bring disadvantages to the whole training process due to the huge increase in model size and local training time.
Thus, a proper setting of T is necessary and there are many things to take into consideration, e.g., the computation overhead, communication overhead, and model performance.
C ONCLUSION
In this paper, we have investigated the potential conflicts caused by the irrelevant knowledge transfer from heterogeneous clients in pFL, and integrated a novel model disentanglement method (pFedC) into the pFL training framework, which can improve the training efficiency and model performance during the learning process over non-IID data.
It is shown that pFedC can capture the relevant knowledge and circumvent the irrelevant knowledge transfer from other clients.
Extensive empirical experiments have been conducted over various models and datasets to verify the effectiveness and superior performance of pFedC framework.
A PPENDIX 7. C OMPARISON WITH OTHER P FL METHODS
[Figure]
(a)
[Figure]
𝜃" "
Client𝜙 (b) FedRep 𝜙 𝜔 !
Client𝜙 𝜔 "
Client𝜙 𝜙 (c) pFedC (Ours)
Client
Client𝜔 !!
𝜔 !" 𝜔 !# … 𝜔 "!
𝜔 "" 𝜔 "# … 𝜙 𝜔# !!
𝜔# !"
𝜔# !
# … 𝜔# "!
𝜔# "" 𝜔# "# …
[Figure]
As shown in Figure 3, different from previous pFL methods, e.g., FedFomo Zhang et al. (2021b) and
FedRep Collins et al.
(2021), pFedC disentangles the classifier at the semantic dimension to further improve the FL personalization.
We demonstrate that pFedC can let clients selectively acquire relevant knowledge from others while excluding irrelevant knowledge transfer to avoid potential conflicts.
The proposed pFedC can even deal with the extreme case of user heterogeneity, e.g., each client needs to solve different classification tasks.
7. A DDITIONAL E XPERIMENTAL D ETAILS
In this section, we provide more concrete experimental settings and more numerical results to examine the advantages of our proposed pFedC, i.e., the convergence rate.
7.2. N ON -IID D ATA S ETTING
Pathological non-IID setting: each client is randomly assigned two classes with the same amount of data on each class.
Real-world non-IID setting
: each client contains all classes, while the data on each class is not uniformly distributed, i.e., using a Dirichlet distribution Dir ( α ), where α indicates the non-IID level, and a smaller value of α indicates higher data heterogeneity.
We visualize the effects of adopting different α on the data heterogeneity for the MNIST dataset in
Figure 4.
7.2. D ETAILED A RCHITECTURES USED FOR P F ED C
For MNIST and EMNIST dataset, we use a CNN architecture based on LeNet.
We treat all layers except the last as the shared representation function and put ten fully-connected layers as taskspecific functions.
For CIFAR10 and CIFAR100 dataset, we use an architecture based on ResNet18.
Similarity, all layers except the last one are treated as shared representation for all tasks, and the fully connected layer are used as task-specific functions.
The architectures used for different datasets are visualized in Figure.
7.2. T HE T RAINING P ERFORMANCE ON D IFFERENT N ON -IID S ETTINGS
We compare pFedC with state-of-the-art baselines on four datasets: MNIST, EMNIST, CIFAR and CIFAR100.
There are two non-IID data settings: 1) Pathological non-IID setting, where the
[Figure]
data on each client only contains the specific number of labels (maybe only two labels); 2) Realworld non-IID setting, where the number of for client randomly chosen
Figure 6, 7, 8, 9 shows the empirical convergence results of pFedC along with other baselines.
Specifically, we focus on the changes of average test accuracy of these algorithms in each communication round.
It is obvious that pFedC converges to higher average test accuracy on all four datasets than baselines.
This phenomenon validate the effectiveness of the model decomposition based training and aggregation on pFL.
Moreover, pFedC can achieve faster convergence speed than baselines in most cases, which further verifies the advantages of model decomposition and conflict-aware personalized aggregation on training performance. pFedC) on CIFAR-100 with different T . Setting classes per client = α = 0 .α
{+= 0 .α = T =244 |250 |200 |
 |+}
{+T =244 |250 |200 |
+}
|
{+T =244 |250 |200 |
+}
|
{+T =244 |250 |200 |
+}
| 7.2. T HE R EQUIRED C OMMUNICATION R OUNDS ON D IFFERENT D ECOMPOSITION S TRATEGIES
Finally, we conduct experiments to analyze the convergence rate.
From table 5, the required communication rounds to reach the convergence in pFedC is far less than baseline methods.
C onv ( K e r n a l: 5 ,
F ilt e r : 10 )
R e L U
M a x P oo l ( 2x2 )
C onv ( K e r n a l: 5 ,
F ilt e r : 20 )
R e L U
M a x P oo l ( 2x2 )
F C L a y e r ( 320x )
FC Layer (50x1) FC Layer (50x1)
FC Layer (50x1) FC Layer (50x1) …
Task 1: (binary)
Task 2: (binary)
Task 3: (binary)
Task C: (binary)
Input
Image28x ( [Figure] MNIST and EMNIST.
FC Layer (50x1) FC Layer (50x1)
FC Layer (50x1) FC Layer (50x1) …
Task 1: (binary)
Task 2: (binary)
Task 3: (binary)
Task C: (binary)
[Figure] 32x32x
ResNet-18: ﻿ ResNet [Figure] [Figure] the
[Figure]
fully connected layer
[Figure]
Linear Model
[Table]
7.2. C OMPARISON WITH MORE P FL M ETHODS
To further demonstrate the efficiency of MTL process in pFedC, we compare it with MOCHA Smith et al.
(2017b)
and
FedEM Marfoq et al.
Since MOCHA only works for linear models, we conduct experiments on two different models: linear models and CNN models (i.e.,
ReNet-18).
We tune the parameter λ of MOCHA on a holdout validation set via grid search in { 10 1 , 10 0 , 10 1 , 10 2 , 10 3 } , and find that the optimal value of λ is 10 0 .
Unless otherwise stated, the number of components considered by FedEM is M = 3 .
From Table 6, we show that pFedC can obtain the best performance when compared with other federated MTL methods.
7.2. E FFECT OF L OCAL E POCHS
Figure 10 demonstrates the performance of pFedC under different settings of local epochs (the local epochs E in each communication round are set to 3, 5, 10 and 20 respectively).
It is observed that larger values of E are not always optimal for obtaining the best performance.
For example, the model can achieve the highest averaged accuracy( e.g., 96.86%) on MNIST with α = 1 when pFedC performs 5 epochs during one communication round, while 10 epochs are better for CIFAR ( α = 1 ).
[Figure]
cc
[Figure]
ac
[Figure]
7. C ONVERGENCE A NALYSIS OF P F ED C
The personalized performance convergence analysis for each client is the same, so we only focus on one client i, 1 ≤ i ≤ N .
Consider that all clients participate in the aggregation phase in pFedC.
7.3. A DDITIONAL N OTATION
Let θ i,t be the model parameters of i -th client at t -th step.
Let E be the number of local updates in each round.
Let I E be the set of synchronization steps, i.e., I E = { nE | n = 1 , 2 , · · · } .
If t +1 ∈ I E , i.e., the step to communication.
Then, the model update with all participants can be described as:
[Equation]
where
[Equation]
and p n = D n D , ϵ c
{+= (cid:80) Nn =1 m cn , v i,t +
+}
=
{ ϕ i,t +1 ◦ ω
i,t +1 } .
Here, an additional variable v i,t +1 is introduced to represent the immediate result of one step SGD update from θ i,t +1 .
We regard θ i,t + as the parameter obtained after communication steps.
In following analysis, we define two virtual sequences ¯v i,t
=
Agg (v i,t | i = 1 , · · · , N ) and ¯ θ i,t =
Agg ( θ i,t | i = 1 , · · · , N ) .
As the whole original model is divided into two different roles in our proposed pFedC: global feature extractor (i.e., ϕ ) and task-specific classifier (i.e., ω ).
According to
Eq.
11, the update policy for ϕ is same as FedAvg algorithm, so we omit the analysis on ϕ and only consider the task-specific parts.
Moreover, for each task c , π c i has two different values, i.e., 0 or
[Figure]
cc
[Figure]
ac
[Figure]
For convenience, we only consider the case with π ci = 1 .
Therefore, in the rest of the analysis, v , ¯v , θ and ¯ θ denote the parameter related to ω c unless other specified.
It is noticed that ¯v i,t +1 comes from an single step of SGD from ¯ θ i,t .
We also define ¯ g ci,t ( ω ci ) =
(cid:80) Nn =1 p n π cn
ϵ c ∇ ω cn L n ( θ n,t ) and g ci,t ( ω ci ) =
(cid:80) Nn =1 p n π cn
ϵ c ∇ ω cn L n ( θ n,t , ξ n,t ) .
For simplicity, we omit the w ci in the above expression of ¯ g ci,t
and g ci,t .
Thus, we have ¯v i,t
{++
+}
= ¯ θ i,t − η t g i,t and E g i,t = ¯ g i,t .
7.3. K EY L EMMAS
To clearly show our proof, it is necessary to define some lemmas before the main theorem.
The proof of these lemmas can be easily found in Li et al.
and we only focus on the main theorem.
Lemma 1.
(Results of one step SGD).
Assume Assumption 1 and 2 hold, we have
[Equation]
where
Γ =
L ∗ i − (cid:80)
Nn =1 p n π cn ϵ c L ∗ n ≥ 0 .
Lemma 2.
(Bounding the variance).
Assume Assumption 3 holds.
It follows that
[Equation]
Lemma 3.
(Bounding the divergence of θ n,t ).
Assume Assumption 4 holds, that η t is non-increasing and η t ≤ 2 η t + E for all t ≥ 0 .
It follows that
[Equation]
4 η
[Equation]
[Figure]
r a
[Figure]
e
[Figure]
t
A cc u r ac
[Figure]
cc
[Figure]
ac
[Figure]
7.3. F ULL P ROOF OF T HEOREM
Proof.
It is clear that no matter whether t + 1 ∈ I E or t + 1 / ∈ I E
[Figure]
we always have the following equation: ¯ θ i,t +1 = ¯v
i,t +1 .
Let ∆ t = E ∥
¯ θ i,t − θ ∗ i ∥ 2 .
From Lemma 1, Lemma 2 and Lemma 3, it follows that
[Equation]
where
[Equation]
=
[Equation]
For a diminishing stepsize, η t = βt + γ for some β > 1 µ and γ > 0 such that η 1 ≤ min { 1 µ , 14 L } and η t ≤ 2 η t + E .
We prove that ∆ t ≤ vγ + t where v = max { β 2 B βµ − 1 , ( γ + 1)∆ 1 } .
First, the definition of v ensures that it holds for t = 1 .
Assume the conclusion holds for some t , it follows that
[Equation]
Then, according to the L -smoothness of L i ( · ) , we have
[Equation]
[Figure]
A v e r a g e
T e s t
A cc u r ac y
[Figure]
v
[Figure]
r a g e
T e s t
A cc u
[Figure]
ac y
[Figure]
[Figure]
u
[Figure]
Specifically, if we choose β = 2 µ , γ = max { 8 Lµ − 1 , E } and denote κ =
Lµ , then η t = 2 µ 1 η + t .
One can verify that the choice of η t ≤ 2 η t + E for t ≥ 1 .
Then we have
[Equation]
(21) and
[Equation]
The proposed pFedC requires the exchange of local label distribution between the server and the clients. This property may bring privacy concerns to the user profile or other sensitive information. In this section, we integrate the homomorphic encryption technique with pFedC to further protect the local data distribution.
Specifically, we propose an evolutionary version of pFedC, dubbed PrivacypFedC, to achieve a privacy-preserving training framework.
Instead of directly uploading shared feature extractor ϕ , multiple task-specific parameters ω i and local label distributions Π i to the server, the results of ω i · Π ∗ i will be encrypted and sent to the server, where Π ∗ i =
[1 , π 1 i , · · · , π C i ] .
The server will aggregate the local updates to obtain the encrypted personalized parameters and then distribute the updated results back to the clients.
Then, the clients can decrypt the updated results to obtain personalized models and then continue with the next round of training.
To this end, the privacy of the training framework can be guaranteed.
[-[Table] T =244 |250 |200 |
 | T =244 |250 |200 |
-]Attackers from the server cannot reconstruct the raw data of clients from aggregated results or encrypted results without access to the private key on each client.
[-| T =244 |250 |200 |
 |-]The detailed workflow of Privacy-pFedC is illustrated in Algorithm 2.
Require: {D 1 , D 2 , . . .
, D N } , η , E , T , { Π ∗ 1 , Π ∗ 1 , · · · , Π ∗ N } , Encryption function Enc ( · ) , Decryption function Dec ( · ) .
Ensure: Trained personalized models { θ 1 , θ 2 , . . .
, θ N } .
1: Generate public key pk and secret keys sk 2: procedure S ERVER EXECUTES 3: for each communication round t ∈ { 1 , . .
, T } do 4:
Receive ϕ i and { (cid:101) ω
i } Ni =1 from the selected clients 5: ϕ ← (cid:80) Ni
=1 D i D
ϕ
i 6: for each task c do 7: (cid:101) ω
c ← 1 (cid:80) Nn =1 m ci,n (cid:80) N n =1 (cid:101) ω c n · m c i,n 8:
Distributed ϕ and { (cid:101)
ω c } Cc =1 to selected clients 9: for each client i in parallel do 10: θ i ← CL IENT U PDATE ( ϕ, { (cid:101)
ω c } Cc =1 ) 11: (cid:101) ω
i ← Enc ( ω i · Π ∗ i , pk ) return Trained personalized models { θ 1 , . .
, θ N } 12: procedure C LIENT U PDATE ( ϕ, { (cid:101)
ω c } Cc =1 ) 13: { ω c } Cc =1 ← Dec ( { (cid:101) ω c } Cc =1 , sk ) 14: for each task c do 15: if π ci =
= 1 then
: ω ci ← ω c 17: else 18: ω ci ← ω ci 19: θ i ← ϕ ◦ ω i 20: for each local epoch e ∈ { 1 , . . .
, E } do 21: for mini-batch ξ t ⊆
D i do 22:
Update model parameters θ i via: θ i ← θ i − η ∇ θ
i L i ( θ i ) 23: return θ i