Title: Fed-CO$_{2}$: Cooperation of Online and Offline Models for Severe Data Heterogeneity in Federated Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Federated Learning

Data Heterogeneity

Model Cooperation

Mutual Learning

Knowledge Transfer
Abstract: Federated Learning (FL) has emerged as a promising distributed learning paradigm that enables multiple clients to learn a global model collaboratively without sharing their private data. However, the effectiveness of FL is highly dependent on the quality of the data that is being used for training. In particular, data heterogeneity issues, such as label distribution skew and feature skew, can significantly impact the performance of FL. Previous studies in FL have primarily focused on addressing label distribution skew data heterogeneity, while only a few recent works have made initial progress in tackling feature skew issues. Notably, these two forms of data heterogeneity have been studied separately and have not been well explored within a unified FL framework. To address this gap, we propose Fed-CO$_2$, a universal FL framework that handles both label distribution skew and feature skew within a Cooperation mechanism between the Online and Offline models. Specifically, the online model learns general knowledge that is shared among all clients, while the offline model is trained locally to learn the specialized knowledge of each individual client. To further enhance model cooperation in the presence of feature shifts, we design an intra-client knowledge transfer mechanism that reinforces mutual learning between the online and offline models, and an inter-client knowledge transfer mechanism to increase the models’ domain generalization ability. Extensive experiments show that our Fed-CO$_2$ outperforms a wide range of existing personalized federated learning algorithms in terms of handling label distribution skew and feature skew, both individually and collectively. The empirical results are supported by our convergence analyses in a simplified setting.

Fed-CO 2 : Cooperation of Online and Offline Models for Severe Data Heterogeneity in Federated Learning

Abstract

Federated Learning (FL) has emerged as a promising distributed learning paradigm that enables multiple clients to learn a global model collaboratively without sharing their private data. However, the effectiveness of FL is highly dependent on the quality of the data that is being used for training. In particular, data heterogeneity issues, such as label distribution skew and feature skew, can significantly impact the performance of FL. Previous studies in FL have primarily focused on addressing label distribution skew data heterogeneity, while only a few recent works have made initial progress in tackling feature skew issues. Notably, these two forms of data heterogeneity have been studied separately and have not been well explored within a unified FL framework. To address this gap, we propose Fed-CO 2 , a universal FL framework that handles both label distribution skew and feature skew within a C ooperation mechanism between the O nline and O ffline models. Specifically, the online model learns general knowledge that is shared among all clients, while the offline model is trained locally to learn the specialized knowledge of each individual client. To further enhance model cooperation in the presence of feature shifts, we design an intra-client knowledge transfer mechanism that reinforces mutual learning between the online and offline models, and an inter-client knowledge transfer mechanism to increase the models' domain generalization ability. Extensive experiments show that our Fed-CO 2 outperforms a wide range of existing personalized federated learning algorithms in terms of handling label distribution skew and feature skew, both individually and collectively. The empirical results are supported by our convergence analyses in a simplified setting.

1 Introduction

Federated Learning (FL) [1, 2] is a distributed learning framework that involves collaboratively training a global model with multiple clients while ensuring privacy protection. The pioneering work FedAvg [2] learns the global model by aggregating the local client models and obtains satisfactory performance when the client data are independently and identically distributed (IID). However, when the data are heterogeneous and non-IID among clients, performance with the global model can degrade substantially [3, 4]. Common data heterogeneity issues include label distribution skew and feature skew. Clients experiencing label distribution skew exhibit varying class distributions within the same domain, while clients encountering feature skew maintain consistent class distributions

Ye Shi ∗

Jingya Wang

but belong to different domains. As a promising solution to data heterogeneity issues, Personalized Federated Learning (PFL) has emerged, where personalized models are trained for individual clients. Most previous studies have focused on addressing label distribution skew [5, 6, 7], while only a few recent works [8, 9] have made initial progress in tackling feature shift issues. So far, these two forms of data heterogeneity have been studied separately and have not been addressed within a unified FL framework. Therefore, in this paper, we propose a universal FL framework that effectively handles data heterogeneity issues arising from label distribution skew, feature skew, or their combination.

Several algorithms personalized some parts of the model in FL with label distribution imbalance or feature shifts to retain and leverage some of the local offline information [10, 11, 12, 13, 8]. However, in realistic FL scenarios, where extreme label distribution skew, severe feature skew, or even both are present, these algorithms fail to effectively harness local specialized knowledge for satisfactory adaptation. In fact, in some cases of extreme heterogeneity, models trained by those personalized algorithms may even perform worse than the locally-trained model, as the locally-trained model excels in capturing offline specialized knowledge. On the other hand, in FL scenarios with milder heterogeneity, partially personalized models trained by PFL algorithms perform better due to their ability to access online general information from other clients.

So, given that the model with partially personalized parameters and the locally trained model each perform better in different cases, the question is: Is there a more effective approach to fuse the online general knowledge and the offline specialized knowledge for better performance? From our investigations, the answer is yes. Accordingly, we propose a novel universal cooperation framework with these two models to address this challenge for both label distribution skew and feature skew data heterogeneity, referring to the model with partially personalized parameters as the online model, and the locally trained model as the offline model. Specifically, we personalize Batch Normalization layers in the online model and fuse the online and offline models' predictions as the final prediction. The prediction fusion between the online and offline models rectifies errors in their respective individual predictions and exhibits better performance.

In FL scenarios with feature skew, the general knowledge learned by the online model is domaininvariant, whereas the specialized knowledge learned by the offline model is domain-specific. Cooperation that occurs solely by fusing predictions is not sufficient as this process does not let the online and offline models communicate during the training phase. In turn, this impedes the transfer of domain-invariant and domain-specific knowledge between the models. Hence, to further encourage cooperation between the models, we propose two novel knowledge transfer mechanisms - one intraclient and one inter-client - which work at the model and client level, respectively. The intra-client knowledge transfer mechanism facilitates mutual learning between the online and offline models via knowledge distillation, which enables the online and offline models to benefit from both online domain-invariant and offline domain-specific knowledge. Conversely, the inter-client knowledge transfer mechanism enhances the model's domain generalization ability by introducing classifiers from the offline models of other clients to each local client.

Contribution. We propose Fed-CO 2 , a universal cooperative FL framework for severe data heterogeneity including both label distribution skew and feature skew. By simply fusing the online and offline models, Fed-CO 2 can handle severe label distribution skew effectively. To enhance model performance in the presence of severe feature skew, Fed-CO 2 involves an intra-client knowledge transfer mechanism that improves model cooperation and an inter-client knowledge transfer mechanism that increases client cooperation. We theoretically show that Fed-CO 2 has a faster convergence rate than FedBN [10] in a simplified setting. Besides, extensive experiments on five benchmark datasets demonstrate that our Fed-CO 2 framework has a prominent edge over a range of state-of-the-art algorithms where the data contain label distribution skew, feature skew, or both 1 .

2 Related Work

Federated Learning for Label Distribution Skew Data Heterogeneity. In real-world FL scenarios, label distribution imbalance among clients is a common phenomenon that poses challenges to learning a single model that can effectively cater to all clients. A straightforward approach to personalizing the global model is fine-tuning it on local datasets [14, 15, 16, 17]. Other approaches attempt to overcome distribution heterogeneity by exerting a proximal regularization term on the global model,

as seen in examples like FedProx [4], pFedMe [18], MOON [7], and Ditto [19]. Instead of adapting the single global model, learning personalized models is another main track. Here, the more explicit methods [11, 20, 21, 22] personalize some of the model parameters while leaving other parameters for aggregation. Similarly, FedMask [23] learns distinct binary masks for the last several layers of the local models. Further, inspired by Hypernetworks [24], some methods [6, 13, 25] apply a hypernetwork to produce client-specific parameters for participants. Beyond directly personalizing parts of the model parameters, some of the newer methods involve a two-branch architecture where the aim is to calibrate the model's predictions [12, 26, 27]. For example, FedRoD [12] uses a personal head to strengthen local learning, while PerFCL [26] splits the features in the local clients into shared parts and personalized parts. FedProc [27] uses prototypes as global knowledge to correct local training. However, unlike these methods, Fed-CO 2 maintains two separate personalized models for each client - one that learns online general knowledge and the other that learns offline specialized knowledge - with a focus on facilitating cooperation between the two models.

Additionally, knowledge distillation, which transfers dark knowledge between models, has played an important role in addressing label distribution skew [28, 29, 30, 31, 32, 22]. Most of previous works in this arena rely heavily on a public dataset to transfer knowledge between the server and the local clients [29, 30, 31]. Only a limited number of studies have delved into the synergistic interplay between learning global knowledge and learning local knowledge, as well as strategies to optimize the mutual benefits of these two processes. In FML [28], the global model and the local model engage in mutual learning, but only the local model is used for personalized predictions. A very recent study, CD 2 -pFed [22], employs cyclical distillation as a regularization term in conjunction with the local training cross-entropy loss to mitigate the gaps between the learned representations from local weights and those from global weights. By contrast, our method utilizes knowledge distillation to facilitate the exchange of beneficial knowledge between the online and offline models before local training, as well as to ensure effective and adequate intra-client knowledge transfer.

Federated Learning for Feature Skew Data Heterogeneity. While various methods have shown promising results on label distribution data heterogeneity, they often suffer from significant performance degradation when confronted with the more challenging feature shift data heterogeneity. To adapt to the feature shifts among different domains, FedBN [10] personalizes Batch Normalization (BN) layers for each participant. Beyond personalizing BN layers, PartialFed [8] explores the strategy of selecting personalized parameters based on the feature characteristics of different clients. To handle the feature skew, FedAP [9] learns the similarities between participants by analyzing their private parameters in the BN layers. In a recent work called TrFedDis [33], the concept of feature disentangling is used to capture both domain-invariant and domain-specific knowledge. Our Fed-CO 2 deviates from such disentangling approaches and instead focuses on facilitating mutual learning between the online and offline models.

3 METHOD

In this section, we begin by formulating the research problem and subsequently introduce Fed-CO 2 , our novel universal FL framework. Finally, we will present the intra-client and inter-client knowledge transfer mechanisms designed to further enhance model cooperation in FL with feature skew.

3.1 Problem Formulation

We aim to train a set of N models for each client to fit local data distribution. Each client i has its own data distribution P i with its private dataset D i = { ( x k i , y k i ) : k ∈ { 1 , . . . , m i }} , i ∈ { 1 , . . . , N } and m i is the sample number in the private dataset. Due to label distribution skew or feature skew, the data distribution P i ( x, y ) differs on a client-by-client basis. Our goal is to learn a good model F i ( · ) with parameters θ i for each client according to information in all clients:

where l ( · , · ) is a sample-wise loss function, F i is constructed by a feature extractor f i ( · ) with parameters η i and a classifier C i ( · ) with parameters ϕ i . Accordingly, θ i = { η i , ϕ i } .

Server

OnlineModel

OnlineModel

OnlineModel

OfflineModel

OfflineModel

OfflineModel

Client 1

Client 2

Client N

SharedParameters

PersonalizedParameters

(b) Intra-Client Knowledge Transfer

(c) Inter-Client Knowledge Transfer

OfflineModellastround

FixedClassifiersfromotherclients

Lonr

gen

OnlineModel

KnowledgeTransfer

OnlineModel

→LCE

OfflineModel

→LCE

OnlineModel lastround

KnowledgeTransfer

Loff

gen

FixedClassifiersfromotherclients

3.2 Cooperation of Online and Offline Models

As shown in Fig. 1(a), we propose Fed-CO 2 : a novel universal collaboration framework between the online and offline models to adapt to local data distribution in the presence of label distribution skew, feature skew, or both. Specifically, for each client i , we train a partially personalized model, referred to as the online model F on i , and a locally-trained model, referred to as the offline model F off i . With two models in each client F i = { F on i , F off i } , cooperation is achieved through prediction fusion of the online and offline models. In this manner, our framework aims to exploit both online general and offline specialized knowledge, enabling consistent high performance in FL across various cases of local data heterogeneity.

For the online model, we personalize only a few critical parameters to enable the online model to learn knowledge from other clients through model aggregation on the server, while still retaining essential local knowledge. Prior work FedBN [10] has shown that Batch Normalization (BN) layers can capture feature distribution in local clients. Inspired by such discovery, we personalize the BN layers in the online model and split its parameters θ on i into { θ on p,i , θ on g } , where θ on p,i denotes all BN layer parameters and θ on g denotes other parameters. Model aggregation is done on the server to obtain the shared ˜ θ on g among all clients with the formula:

For the offline model, all its parameters are personalized and do not engage in server aggregation. As a result, the offline model is able to learn local specialized knowledge without forgetting or being contaminated by other irrelevant information. With online and offline models, we employ a simple fusion technique and obtain the final prediction for model F i :

It's worth noting that the online-offline model cooperation does not require extra communication costs compared to standard federated aggregation operations since offline models are not uploaded to the server.

When feature skew data heterogeneity is present, relying solely on prediction fusion for cooperation is insufficient. This approach fails the communication between the online and offline models during the training phase, impeding the transfer of online domain-invariant and offline domain-specific knowledge. Therefore, to enhance collaboration among the models in FL with feature skew, we propose novel intra-client and inter-client knowledge transfer mechanisms. In the next two parts, we will introduce these mechanisms and discuss their role.

3.3 Intra-Client Knowledge Transfer

In our framework, each client i is equipped with two models: an online model for capturing general domain-invariant knowledge and an offline model for capturing specialized domain-specific knowledge. Unlike previous algorithms that utilized a two-branch framework, aiming to make the knowledge learned by each branch as disjoint as possible through methods such as Contrastive Learning [26] or Feature Disentangling [33], we introduce a novel intra-client knowledge transfer mechanism to enhance the collaboration of the two models. This mechanism employs knowledge distillation to facilitate mutual learning between the two models during the training phase. However, transferring online general knowledge to the offline model could potentially conflict with the local training objective. The former emphasizes learning global information among { D j } j = i and avoiding over-fitting on local data, whereas the latter focuses on fitting local information in D i . As a result, neither of the two goals can be fully accomplished. Furthermore, during local training, there is a risk of the online model losing general domain-invariant knowledge, which leads to a reduction in the amount of valuable information accessible to the offline model.

̸

To address these challenges, we divide the local training process into two distinct phases: a mutual learning phase and a local adaptation phase to achieve their respective goals. Concretely, in the mutual learning phase, we create a duplicate of the initial online model, which consists of its personalized BNs from the last round and other parameters updated with the global aggregation model. Additionally, we retain a copy of the initial offline model from the last round. As illustrated in Fig. 1(b), these duplicated models are frozen and utilized as teacher networks, with the copied online model serving as a teacher for the offline model, and the copied offline model guiding the online model. We denote the duplicated initial online model and copied initial offline model as ¯ F on i and ¯ F off i , respectively, with frozen parameters ¯ θ on i and ¯ θ off i . In training step s , given input sample x k , online and offline models conduct intra-client knowledge transfer and update their parameters in the mutual learning form:

where KL ( · , · ) denotes the Kullback-Leibler (KL) divergence and α is the learning rate. Through mutual learning, the online model and the offline model engage in model-level collaboration and share the knowledge acquired in the previous round.

3.4 Inter-Client Knowledge Transfer

After the intra-client knowledge transfer in the first phase, both the online and offline models have gained beneficial knowledge from each other, and now it is opportune to let them adapt to the local data distributions. The pertinent question is whether we can leverage additional knowledge from other clients to bridge the domain gaps in the local training. Currently, the approach of acquiring knowledge from other clients is limited to aggregating shared parameters on the server. In comparison to various advanced approaches used in Domain Generalization (DG) research, parameter aggregation alone may struggle to learn effective domain-invariant features. The primary challenge stems from the isolation of private data, as FL prohibits cross-client access to data from other domains. Without data from multiple domains, it is tough to directly apply techniques in DG.

Considering that the offline model in each client is trained to capture local specialized domain-specific knowledge, we propose to leverage these lightweight classifiers of the offline model to transfer domain

knowledge among clients, as illustrated in Fig. 1(c). The classifiers of the offline model from other clients are introduced to each local client, enabling the feature extractors of its online and offline models to generate robust and well-generalized features that can be effectively recognized by these introduced classifiers. We share similar design intuitions with COPA [34], but here our inter-client knowledge transfer mechanism stresses more on making use of knowledge from other clients to benefit the personalized model for each client rather than serving unseen novel clients.

To be specific, in Fed-CO 2 , each client uploads its classifier C off i from its offline model to the server, along with its shared parameters (BNs are not included) from its online model. These uploaded heads construct a classifier set for inter-client knowledge transfer: { C off j } N j =1 with parameter set { ϕ off j } N j =1 and are transmitted to each client in the next round. In communication round t , each client has its own personalized model F i = { F on i , F off i } and the downloaded knowledge-transfer classifier set { C off j } N j =1 . Similar to teacher models in the mutual learning phase, we freeze these knowledge-transfer classifiers to keep their knowledge unchanged. With input data x k and its label y k our inter-client knowledge transfer loss function L gen for models on client i is formulated as:

̸

Combined with its classification loss to adapt to local private data, the total loss in local adaptation phase for each model is:

where µ is a penalized factor set as 1 by default. After local training, we upload parameters to the server and aggregate them with Eq. 2. Algorithm 1 in Appendix demonstrates the procedures of our Fed-CO 2 algorithm.

4 Theoretical Analysis

Here, we provide convergence analyses, which compare our Fed-CO 2 and FedBN [10] with the neural tangent kernel (NTK) [35] theory, to illustrate the effectiveness of our algorithm. For simplification, only the cooperation of online and offline models is considered while the intra-client and inter-client knowledge transfer mechanisms are ignored.

4.1 Formulation

Before our analyses, we clarify the simplified setup and basic assumptions. Suppose we have N clients, each with M training examples, and we aim to jointly train them for T communication rounds. In each round, each model is locally trained for one epoch. For simplification, we assume that all clients share an identical two-layer neural network for a regression task, with the only source of heterogeneity being the feature skew. Namely, each client i has its private dataset D i = { ( x i j , y i j ) }} j M =1 with x i j ∈ R d and y i j ∈ R . For convenience, we adopt the same data distribution assumption as in [10]:

̸

Assumption 4.1 For each client i ∈ { 1 , . . . , N } , the inputs x i j are centered, meaning that E [ x i ] = 0 , and they have a covariance matrix S i = E [ x i ( x i ) ⊤ ] . S i is independent of the label y and varies for each client i . Not all S i matrices are identity matrices. For any index pair p and q ( p = q ), we have x p = w · x q for any non-zero w .

̸

Let v k ∈ R d represent the parameters of the first layer, where k ∈ [ m ] and m is the width of the hidden layer. We define ∥ v ∥ S := √ v ⊤ Sv as the induced vector norm for a positive definite matrix S and a l 2 vector norm ∥ v ∥ 2 . The projections of x onto v and v ⊥ are defined as x v := vv ⊤ x ∥ v ∥ 2 2 , x v ⊥ := ( I -vv ⊤ ∥ v ∥ 2 ) x .

Based on Assumption 4.1, for client i , the output of the first layer is normalized as v ⊤ k x i ∥ v k ∥ S i . The shift parameter of BN is omitted. Further, we denote the scaling parameter of BN γ ∈ R m × N and the second layer parameter c ∈ R m .

With these parameters, we proceed to train the online model:

where σ ( · ) is the ReLU activation function and γ on k,i is personalized BN parameters. For the offline model with all parameters personalized, we train:

In our analysis, we adopt one random strategy [36] to initialize parameters:

where α 2 controls the magnitude of v on k and v off k,i at initialization. The initialization of the BN parameters γ on k,i and γ off k,i are independent of α . We use the MSE loss to train our model F :

4.2 Convergence Analysis

We employ NTK [35] to analyze the trajectory of networks F learned by Fed-CO 2 and networks F on learned by FedBN (the online model in Fed-CO 2 ). Existing studies [37, 38] have validated that the convergence rate of finite-width over-parameterized networks is controlled by the least eigenvalue of the induced kernel in the training process. Following the discovery in [38], we can decompose the NTK into a direction component V ( t ) /α 2 and a magnitude component G ( t ) :

The specific forms of V ( t ) and G ( t ) are provided in the Appendix. Here, let λ min ( H ) denote the minimum eigenvalue of matrix H . It is worth noting that both matrices V ( t ) and G ( t ) are positive semi-definite as they can be interpreted as covariance matrices. Based on this, we can deduce that λ min ( Λ ( t )) ≥ max { λ min ( V ( t ) /α 2 ) , λ min ( G ( t )) } . From the NTK theory, the value of λ min ( Λ ( t )) controls the convergence rate. Then, considering that α is the pre-defined parameter, for α ∈ (0 , 1) , convergence is dominated by V ( t ) . Let Λ ( t ) , Λ on ( t ) and Λ off ( t ) denote the evolution dynamics of Fed-CO 2 , the online model and the offline model, respectively. Based on the work in [10], the auxiliary version of the Gram matrices V ∞ , V ∞ on , and V ∞ off for three models are strictly positive definite. Let the least eigenvalues λ min ( V ∞ on ) := µ on , λ min ( V ∞ off ) := µ off , and λ min ( V ∞ ) := µ , where µ on , µ off , and µ are all positive values. To be specific, we define Gram matrices V ∞ on and V ∞ off in Definition 4.2.

Definition 4.2 Given sample points { x p } NM p =1 , we define the auxiliary Gram matrices V ∞ on ∈ R NM × NM and V ∞ off ∈ R NM × NM as

✶ With the prerequisite provided in Appendix, the convergence performance of the online model, the offline model, and our Fed-CO 2 can be analyzed by comparing λ min ( V ∞ on ) , λ min ( V ∞ off ) , and λ min ( V ∞ ) . Our main theoretical result is given in Theorem 4.3.

Theorem 4.3 For the V-dominated convergence, the convergence rate of Fed-CO 2 is faster than that of FedBN (the online model in Fed-CO 2 ).

Proof sketch The core is to show λ min ( V ∞ on ) ≤ λ min ( V ∞ ) . To prove the theorem, we first show that λ min ( V ∞ on ) ≤ λ min ( V ∞ off ) . Comparing Eq. 14 and 15, V ∞ off takes the M × M block matrices on the diagonal of V ∞ on . Let V ∞ i be the i-th M × M block matrices on the diagonal of V ∞ on . According to linear algebra, we have λ min ( V ∞ i ) ≥ λ min ( V ∞ on ) , for i ∈ [ N ] . With V ∞ off = diag ( V ∞ 1 , · · · , V ∞ N ) , it can be obtained that λ min ( V ∞ off ) ≥ λ min ( V ∞ on ) . Then, based on Eq. 10, 12, and 13, we can obtain λ min ( V ∞ ) = λ min ( 1 2 ( V ∞ on + V ∞ off )) . Therefore, we have λ min ( V ∞ ) ≥ ( 1 2 λ min ( V ∞ on ) + 1 2 λ min ( V ∞ off )) . Since we have proven that λ min ( V ∞ off ) ≥ λ min ( V ∞ on ) , the result λ min ( V ∞ ) ≥ λ min ( V ∞ on ) can be achieved.

5 EXPERIMENTS

5.1 Experiment Settings

Dataset and Data Heterogeneity. We followed the footstep of prior research [39, 7, 5, 13] to study the label distribution skew heterogeneity with image datasets: CIFAR10 and CIFAR100 [40]. Specifically, we considered two different label distributions among participants: 1) Pathological Distribution, each client is randomly assigned 2 classes per client in CIFAR10 (10 classes per client in CIFAR100); 2) Dirichlet Distribution, each client gets its private data through partitioning of the datasets using a symmetric Dirichlet distribution with a default parameter α = 0 . 3 .

For feature skew heterogeneity, we conducted extensive experiments on three datasets: Digits, OfficeCaltech10 [41], and DomainNet [42]. In this non-IID data setting, each domain serves as a client, with each client having access to all the data from its respective domain. Digits is composed of 5 different digit datasets with feature shift: SVHN [43], USPS [44], SynthDigits [45], MNIST-M [45] and MNIST [46]. Office-Caltech10 [41] owns data from four different domains including Amazon, DSLR, WebCam, and Caltech-256. DomainNet [42] is a challenging dataset that comprises six distinct domains, namely Clipart, Infograph, Painting, Quickdraw, Real, and Sketch. To explore the more realistic FL scenarios with both label distribution skew and feature skew, we exert a Dirichlet Distribution to datasets with feature shift. More details are supplemented in Appendix.

Compared Benchmarks. We compared our universal framework Fed-CO 2 with fundamental algorithms including SingleSet , FedAvg [2] and FedProx [4], together with some refined FL algorithms: FedPer [11], MOON [7], FedBN [10], and FedRoD [12]. Here, SingleSet means each client trains their own model on their private data. We applied the linear version for FedRoD . For FL with feature skew issues, we further compared COPA [34], which was originally proposed to solve FDG and FDA problems. We now train the model on all domains in the dataset and evaluate its performance on each individual training client using the learned universal model.

5.2 Performance Evaluation

Model Evaluation on Feature Skew. Experiments are conducted on datasets Digits, OfficeCaltech10, and DomainNet. Here, we present the results of experiments on Office-Caltech10, and DomainNet in Table 1 and Table 2. Full experiments are shown in Appendix. The performance of our Fed-CO 2 exhibits a dominant edge over state-of-the-art algorithms in all experiments, where the average accuracy on each domain increases by nearly 4 % . The results confirm that our novel collaboration framework between online and offline models effectively addresses FL challenges associated with severe feature skew data heterogeneity.

Remarkably, when compared to algorithms FedBN [10] and COPA [34], which are explicitly designed to handle feature shift issues, our method consistently outperforms them across all sub-datasets. This phenomenon provides compelling evidence that our cooperation mechanism adapts to local feature shifts more effectively by leveraging both general domain-invariant and specialized domain-specific knowledge.

Model Evaluation on Label Distribution Skew. As shown in Table 3, our Fed-CO 2 outperforms a variety of state-of-the-art PFL algorithms in terms of average test accuracy across almost all cases in experiments with label distribution skew. Our state-of-the-art results across almost all cases

prove that our method can more effectively excavate and utilize local information to overcome data heterogeneity caused by label distribution imbalance than previous methods. It is worth highlighting that our Fed-CO 2 outperforms both SingleSet and FedBN, regardless of which one performs better individually. This observation confirms that Fed-CO 2 effectively integrates the general knowledge obtained by the online model and the specialized knowledge acquired by the offline model, resulting in enhanced performance.

Model Evaluation on Label Distribution Skew and Feature Skew. Here, we evaluate our Fed-CO 2 and benchmark algorithms in FL scenarios with label distribution skew and feature skew. We exhibit the experiment results on Digits in Table 4, where we exerted a Dirichlet Distribution on each client with α = 0 . 3 . From the results, it is evident that, even in the presence of two types of data heterogeneity, our Fed-CO 2 outperforms various state-of-the-art algorithms with a significant margin in every sub-dataset. Therefore, we can conclude that under our universal cooperation framework, local clients not only make the most use of their local offline information but also benefit from the general knowledge shared by other clients for a better adaptation to local data, resulting in state-ofthe-art performance in FL with both label distribution skew and feature skew. More experiments are provided in Appendix.

5.3 Component Analysis

In the presence of feature skew, we add intra-client and inter-client knowledge transfer mechanisms to foster the collaboration between the online and offline models for better local adaptation. Here, we evaluate the efficacy of these two knowledge transfer mechanisms on the challenging DomainNet dataset through three additional experiments, as presented in Table 5. In these experiments, we selectively removed specific mechanisms from Fed-CO 2 to investigate their individual contributions and functionalities.

According to the experimental results, we can observe that: 1) In the absence of intra-client and inter-client knowledge transfer mechanisms, Fed-CO 2 exhibits a significant decline in performance, resulting in a performance level comparable to that of benchmark methods. Therefore, cooperation relying solely on prediction fusion proves to be inadequate for FL with feature skew. 2) The removal of either the intra-client or the inter-client knowledge transfer mechanisms will also result in an average decline in performance. This phenomenon validates the effectiveness of our intra-client and inter-client knowledge transfer mechanisms, both of which contribute to enhancing the models' capability to adapt to local data distributions. 3) Compared with all three ablation experiments, our Fed-CO 2 consistently achieves the highest performance on average and across the majority of subdatasets. These results firmly demonstrate that our cooperation framework effectively leverages both model-level and client-level collaboration through the intra-client knowledge transfer and inter-client knowledge transfer mechanisms, respectively, to adapt to feature shifts resulting from domain gaps. Further analysis of these two cooperation mechanisms is supplemented in the Appendix.

6 CONCLUSIONS

In this work, we proposed a new universal FL framework called Fed-CO 2 that is capable of handling label distribution skew and feature skew even when both are present in the same data. The core of our approach is to foster cooperation between the online and offline models, leveraging the benefits of online general knowledge and offline specialized knowledge to effectively adapt to local data distribution. To improve model collaboration in FL with feature shifts, we designed two novel knowledge transfer mechanisms: one intra-client and the other inter-client. These mechanisms facilitate mutual learning between the online and offline models, concurrently enhancing the model's capacity to generalize across different domains. Comparisons with a wide range of state-of-the-art methods on five benchmark datasets consistently show that Fed-CO 2 yields superior performance in addressing both label distribution skew and feature skew challenges, both individually and in combination. Extending Fed-CO 2 to scenarios with noisy training data is under consideration in our future work.

References

Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST) , 10(2):1-19, 2019.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning , 14(1-2):1-210, 2021.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems , 2:429-450, 2020.

Othmane Marfoq, Giovanni Neglia, Richard Vidal, and Laetitia Kameni. Personalized federated learning through local memorization. In International Conference on Machine Learning , pages 15070-15092. PMLR, 2022.

Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernetworks. In International Conference on Machine Learning , pages 9489-9502. PMLR, 2021.

Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10713-10722, 2021.

Benyuan Sun, Hongxing Huo, Yi Yang, and Bo Bai. Partialfed: Cross-domain personalized federated learning via partial initialization. Advances in Neural Information Processing Systems , 34:23309-23320, 2021.

Wang Lu, Jindong Wang, Yiqiang Chen, Xin Qin, Renjun Xu, Dimitrios Dimitriadis, and Tao Qin. Personalized federated learning with adaptive batchnorm for healthcare. IEEE Transactions on Big Data , 2022.

Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. International Conference on Learning Representations , 2021.

Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arXiv preprint arXiv:1912.00818 , 2019.

Hong-You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. International Conference on Learning Representations , 2022.

Hongxia Li, Zhongyi Cai, Jingya Wang, Jiangnan Tang, Weiping Ding, Chin-Teng Lin, and Ye Shi. FedTP: Federated Learning by Transformer Personalization. IEEE Transactions on Neural Networks and Learning Systems , pages 1-15, 2023.

Kangkang Wang, Rajiv Mathews, Chloé Kiddon, Hubert Eichner, Françoise Beaufays, and Daniel Ramage. Federated evaluation of on-device personalization. arXiv preprint arXiv:1910.10252 , 2019.

Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002.10619 , 2020.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems , 33:3557-3568, 2020.

Mikhail Khodak, Maria-Florina F Balcan, and Ameet S Talwalkar. Adaptive gradient-based meta-learning methods. Advances in Neural Information Processing Systems , 32, 2019.

Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems , 33:21394-21405, 2020.

Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning , pages 6357-6368. PMLR, 2021.

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International Conference on Machine Learning , pages 2089-2099. PMLR, 2021.

Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01523 , 2020.

Yiqing Shen, Yuyin Zhou, and Lequan Yu. Cd2-pfed: Cyclic distillation-guided channel decoupling for model personalization in federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10041-10050, 2022.

Ang Li, Jingwei Sun, Xiao Zeng, Mi Zhang, Hai Li, and Yiran Chen. Fedmask: Joint computation and communication-efficient personalized federated learning via heterogeneous masking. In Proceedings of the 19th ACM Conference on Embedded Networked Sensor Systems , pages 42-55, 2021.

David Ha, Andrew Dai, and Quoc V Le. Hypernetworks. arXiv preprint arXiv:1609.09106 , 2016.

Xiaosong Ma, Jie Zhang, Song Guo, and Wenchao Xu. Layer-wised model aggregation for personalized federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10092-10101, 2022.

Yupei Zhang, Yunan Xu, Shuangshuang Wei, Yifei Wang, Yuxin Li, and Xuequn Shang. Doubly contrastive representation learning for federated image recognition. Pattern Recognition , 139:109507, 2023.

Xutong Mu, Yulong Shen, Ke Cheng, Xueli Geng, Jiaxuan Fu, Tao Zhang, and Zhiwei Zhang. Fedproc: Prototypical contrastive federated learning on non-iid data. Future Generation Computer Systems , 143:93-104, 2023.

Tao Shen, Jie Zhang, Xinkang Jia, Fengda Zhang, Gang Huang, Pan Zhou, Kun Kuang, Fei Wu, and Chao Wu. Federated mutual learning. arXiv preprint arXiv:2006.16765 , 2020.

Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv preprint arXiv:1910.03581 , 2019.

Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems , 33:2351-2363, 2020.

Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang, Wenchao Xu, and Feijie Wu. Parameterized knowledge transfer for personalized federated learning. Advances in Neural Information Processing Systems , 34:10092-10104, 2021.

Huancheng Chen, Haris Vikalo, et al. The Best of Both Worlds: Accurate Global and Personalized Models through Federated Learning with Data-free Hyper-Knowledge Distillation. arXiv preprint arXiv:2301.08968 , 2023.

Meng Wang, Kai Yu, Chun-Mei Feng, Yiming Qian, Ke Zou, Lianyu Wang, Rick Siow Mong Goh, Xinxing Xu, Yong Liu, and Huazhu Fu. TrFedDis: Trusted Federated Disentangling Network for Non-IID Domain Feature. arXiv preprint arXiv:2301.12798 , 2023.

Guile Wu and Shaogang Gong. Collaborative optimization and aggregation for decentralized domain generalization and adaptation. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 6484-6493, 2021.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems , 29, 2016.

Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054 , 2018.

Yonatan Dukler, Quanquan Gu, and Guido Montúfar. Optimization theory for relu neural networks trained with normalization layers. In International conference on machine learning , pages 2751-2760. PMLR, 2020.

Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koneˇ cn` y, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295 , 2020.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Boqing Gong, Yuan Shi, Fei Sha, and Kristen Grauman. Geodesic flow kernel for unsupervised domain adaptation. In 2012 IEEE conference on computer vision and pattern recognition , pages 2066-2073. IEEE, 2012.

Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision , pages 1406-1415, 2019.

Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.

Jonathan J. Hull. A database for handwritten text recognition research. IEEE Transactions on pattern analysis and machine intelligence , 16(5):550-554, 1994.

Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference on machine learning , pages 1180-1189. PMLR, 2015.

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998.

Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. arXiv preprint arXiv:2111.04263 , 2021.

Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advances in Neural Information Processing Systems , 34:8392-8406, 2021.

Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM , 60(6):84-90, 2017.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

A Dataset Setup Details

The details of non-IID settings of mentioned datasets are presented in the section. Our two settings of label distribution skew, namely the Pathological setting and the Dirichlet setting are similar to preceding works [5, 13]. In the Pathological setting, where each client i is randomly assigned two (ten) classes from CIFAR10 (CIFAR100), we first sampled q i c ∼ U (0 . 4 , 0 . 6) for the selected class

c

and then calculated the sample rate for this class

c

on client

i

with sr

i

=

q

∑

c

i

c

j

q

j

c

setting, we partitioned the datasets randomly utilizing a symmetric Dirichlet distribution with default parameter α = 0 . 3 . Specifically, for CIFAR100, we leveraged its coarse labels following previous works [5, 39]. In detail, a two-stage Pachinko allocation method was adopted to partition samples in CIFAR100. For each client, this method first generates a Dirichlet distribution with default parameter α = 0 . 3 over the coarse labels and then generates another Dirichlet distribution with parameter β = 10 over the corresponding fine labels. The local training batch size was set to 64. In both non-IID partitions, the classes and their distribution in each client's training and test sets remain consistent.

We adopted the setting used in FedBN [10] for feature skew heterogeneity, which involved adjusting the local training batch size to 32 to maintain consistency. Additionally, we followed the implementation of selecting only the top 10 classes based on data amount from DomainNet to create the dataset for our experiments. For experiments with label distribution skew and feature skew, we partitioned each sub-dataset into several parts with a Dirichlet Distribution and selected one part as the new training set. Then we adjusted the label distribution in the test set to the same distribution in the new training set. Specifically, based on the number of training images in each dataset, we divided each sub-dataset in Digits into three parts, each sub-dataset in Office-Caltech10 into two parts, and each sub-dataset in DomainNet into five parts. Table 6 summarizes the datasets and the number of clients.

B Implementation Details

Network Modules. Similar to prior works [47, 48], we adopted a ConvNet [46] with two convolutional layers and three fully-connected layers for methods in experiments conducted under label distribution skew setting. With regards to feature skew setting, experiments were performed with a backbone similar to FedBN [10]. In other words, for the Digits dataset, we applied a ConvNet [46] with four convolutional layers and three fc layers. For Office-Caltech10 and DomainNet dataset, we applied one Alexnet [49] as the default model backbone.

In our experimental evaluations of model performance under label distribution skew heterogeneity, we added an additional BN layer after each convolutional layer and fully connected (fc) layer, except for the last layer to algorithms: FedBN [10] and Fed-CO 2 . Meanwhile, in experiments pertaining to feature skew heterogeneity, as our default network incorporated a BN layer after each convolutional and fc layer (except for the last layer), we omitted the BNs in our classifiers for methods other than FedBN [10] and Fed-CO 2 . All other Convnet network architectures used for the Digits dataset, as well as the Alexnet architecture employed for Office-Caltech10 and DomainNet, are identical to those used in FedBN [10].

Concrete Implementation. For experiments evaluating model performance on label distribution skew, we set up 100 local clients and randomly sampled 5 % of them to participate in each communication round. We trained every algorithm with 1500 communication rounds. Meanwhile, in experiments about feature skew non-iid heterogeneity, all clients participated in each communication round and every algorithm was trained with 300 communication rounds. All training tasks were optimized with an SGD optimizer with the default learning rate 0.01. We used PyTorch [50] to implement our algorithms with an RTX 2080 Ti GPU.

. In the Dirichlet

For experiments on CIFAR10 and CIFAR100, we tested the model performance every 5 rounds during its last 200 communication rounds. The average and variance of test accuracy were reported. For all experiments with feature skew, we run 5 times and reported its mean and variance.

Evaluation Metrics. We evaluated the learned personal models on the private test data of each client. We examined the model's performance on each domain using the same method as FedBN [10] in FL with feature skew. Furthermore, for experiments on label distribution skew, we calculated the accuracy using the same protocol as FedTP [13].

C Convergence

In this section, we first provide the preliminary and complete theorem proof. Then, we empirically demonstrate that our Fed-CO 2 shows faster and more robust convergence performance than FedBN and other benchmark methods on FL issues with Feature Skew.

C.1 Preliminary

Following the V-dominated convergence rate for FedAvg (Theorem C.1) in [38], we can derive the convergence rate for the online model, the offline model, and Fed-CO 2 Corollary C.2.

Theorem C.1 ( V -dominated convergence for FedAvg [38]) Suppose the network is initialized as in Eq. 11 with α ≤ 1 , trained using gradient descent and Assumption 4.1 holds. Assuming the loss function used for training the neural network is the square loss, and the target values y satisfy ∥ y ∥ ∞ = O (1) . If m = Ω( N 4 M 4 log ( NM/δ ) /µ 4 0 ) , then with probability 1 -δ ,

for iterations t = 0 , 1 , · · · , the evolution matrix Λ ( t ) satisfies λ min ( Λ ( t )) ≥ µ 0 2 α 2 ;

training with gradient descent of step-size η = O ( α 2 ∥ V ∞ ∥ 2 ) converges linearly as

Corollary C.2 Under the same assumption in Theorem C.1, with probability 1 -δ , for iterations t = 0 , 1 , · · · , we have

The evolution matrix Λ on ( t ) satisfies λ min ( Λ on ( t )) ≥ µ on 2 α 2 and training with gradient descent of step-size η = O ( α 2 ∥ V ∞ on ∥ 2 ) converges linearly as ∥ F on ( t ) -y ∥ 2 2 ≤ ( 1 -ηµ on 2 α 2 ) t ∥ F on (0) -y ∥ 2 2 .

The evolution matrix Λ off ( t ) satisfies λ min ( Λ off ( t )) ≥ µ off 2 α 2 and training with gradient descent of step-size η = O ( α 2 ∥ V ∞ off ∥ 2 ) converges linearly as ∥ F off ( t ) -y ∥ 2 2 ≤ ( 1 -ηµ off 2 α 2 ) t ∥ F off (0) -y ∥ 2 2 .

The evolution matrix Λ ( t ) satisfies λ min ( Λ ( t )) ≥ µ 2 α 2 and training with gradient descent of step-size η = O ( α 2 ∥ V ∞ ∥ 2 ) converges linearly as ∥ F ( t ) -y ∥ 2 2 ≤ ( 1 -ηµ 2 α 2 ) t ∥ F (0) -y ∥ 2 2 .

Therefore, the exponential factor of convergence for the online model, the offline model, and Fed-CO 2 are controlled by the smallest eigenvalue of V on ( t ) , V off ( t ) , and V ( t ) , respectively. The convergence performance of the online model, the offline model, and our Fed-CO 2 can be analyzed by comparing λ min ( V ∞ on ) , λ min ( V ∞ off ) , and λ min ( V ∞ ) .

C.2 Theorem Proof

To prove Theorem 4.3, we will first prove SingleSet (the offline model in Fed-CO 2 ) converges faster than FedBN (the online model in Fed-CO 2 ). Gram matrix V on ( t ) and G on ( t ) of FedBN can be

obtained from work [10] as:

To compare the convergence rates of SingleSet and FedBN, we compare the exponential factor in the convergence rates, which are (1 -ηµ off 2 α 2 ) and (1 -ηµ on 2 α 2 ) , respectively. Then, considering that α is the pre-defined parameter, it reduces to comparing µ on = λ min ( V ∞ on ) and µ off = λ min ( V ∞ off ) . Comparing Eq. 16 and 18, V ∞ off takes the M × Mblock matrices on the diagonal of V ∞ on :

where V ∞ i is the i-th M × M block matrices on the diagonal of V ∞ on . Therefore, we have:

Since the eigenvalues of V ∞ off are exactly the union of eigenvalues of V ∞ i , we get:

Hence, (1 -ηµ on 2 α 2 ) ≥ (1 -ηµ off 2 α 2 ) and we get the first-stage conclusion that the convergence rate of method SingleSet is faster than that in FedBN.

Then we go to prove Fed-CO 2 is faster than FedBN, as well. With Eq. 10, 12, and 13, we get:

Based on Eq. 19, we get the following relationships among three models:

To compare convergence rates of Fed-CO 2 and FedBN, we need to compare their exponential factor (1 -ηµ 2 α 2 ) and (1 -ηµ on 2 α 2 ) , as well. Similar to the former proof, the problem reduces to compare µ = λ min ( 1 2 ( V ∞ on + V ∞ off ) ) with µ on = λ min ( V ∞ on ) . With previous proof result λ min ( V ∞ off ) ≥ λ min ( V ∞ on ) , we get

Therefore, we have (1 -ηµ on 2 α 2 ) ≥ (1 -ηµ 2 α 2 ) . This theoretical conclusion guarantees a faster convergence in our Fed-CO 2 than FedBN.

C.3 Experimental Performance

Here, we empirically analyze the convergence of test accuracy for various algorithms, as well as the convergence behavior of both the online and offline models. The results for sub-dataset WebCam and SVHNare illustrated in Fig. 2 as examples. As Fig. 2 exhibits, among various FL algorithms, our FedCO 2 achieves the highest accuracy with significantly faster convergence speed and demonstrates much more robust behavior. Compared with its online and offline models, Fed-CO 2 consistently surpasses both models, exhibiting a smoother curve, which proves the success in fusing domain-invariant and domain-specific knowledge.

0A

(b)

WebCam

SVHN

1.0

0.9

0.8

0.8

0.7

0.7

0.6

Test Accuracy

0.5

0.5

FedAvg

FedAvg

0.3

FedProx

FedPer

E0

FedProx

FedPer

0.2

MOON

MOON

FedRoD

0.2

FedRoD

0.1

FedBN

FedBN

0.0

Fed-CO2

0.1

Fed-CO2

CommunicationRound

Communication Round

WebCam

SVHN

1.0

0.9

0.8

0.8

0.7

0.7

0.6

Test Accuracy

0.5

0.5

0.4

0.4

0.3

0.3

0.2

0.1

Online Model

offline Model

0.2

Online Model

Offline Model

0.0

Fed-CO2

0.1

Fed-CO2

CommunicationRound

Communication Round

D Fed-CO 2 Algorithm

Here, we provide a detailed description of the algorithm for our universal FL framework, Fed-CO 2 . The learning process for FL with and without Feature Skew is presented in Algorithm 1 and Algorithm 2, respectively.

Algorithm 1: Fed-CO 2 for FL with Feature Skew

E More Experimental Results On Benchmark Datasets

Here, we supplement more experimental results on benchmark datasets. Firstly, we present the experimental results on Digits dataset for FL with feature skew in Table 7. As the table shows, our Fed-CO 2 surpasses other various state-of-the-art algorithms with a substantial margin, which further proves the effectiveness of our framework in adapting local feature shifts.

Then, we show experimental results in FL with both label distribution skew and feature skew on DomainNet and Office-Caltech10 in Table 8 and Table 9. We applied a Dirichlet Distribution with α = 2 . 0 to DomainNet and another Dirichlet Distribution with α = 1 . 0 to Office-Caltech10. It can be observed that when facing two severe forms of data heterogeneity, algorithms demonstrated significantly lower performance on Office-Caltech10 and DomainNet compared to FL scenarios with feature skew alone. Even in such a challenging task, our Fed-CO 2 still achieves the highest average accuracy and outperforms locally-trained models. This result proves that Fed-CO 2 succeeds in making use of both general domain-invariant and specialized domain-specific knowledge to adapt to extreme local data heterogeneity.

Additionally, to explore broader real-life data heterogeneity issues with both feature skew and label distribution skew, we conducted further experiments by exerting different levels of label distribution skew on the Digits dataset. Specifically, we partitioned each sub-dataset of Digits using the Dirichlet Distribution, with different values of α ∈ { 0 . 3 , 0 . 7 , 1 . 0 , 1 . 5 , 2 . 0 } . A smaller α value indicates a more severe label distribution imbalance. The experimental results are demonstrated in Fig. 3. It is evident that in every instance of label distribution imbalance, our Fed-CO 2 exhibits a significant advantage over other benchmark methods in its sub-dataset as well as the overall average performance. Therefore, we can firmly conclude that our universal FL framework, Fed-CO 2 , effectively addresses a wide range of real-life data heterogeneity issues.

F Additional Experiments and Analyses

F.1 Effectiveness of Cooperation Framework

The experimental results of testing accuracy on five benchmark datasets have validated that our Fed-CO 2 effectively combines online general knowledge and offline specialized knowledge, resulting in improved adaptation to local data distribution in FL scenarios with label imbalance and feature

(a)TestAccuracy onSynthDigitssub-dataset

(b)AverageTestAccuracyonDigits dataset

0.95

0.95

0.90

0.90

0.85

TestAccuracy

0.85

0.80

0.80

SingleSet

SingleSet

FedAvg

FedAvg

0.75

口

FedPer

0.75

FedPer

FedRoD

FedRoD

FedBN

FedBN

0.70

COPA

0.70

COPA

Fed-CO2

Fed-CO2

0.5

1.0

1.5

2.0

0.5

1.0

1.5

2.0

αinDirichletDistribution

αinDirichletDistribution

shifts. To further confirm the effectiveness of the cooperation between the online and offline models, we conducted additional experiments for FL with label imbalance and feature shifts in the following parts.

Online and Offline Models in FL with Label Distribution Skew. For FL with label imbalance, we split the CIFAR10 dataset for 20 clients with various types and degrees of label distribution skew. In the Pathological case, each client i was randomly assigned two classes following the strategies mentioned in Appendix A. In the Dirichlet case, we partitioned the dataset randomly utilizing the Dirichlet distribution with different parameters α ∈ { 0 . 3 , 0 . 7 , 1 . 0 } to evaluate model performances under various extents of label distribution imbalance. In these experiments, we compared the model trained by our Fed-CO 2 with its online and offline models.

The experimental results have been exhibited in Fig. 4. It is evident that the online model outperforms the offline model when the label distribution imbalance is moderate, while the offline model excels when the imbalance is severe. Across all extensive experiments, Fed-CO 2 consistently surpasses both the standalone online model and the standalone offline model in terms of test accuracy statistics, exhibiting higher average and median values. At the client level, Fed-CO 2 improves prediction accuracy for the majority of local clients. Although Fed-CO 2 may not surpass the superior model between the online and offline models on a limited number of clients, it consistently outperforms the weaker model and achieves performance levels approaching that of the superior model. Therefore, we can conclude that our collaborative FL framework, Fed-CO 2 , which combines the strengths of online and offline models, is exceptionally effective. It successfully integrates online general knowledge and offline specialized knowledge, resulting in an improved adaptation to local data distribution under various cases.

Online and Offline Models in FL with Feature Skew. For FL with feature skew, we conducted experiments that evaluate the performance of the online and offline models on DomainNet. Based on the experimental results shown in Table 10, we have the following observations: First, the online model outperforms the offline model on most clients, except for the Quickdraw client. This phenomenon validates our hypothesis that when the feature shift is severe, the model aggregation will lead to the loss of important local offline information, resulting in the model's failure to adapt to such significant data heterogeneity. Second, prediction fusion is even inferior to the online and offline models for some clients. This result reflects that prediction fusion is not sufficient to fuse online general knowledge and offline specialized knowledge in FL with severe data heterogeneity issues. Therefore, for FL with feature skew, intra-client and inter-client knowledge transfer mechanisms are required to boost model and client cooperation under our novel framework.

(a)Pathological

Label Distribution

TestAccuracyStatistics

Performance on each participant

14000.975

1.00 -

1200

0.950

0.95 -

1000

0.925

0.90 -

800

0.900

0.875

600

0.80-

0.850

400

0.75 -

0.825

200

0.800

0.70 -

10

Client IDs

Offline Model

Fed-CO2

0.65 -

Client IDs

19

(b)Dirichlet α=0.3

Fed-CO2

Online Model

offline Model

Label Distribution

2000

0.90

1.0 -

Performance on each participant

1750

0.85

0.9 -

1500

0.80

0.8 -

1250

0.75

0.7-

1000

750

0.70

0.6-

500

0.65

0.5 -

250

0.60

0.4 -

6 78 91011121314151617181920

Client IDs

0

Online Model

offine Model

Fed-CO2

0.3 -

Client IDs

Fed-CO2

Online Model

offline Model

(c)Dirichlet α=0.7

Label Distribution

1400

0.85

1.0 -

Performance oneach participant

1200

0.80

0.9 -

1000

0.75

0.8 -

800

0.70

0.7-

600

0.65

0.6

0.60

400

0.5 -

0.55

200

0.50

0.4 -

10

Client IDs

-0

Fed-CO2

0.3 -

ClientIDs

1011121314151617181920

(d)Dirichlet α=1.0

Fed-CO2

Online Model

offline Model

Label Distribution

Test AccuracyStatistics

Performance on eachparticipant

1200

1.0 -

1000

0.75

0.9-

0.70

0.8-

800

0.65

0.7

600

400

0.60

0.6

0.5-

200

0.55

0.4-

10

0.50

Client IDs

-0

Online Model

offline Model

Fed-CO2

0.3-

2

ClientIDs

1011121314151617181920

Fed-CO2

Online Model

Offline Model

F.2 Knowledge Transfer Mechanisms in Feature Skew

As mentioned in the main paper, intra-client and inter-client knowledge transfer mechanisms foster model-level and client-level collaboration to better make use of general domain-invariant and specialized domain-specific knowledge for a better adaptation to local data distribution. In this section, we did further ablation studies on the challenging DomainNet dataset to investigate two mechanisms from a fine-grained angle.

Ablation Study in Intra-client Knowledge Transfer Mechanism. In this section, we investigated the impact of the intra-client knowledge transfer mechanism on the online and offline models by excluding the inter-client knowledge transfer mechanism from Fed-CO 2 . Specifically, we performed separate intra-client knowledge transfers: from the offline model to the online model and from the online model to the offline model. The results are demonstrated in Table 11 and we can discover that knowledge transfer is beneficial to both online and offline models. By transferring specialized domain-specific knowledge to the online model, we can observe an improvement in the performance of the Quickdraw client, which has a more extreme non-IID data distribution and relies more on local offline information. Conversely, when transferring general domain-invariant knowledge to the offline model, we could notice enhanced performance in the remaining clients that present milder data heterogeneity and require more global information. Hence, by incorporating the full intra-client knowledge transfer mechanism, our Fed-CO 2 effectively utilizes both online domain-invariant and offline domain-specific knowledge, resulting in improved average test accuracy.

Ablation Study in Inter-client Knowledge Transfer Mechanism. In this part, we removed the intra-client knowledge transfer mechanism from Fed-CO 2 and did experiments to investigate the impact of inter-client knowledge transfer mechanism on online and offline models. The experiments included exerting inter-client knowledge transfer mechanism on the online model, the offline model, and both models. Table 12 presents the experimental results. The results demonstrate that knowledge from other domains benefits both the online and offline models, with the online model showing more remarkable improvement. This observation aligns consistently with the purpose of the online model, which focuses on acquiring general domain-invariant knowledge that can be applied across various domains. Furthermore, the application of the inter-client knowledge transfer mechanism to both the online and offline models yields optimal performance, which means simultaneously enhancing domain generalization ability for online and offline models can exploit domain-invariant knowledge more effectively. In summary, our inter-client knowledge transfer mechanism enables each local client to acquire and exploit more domain-invariant knowledge resulting in a better adaptation to local data distribution.

F.3 Data Heterogeneity of Noise

Noise is another common factor causing data heterogeneity issues in real FL scenarios, where each client has the same label distribution but is subject to various degrees of noise. Therefore, we implemented additional experiments on this new scenario to make the initial attempt. In detail, we divided the dataset CIFAR10 into 10 clients in a way to make them independently and identically distributed. Then we applied the increasing level of random Gaussian noise with mean µ i = 0 and standard deviation σ i = σ M N -1 ∗ i , where i ∈ { 0 , 1 , · · · , N -1 } . In this series of initial experiments,

we let the client number and the sample rate be 0.5 with σ M ∈ { 0 , 5 , 10 , 15 , 20 } . The backbone of each method is the same as that in other experiments.

As presented in Fig. 5, the experimental results reveal that our Fed-CO 2 consistently exhibits better and more robust performance across nearly all experiments, particularly in scenarios with severe noise. These experiments will also serve as a source of inspiration for our future work.

TestAccuracyofBenchmarkMethodswithVariousNoiseLevels

0.65

FedAvg

MOON

0.60

FedProx

0.55

FedBN

Fed-CO2

0.50

0.45

0.40

0.35

0.30

0.25

0.20

0.15

0.10

0

5

10

15

20

NoiseLevel

F.4 Effects of Data Number and Client Number

The sample size of the training set and the client number are significant factors. Therefore, in this section, we conducted a series of experiments to explore their influence.

To investigate the effect of the training sample number, we utilized dataset Digits and set 10 % of its training data as the original training dataset. Then, we selected a certain ratio γ of the original training dataset as a new training dataset for each client, where γ ∈ { 0 . 2 , 0 . 4 , 0 . 6 , 1 . 0 , 2 . 0 } . The results shown in Fig. 6 demonstrate that our Fed-CO 2 have an edge over various benchmark methods in almost every case, especially when the number of training data is very limited.

(a)Test Accuracy on MNIST-M sub-dataset

(b)Average Test Accuracy on Digits dataset

0.95

0.95

0.90

0.85

0.90

0.80

0.85

MLcuraLy

0.75

0.80

0.70

SingleSet

0.75

SingleSet

0.65

FedAvg

FedAvg

0.60

FedRoD

MOON

0.70

?

MOON

FedRoD

FedBN

FedBN

0.55

COPA

0.65

COPA

Fed-CO2

Fed-CO2

0.50

0.2

0.4

0.6

1.0

2.0

4.0

0.2

0.4

0.6

1.0

2.0

4.0

Training number sample ratio

Training number sampleratio

As for the effect of client number, we employed dataset CIFAR10 with Dirichlet setting ( α = 0 . 3 ) to conduct a series of experiments on benchmark methods with client number N ∈ { 50 , 100 , 150 , 300 , 500 } . The experimental results shown in Table 13 demonstrate that our FedCO 2 consistently outperforms a bunch of SOTA methods in every case even if the client number is substantial. The results in Table 13 and Fig. 6 prove that our Fed-CO 2 owns higher scalability than benchmark methods.

F.5 Personalized Parts in the Online model

In order to enhance the adaptation to local data distribution, it is necessary to personalize certain critical parameters in the online model to preserve the essential offline specialized knowledge. In this section, we investigated and analyzed which specific parts of our online model should be personalized for optimal performance in FL with feature skew. The most common strategy is to personalize Batch Normalization layers or the Classifier Head. Therefore in our Fed-CO 2 approach, we conducted three experiments with different personalized components in the online model to compare and determine the optimal strategy. The results for three datasets, namely DomainNet, Office-Caltech10, and Digits, are illustrated in Fig. 7. It is evident from the observations that personalizing the Classification Head has a negative impact on almost every sub-dataset, resulting in lower classification accuracy. Based on the results, we can conclude that personalizing Batch Normalization layers in our online model is the best strategy to adapt to feature shifts.

(a)

(b)

Performance Comparison on Office-Caltech10

1.0 -

1.0 -

1.0 -

-8°0

0.8 -

0.8 -

0.6 -

0.6 -

0.6 -

0.4 -

0.4 -

0.4 -

0.2 -

0.2 -

0.2 -

0.0-

Qipart

InfographPainting

Quickdraw

Real

Sketch

Average

0.0 -

Amazon

Caltech

DSLR

WebCam

Average

0.0 -

MNIST

SVHN

USPS

SynthDigits

MNIST-M

Average

Personalize BN

Personalize Head

Personalize BN and Head

Personalize BN

Personalize Head

Personalize BN and Head

Personalize BN

Personalize Head

Personalize BN and Head

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors focus on tackling the feature skew issue, and introduces a new approach that utilizes knowledge distillation. Specifically, the proposed method transfers the knowledge between the FL online model and the local offline model to enable them to have both domain invariant and domain specific knowledge.\n\nSTRENGTHS:\nStrength:\n\n1. The proposed approach of mutual learning between the offline and online models is novel.\n\n2. Theoretical analysis of the proposed method is provided (though they are all in the appendix)\n\n3. Experimental results on multiple datasets with both label distribution skew and feature skew demonstrate promising outcomes.\n\nWEAKNESSES:\n1. Lack of important baselines: \n\t\n\ta. Ditto: Since Ditto is a very classic pFL method, it should be included as the baseline. \n\n\tb. FedMask [1]: There is no hypernetwork based method that directly personalizing parts of the model parameters. FedMask is one of the SOTA method, which learns distinct binary masks for the last several layers of the local models, and aggregates the masks with an intersection operation. \n\n2. Lack of some experiments:\n\n\tThe proposed method actually contains three components, intra-client, inter-client knowledge distillation, and the fuse of online, offline models. In the ablation study, the authors focus on the experiment of the previous two parts, and the fuse part is only mentioned in the label distribution skew setting (Figure 4 in the appendix). How about conducting experiments to evaluate the performance of the online and offline models in the feature distribution skew setting, i.e., the DomainNet dataset?\n\nQUESTIONS:\n1. Refer weakness 1 and weakness 2. \n\n2. Regarding the order of two knowledge transfers:  Why is Intra-Client Knowledge Transfer executed before Inter-Client Knowledge Transfer？\n\n3. Why is the theoretical analysis part only included in the appendix? In my opinion, at least the main theorem should be presented in the main body of the paper, while the assumptions, lemmas, and proofs could be included in the appendix.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: SUMMARY:
The paper proposes a personalized federated learning (PFL) method to handle both label and feature distribution skew.Specifically, each clients holds a partially perosnalized model (personalization happens at the level of batch-normaizliation layers), and a fully personalized local model.  The paper conducts a series of numerical sumilations to evaluate the performance of the proposed method and to compares with state-of-the-art PFL methods.

STRENGTHS:
- The paper is overall well-written and easy to follow. 
- The proposed method shows promising performance and outperforms competitors in most cases. 
- The proposed method does not induce any communication overhead compared to standard federated aggregation operations since the fully personalized offline models are not uploaded to the server.

WEAKNESSES:
- The paper does not provide stong explantion to motivate the proposed method. In particular, it is unclear where (6) and (7) come from, and why they take this specific form. 
- The proposed method induces a memory overhead in comparison to  standard federated averaging.
- Apart from proposing a PFL with good performance, the paper does not provide any insights helping to understand personalization.
- The paper does not provide any theoretical guarantees. 
- I think that (3) is not accurate. Maybe a division by 2 in the RHS of (3) is missing.

QUESTIONS:
- Could you please provide more intuition on the particular choice of the loss function in (6) and (7)?
- Could you please double check (3)?

LIMITATIONS:
- The proposed method induces a memory overhead in comparison to  standard federated averaging.
- The paper does not provide any theoretical guarantees.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents Fed-CO2, a novel framework for Federated Learning (FL), a distributed learning paradigm that allows multiple clients to collectively learn a global model without sharing their private data. FL's effectiveness depends heavily on the quality of the data used for training. Data heterogeneity issues, like label distribution skew and feature skew, can significantly impact FL's performance. Traditionally, most studies have focused on dealing with label distribution skew, while a few recent ones have started addressing feature skew. These forms of data heterogeneity have been studied separately and have not been integrated within a unified FL framework. \n\nFed-CO2 aims to address this gap by developing a framework that handles both label distribution skew and feature skew. The framework utilizes a cooperation mechanism between the Online and Offline models: the online model learns general knowledge shared among all clients, and the offline model is trained locally to learn each individual client's specialized knowledge.  To further enhance model cooperation in the presence of feature shifts, the authors design an intra-client knowledge transfer mechanism to reinforce mutual learning between the online and offline models. Additionally, they introduce an inter-client knowledge transfer mechanism to enhance the models' domain generalization ability. Through extensive experiments, Fed-CO2 outperforms a wide range of existing personalized federated learning algorithms, handling label distribution skew and feature skew, both individually and collectively.\n\nSTRENGTHS:\n- The cooperation mechanism between the online and offline models is a novel and effective way to deal with data heterogeneity issues. This mechanism allows for general learning shared among all clients, as well as specialized learning unique to each client.\n- The Fed-CO2 framework represents a comprehensive approach to addressing data heterogeneity in Federated Learning. It doesn't just focus on label distribution skew or feature skew, but rather addresses both these challenges within a unified framework.\n- The extensive experiments show that the Fed-CO2 model performs better than a wide range of current personalized federated learning algorithms. This indicates the effectiveness and potential of the Fed-CO2 model.\n\nWEAKNESSES:\n- The model may be complex to implement, especially in environments where the number of clients is very large, due to the dual model structure (online and offline) and the knowledge transfer mechanisms.\n- Despite being designed to handle heterogeneity, the performance of the framework is still reliant on the quality of the data used for training, which may not always be guaranteed, especially in a federated setting.\n- The paper does not clearly articulate how the proposed method scales with increased data, feature dimensions, or numbers of clients. Understanding the scalability of the method is crucial in real-world applications.\n\nQUESTIONS:\n- What are the computational costs associated with the proposed intra-client and inter-client knowledge transfer mechanisms?\n- How does Fed-CO2 handle potential security and privacy concerns associated with inter-client knowledge transfer?\n- How does Fed-CO2 handle scenarios where the label distribution skew and feature skew are extreme?\n- Can the framework handle other types of data heterogeneity beyond label distribution skew and feature skew?\n\nLIMITATIONS:\nSee weaknesses and questions", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces an algorithm to address both label and feature heterogeneity in federated learning. The algorithm entails that each client learns two models: one that has personalized batch normalization parameters and otherwise global parameters, and another that is fully personalized. To learn the parameters, a two-stage knowledge distillation procedure is proposed, where in the first stage knowledge is distilled between the two models on each client, then in the second stage knowledge is distilled across clients. Experiments study the performance of the proposed model against a variety of baselines on image classification in the presence of label heterogeneity, feature heterogeneity, and both label and feature heterogeneity.\n\nSTRENGTHS:\n1. The experimental results show convincing improvement over a variety of relevant baselines in multiple settings with different types of heterogeneity. The procedure is rigorous. \n\n2. The related works are well-discussed.\n\n3. The paper addresses a relevant topic.\n\nWEAKNESSES:\n1. The writing can be improved. There are numerous sentences with improper grammar/semantics, e.g. “Federated learning for Data Heterogeneity of Feature Skew”.  The model parameter notation is not consistent as both \\theta= \\{\\phi, \\eta\\} and \\theta = \\{\\phi, \\xi\\} are used. Label and feature heterogeneity/skew are never defined. The motivation for personalizing specifically the batch normalization parameters is not clear whatsoever, and seems to be simply because it worked well in [10].\n\n2. Label and feature heterogeneity/skew are never defined.\n\n3. The proposed method entails learning many parameters per client. This is likely to degrade performance in settings with few samples per client, but the effect of the number of samples on algorithm performance is not investigated.\n\n4. The experiments are limited to image classification on relatively simple datasets.\n\n5. The idea to encourage retaining global information locally by adding a regularizer that penalizes loss on local samples of local models with concatenated heads of all other clients is interesting, but introduces additional communication and privacy concerns (due to all clients sharing each of their heads with each other) that are not addressed.\n\n-----------\n\nPost-rebuttal: I have raised my overall score from 4 to 5, and Contribution score from 2 to 3, please see comment below.\n\nQUESTIONS:\nn/a\n\nLIMITATIONS:\nSee Weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-15"}