SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fed-  $\mathbf{CO}_{2}$ : Cooperation of Online and Offline Models for Severe Data Heterogeneity in Federated Learning
Zhongyi Cai
ShanghaiTech University
caizhy@shanghaitech.edu.cn
Ye Shi*
shiye@shanghaitech.edu.cn
RIKEN Center for Advanced Intelligence Project
Jingya Wang
wangjingya@shanghaiitech.edu.cn
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-  $\mathrm{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-  $\mathrm{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.
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
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 domain-invariant, 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 intra-client 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  $\mathrm{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,  $\mathrm{Fed - CO_2}$  can handle severe label distribution skew effectively. To enhance model performance in the presence of severe feature skew,  $\mathrm{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  $\mathrm{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  $\mathrm{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<sup>1</sup>.
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,
<sup>1</sup>Our codes are publicly available at https://github.com/zhyczy/Fed-C02.
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- $\mathrm{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,  $\mathrm{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  $\mathrm{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  $\mathrm{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.
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_{i}^{k},y_{i}^{k}):k\in \{1,\ldots ,m_{i}\} \}$ ,  $i\in \{1,\dots ,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}(\cdot)$  with parameters  $\theta_{i}$  for each client according to information in all clients:
$$
\min  _ {\left\{\theta_ {i} \right\} _ {i = 1} ^ {N}} \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {m _ {i}} l \left(F _ {i} \left(\theta_ {i}; x _ {i} ^ {j}\right), y _ {i} ^ {j}\right), \tag {1}
$$
where  $l(\cdot, \cdot)$  is a sample-wise loss function,  $F_{i}$  is constructed by a feature extractor  $f_{i}(\cdot)$  with parameters  $\eta_{i}$  and a classifier  $C_{i}(\cdot)$  with parameters  $\phi_{i}$ . Accordingly,  $\theta_{i} = \{\eta_{i}, \phi_{i}\}$ .
(a) Cooperation of Online and Offline Models
(c) Inter-Client Knowledge Transfer
(b) Intra-Client Knowledge Transfer
Figure 1: The  $\mathrm{Fed - CO_2}$  framework. Fig. 1(a) shows an overview of the cooperation between the online and offline models in  $\mathrm{Fed - CO_2}$ . Figures 1(b) and 1(c) respectively illustrate the intra-client and inter-client knowledge transfer mechanisms on FL with feature skew.
3.2 Cooperation of Online and Offline Models
As shown in Fig. 1(a), we propose Fed-  $\mathrm{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_{i}^{\mathrm{on}}$ , and a locally-trained model, referred to as the offline model  $F_{i}^{\mathrm{off}}$ . With two models in each client  $F_{i} = \{F_{i}^{\mathrm{on}}, F_{i}^{\mathrm{off}}\}$ , 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  $\theta_{i}^{\mathrm{on}}$  into  $\{\theta_{p,i}^{\mathrm{on}},\theta_g^{\mathrm{on}}\}$ , where  $\theta_{p,i}^{\mathrm{on}}$  denotes all BN layer parameters and  $\theta_{g}^{\mathrm{on}}$  denotes other parameters. Model aggregation is done on the server to obtain the shared  $\widetilde{\theta}_{g}^{\mathrm{on}}$  among all clients with the formula:
$$
\widetilde {\theta} _ {g} ^ {\text {o n}} = \sum_ {i = 1} ^ {N} \frac {1}{N} \theta_ {g, i} ^ {\text {o n}}. \tag {2}
$$
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}$ :
$$
F _ {i} \left(\theta_ {i}; x\right) = F _ {i} ^ {\text {o n}} \left(\theta_ {p, i} ^ {\text {o n}}, \theta_ {g} ^ {\text {o n}}; x\right) + F _ {i} ^ {\text {o f f}} \left(\theta_ {i} ^ {\text {o f f}}; x\right). \tag {3}
$$
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\neq 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  $\bar{F}_i^{\mathrm{on}}$  and  $\bar{F}_i^{\mathrm{off}}$ , respectively, with frozen parameters  $\bar{\theta}_i^{\mathrm{on}}$  and  $\bar{\theta}_i^{\mathrm{off}}$ . 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:
$$
\theta_ {i, s} ^ {\text {o n}} = \theta_ {i, s - 1} ^ {\text {o n}} - \alpha \cdot \nabla_ {\theta_ {i, s - 1} ^ {\text {o n}}} K L \left(\bar {F} _ {i} ^ {\text {o f f}} \left(\bar {\theta} _ {i} ^ {\text {o f f}}; x _ {k}\right), F _ {i} ^ {\text {o n}} \left(\theta_ {i, s - 1} ^ {\text {o n}}; x _ {k}\right)\right), \tag {4}
$$
$$
\theta_ {i, s} ^ {\text {o f f}} = \theta_ {i, s - 1} ^ {\text {o f f}} - \alpha \cdot \nabla_ {\theta_ {i, s - 1} ^ {\text {o f f}}} K L \left(\bar {F} _ {i} ^ {\text {o n}} \left(\bar {\theta} _ {i} ^ {\text {o n}}; x _ {k}\right), F _ {i} ^ {\text {o f f}} \left(\theta_ {i, s - 1} ^ {\text {o f f}}; x _ {k}\right)\right), \tag {5}
$$
where  $KL(\cdot, \cdot)$  denotes the Kullback-Leibler (KL) divergence and  $\alpha$  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-  $\mathrm{CO}_{2}$ , each client uploads its classifier  $C_{i}^{\mathrm{off}}$  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:  $\{\overline{C}_j^{\mathrm{off}}\}_{j = 1}^N$  with parameter set  $\{\overline{\phi}_j^{\mathrm{off}}\}_{j = 1}^N$  and are transmitted to each client in the next round. In communication round  $t$ , each client has its own personalized model  $F_{i} = \{F_{i}^{\mathrm{on}}, F_{i}^{\mathrm{off}}\}$  and the downloaded knowledge-transfer classifier set  $\{\overline{C}_j^{\mathrm{off}}\}_{j = 1}^N$ . 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:
$$
L _ {g e n} \left(x _ {k}, y _ {k}\right) = \sum_ {j \neq i} L _ {C E} \left(\bar {C} _ {j} ^ {\text {o f f}} \left(\bar {\phi} _ {j} ^ {\text {o f f}}; f _ {i} \left(\eta_ {i, t}; x _ {k}\right)\right), y _ {k}\right), \tag {6}
$$
Combined with its classification loss to adapt to local private data, the total loss in local adaptation phase for each model is:
$$
L = L _ {C E} \left(x _ {k}, y _ {k}\right) + \mu \cdot L _ {\text {g e n}} \left(x _ {k}, y _ {k}\right), \tag {7}
$$
where  $\mu$  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-  $\mathrm{CO}_{2}$  algorithm.
Here, we provide convergence analyses, which compare our  $\mathrm{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  $\mathbb{D}_i = \{(\mathbf{x}_j^i,y_j^i)\} _j = 1^M$  with  $\mathbf{x}_j^i\in \mathbb{R}^d$  and  $y_{j}^{i}\in \mathbb{R}$ . For convenience, we adopt the same data distribution assumption as in [10]: