Title: Is Heterogeneity Notorious? Taming Heterogeneity to Handle Test-Time Shift in Federated Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Federated Learning

Test-Time Shift

Contrastive Learning
Abstract: Federated learning (FL) is an effective machine learning paradigm where multiple clients can train models based on heterogeneous data in a decentralized manner without accessing their private data. However, existing FL systems undergo performance deterioration due to feature-level test-time shifts, which are well investigated in centralized settings but rarely studied in FL. The common non-IID issue in FL usually refers to inter-client heterogeneity during training phase, while the test-time shift refers to the intra-client heterogeneity during test phase. Although the former is always deemed to be notorious for FL, there is still a wealth of useful information delivered by heterogeneous data sources, which may potentially help alleviate the latter issue. To explore the possibility of using inter-client heterogeneity in handling intra-client heterogeneity, we firstly propose a contrastive learning-based FL framework, namely FedICON, to capture invariant knowledge among heterogeneous clients and consistently tune the model to adapt to test data. In FedICON, each client performs sample-wise supervised contrastive learning during the local training phase, which enhances sample-wise invariance encoding ability. Through global aggregation, the invariance extraction ability can be mutually boosted among inter-client heterogeneity. During the test phase, our test-time adaptation procedure leverages unsupervised contrastive learning to guide the model to smoothly generalize to test data under intra-client heterogeneity. Extensive experiments validate the effectiveness of the proposed FedICON in taming heterogeneity to handle test-time shift problems.

Is Heterogeneity Notorious? Taming Heterogeneity to Handle Test-Time Shift in Federated Learning

Abstract

Federated learning (FL) is an effective machine learning paradigm where multiple clients can train models based on heterogeneous data in a decentralized manner without accessing their private data. However, existing FL systems undergo performance deterioration due to feature-level test-time shifts, which are well investigated in centralized settings but rarely studied in FL. The common non-IID issue in FL usually refers to inter-client heterogeneity during training phase, while the test-time shift refers to the intra-client heterogeneity during test phase. Although the former is always deemed to be notorious for FL, there is still a wealth of useful information delivered by heterogeneous data sources, which may potentially help alleviate the latter issue. To explore the possibility of using inter-client heterogeneity in handling intra-client heterogeneity, we firstly propose a contrastive learning-based FL framework, namely FedICON, to capture invariant knowledge among heterogeneous clients and consistently tune the model to adapt to test data. In FedICON, each client performs sample-wise supervised contrastive learning during the local training phase, which enhances sample-wise invariance encoding ability. Through global aggregation, the invariance extraction ability can be mutually boosted among inter-client heterogeneity. During the test phase, our test-time adaptation procedure leverages unsupervised contrastive learning to guide the model to smoothly generalize to test data under intra-client heterogeneity. Extensive experiments validate the effectiveness of the proposed FedICON in taming heterogeneity to handle test-time shift problems.

1 Introduction

Federated learning (FL) is a promising machine learning paradigm that enables multiple clients or devices to collaboratively train models without sacrificing their data privacy [1, 2]. Practical FL frameworks usually suffer from notorious non-IID (independent and identically distributed) local data distributions, which may lead to severe performance degradation during the test phase [3, 4].

A.ribute Shi,

…

…

…

…

Client 1

Client m

Intra-Client Heterogeneity

Inter-Client Heterogeneity

Original

Distribu>on

Test Data

(a)

MNIST

SVHN

USPS

Synth

MNIST-M

10

20

30

40

50

60

70

80

90

Avg Test Acc (%)

Standalone

Inter-Client Heter

(b)

To improve the performance of the trained model on local data distribution, there are a variety of personalized federated learning (PFL) methods based on different techniques, i.e. , regularization [5, 6], meta-learning [7, 8], knowledge distillation [9, 10], to help the model adapt to local datasets [11].

In existing PFL systems, it is commonly assumed that the training and test datasets within a client share the same or similar distribution. However, this assumption often does not hold in real-world applications where test samples may encounter natural variations or corruptions. Although test-time distribution shifts have received increasing attention in centralized machine learning paradigms [12, 13, 14, 15], the negative effects of these shifts in PFL systems are largely underexplored. Take federated chest X-ray classification [16, 17] among multiple hospitals as an example, while most existing studies acknowledge that different hospitals can have varying data distributions due to diverse patient populations (e.g., races, ages, etc.), they overlook the possibility of variation within a single hospital. For instance, there may be multiple X-ray machines from different vendors, each with its own unique internal parameters (e.g., exposure time, spectrum, etc.). Similar issues can be found in various other domains, such as autonomous vehicles and video surveillance. Consequently, an efficient and effective solution to handle test-time shifts is necessary to achieve satisfactory performance.

The extensively studied heterogeneous distribution in FL, also known as non-IID issue, mainly refers to the inter-client heterogeneity where clients own heterogeneous training data during their federated training phase [2, 3, 18]. Differently, the aforementioned test-time distribution shift during the inference stage of PFL refers to the heterogeneity within a client, which can be defined as intra-client heterogeneity . As illustrated in Fig. 1(a), the test data of a client can differ from its training data in various aspects. This discrepancy results in unexpected features in the model optimized on the training data, ultimately leading to suboptimal performance. In this case, a natural question that arises is: what is the difference between the test-time shift in centralized settings and that in FL? Interestingly, we find that the answer lies in the unique inherent inter-client heterogeneity in FL compared with centralized settings. While inter-client heterogeneity is often considered to be a challenge for FL, in this paper, we pinpoint that heterogeneous sources can also convey a wealth of useful information that potentially helps alleviate test-time shift problems in FL systems. A benefit of this heterogeneity is that it enables the global model to capture robust features that are less sensitive to both inter-client and intra-client shifts. We firstly discover that this information has the potential to help alleviate test-time shift problems in FL systems, as it enables the global model to capture robust features that are less sensitive to both inter-client and intra-client shifts. Through our motivated experiments (Fig. 1(b)), we observe that FL under inter-client heterogeneity can improve the ability of clients to handle test-time shift problems compared to training on homogeneous data. In this way, we turn the problem of inter-client heterogeneity as a blessing in solving test-time shift.

Motivated by the above findings, we introduce Fed erated Learning via I nvarian C e Extracti O n and Shari N g (FedICON), a contrastive learning-based FL approach that enables clients to learn the robust features on heterogeneous data. Our theme is to fully take advantage of the benefit brought by inter-client heterogeneity during training time, and then, at the test phase, tune the local model of each client on its private test data to conquer the intra-client heterogeneity . During the training phase, clients first carry out local representation learning under the guidance of label information

so that invariant information underlying each training sample, as well as the class it belongs to, is extracted. At the global aggregation stage of federated training, the parameters of feature encoders in all clients are aggregated to further boost the ability of clients to learn invariant information among inter-client heterogeneity. Then, during the test phase, we introduce self-supervision signals to the unlabeled test samples and smoothly tune the feature encoder based on the acquired knowledge in identifying invariant properties between diverse training data and shifted test data. Notably, we leverage contrastive learning paradigms for both federated training and local test-time adaptation, enabling the trained model to well adapt to test data distribution.

Our contributions are summarized as follows:

We are the first to comprehensively investigate feature-level test-time shifts in FL. We formulate the problem from the perspective of inter-client heterogeneity and intra-client heterogeneity and propose corresponding experimental settings tailored for this problem.

We propose a novel representation learning-based federated learning algorithm, termed FedICON, that addresses the test-time shift problem by taming inherent heterogeneity in FL for invariance extraction and sharing.

We conduct extensive experiments on benchmark datasets to measure the performance of FedICON in handling feature-level test-time shift problems. The results demonstrate that FedICON achieves significant improvement in terms of test-time adaptation ability.

2 Related Work

Personalized Federated Learning. Personalized federated learning (PFL) aims to enable each client to learn a personalized model that performs well on its local dataset [19, 11]. Existing PFL methods are developed on the basis of various techniques. A branch of work utilizes model decoupling to separate the local model into the globally shared part and personally trained part to realize different levels of personalization [20, 21, 22, 23, 24]. [25, 26, 27, 5, 28] augment the original objective function with additional regularization terms to induce personalization via representation alignment, knowledge transfer, bi-level optimization, etc. [8, 7] incorporate meta-learning to generate the generalized model which can be further optimized as a personalized model. Existing personalized federated learning methods can only personalize the model to local training data, lacking the ability to generalize to the test dataset [29].

Federated Learning with Feature-Level Distribution Shifts. Non-IID issue is the core challenge in both FL and PFL [30]. A common type of non-IID data in real-world scenarios is known as feature-level distribution shifts, including attribute shift, covariate shift, and domain shift [31, 32]. For instance, in FL, the training and test images can be collected in different time period, differnt locations, and even from different photographic equipment. [33] leverages mutual information theory and diversity augmentation to mitigate the negative effects brought by attribute shift and domain shift. [18] proposes to use local batch normalization to alleviate the feature-level distribution shift problems. Enlightened by the idea of domain generalization methods, [34, 35] design different knowledge transfer strategies to learn high-quality domain consensus knowledge across shifted clients. Based on feature-level distribution shift setting, [36, 37] aim to alleviate the inter-domain knowledge forgetting problem and performance degradation problem on unseen clients/domains, respectively.

Test-Time Adaptation. Test-time adaptation methods aim to improve the model performance on out-of-distribution data by adapting the learned model to test data [38, 39]. Several methods propose to adjust the model by minimizing the pre-defined test-time entropy in various forms [40, 15]. [41] presents a classifier-only adjustment module for test domain adaptation, and [12] further shows that selectively fine-tuning part of the model can address different types of distribution shifts. [13, 42] tune the model on unlabelled test data via self-supervised learning objectives. [43] further proposes to use more informative self-supervised learning to solve the test-time performance degradation in the presence of severe distribution shifts. The performance of existing centralized test-time adaptation methods can be hindered due to the inherent heterogeneity in FL during the training phase, which further affects the test adaptation process.

Federated Representation Learning. Representation learning focuses on extracting higher-level representations from raw data using deep neural networks [44]. Combining it with the federated

…

!

!

!

"

Client 1

Local Invariance Extrac0on

"

!

!

!

Feature

Encoder

"

"

!

"

Client m

!

Feature

Encoder

Test Phase on a Client

#

!

#

"

…

Training Phase

Feature

Encoder

Contrastive

"

!

"

"

Contrastive

"

!

"

"

Contras+ve

Global Invariance Sharing

…

!

!

!

"

Client 1

Local Invariance Extrac0on

"

!

!

!

Feature

Encoder

"

"

!

"

Client m

!

Feature

Encoder

Test Phase on a Client

#

!

#

"

…

Training Phase

Feature

Encoder

Contrastive

"

!

"

"

Contrastive

"

!

"

"

Contras+ve

Global Invariance Sharing

(a) Federated Training under Inter-Client Heterogeneity

(b) Test-Time Adaptation

learning framework enables distributed clients to learn powerful representation and thus benefit their downstream tasks [45, 46, 47, 48]. [49] introduces a generalized federated self-supervised learning framework and provides in-depth empirical studies on the effect of key components in federated representation learning. [50] incorporates model decoupling and representation learning for better PFL. [51, 52] aim to learn a common representation model without supervision under the FL framework. Both of them utilize Siamese networks as their local representation module.

3 Notations and Preliminaries

3.1 Notations

Considering m clients in an FL system, P i and Q i are training and test data distributions of the i -th client ( i ∈ { 1 , · · · , m } ) respectively and D P i and D Q i are the corresponding training dataset and test dataset. A sample-label pair is denoted as ( x , y ) where x is the raw input data sample and y is its corresponding label.

3.2 Inter-Client Heterogeneity

̸

̸

For arbitrary client i and client j in an FL system, the inter-client heterogeneity refers to varying P i and P j with respect to their raw data sample x and/or label y . For example, in the cases where P i ( y ) = P j ( y ) while P i ( x | y ) = P j ( x | y ) , there exists an inter-client label distribution shift. The case where P i ( x ) = P j ( x ) while P i ( y | x ) = P j ( y | x ) can lead to inter-client feature distribution shift, which is common in real-world FL systems [18, 33, 36].

3.3 Intra-Client Heterogeneity

Most existing works on FL consider the heterogeneity of training datasets D P i across clients while overlooking the intra-client heterogeneity P i = Q i . [31] examines test distribution in robust FL benchmarks, combining label-level and feature-level shifts. To better understand their distinct impacts, we propose to study these shifts separately. Here, we focus on varying P i and Q i with respect to x rather than y , i.e., addressing the feature-level test-time shift problem in FL. Incorporating real-world scenarios, we thoroughly examine inter-client feature distribution shifts, including covariate shift, attribute shift, and domain shift. We assume that all clients have access to their own labelled training data when participating in the federated training process and test the performance of their personally tuned model on their own test set.

̸

Challenges. It is a challenging task to deal with both inter-client and intra-client heterogeneity in an FL system. Although there is useful information underlying heterogeneous clients, how to effectively extract it and how to further leverage it to benefit the test procedure remain unexplored. Hence, there is a need for feasible solutions to tame inherent inter-heterogeneity to handle the test-time intra-client heterogeneity in FL.

4 Federated Learning via Invariance Extraction and Sharing

Both inter-client heterogeneity and intra-client heterogeneity arise from shifts in data distribution. Our core idea is to take advantage of diverse data distributions across clients so that robust features that remain invariant can be learned to handle intra-client shifts. Motivated by this, we propose Fed erated Learning via I nvarian C e Extracti O n and Shari N g (FedICON) as shown in Figure 2. During the federated training phase, to extract invariant information of each class within the local dataset, we design a local invariance extraction scheme by employing contrastive learning in a supervised manner (Sec. 4.1.1). Then, we perform global invariance sharing under inter-client heterogeneity to further boost the invariance extraction ability for each client (Sec. 4.1.2). During the test phase, each client tunes the personalized model to adapt to its local test set (Sec. 4.2).

Optimization Objective. Since each client has varying training and test data distribution, the global objective of the personalized FL system across m clients can be formulated as a two-stage optimization process,

where Eq. 1 and Eq. 2 are the optimization objective during the federated training phase and test phase, respectively. In Eq. 1, L i, sup and θ i are the local supervised loss function and learnable model parameters for client i , respectively. N P is the total number of training samples among all clients. In Eq. 2, θ i is initialized as θ ′ i and further optimized on the test data of client i via unsupervised learning loss function L i, us . The eventual purpose of the framework is to efficiently learn invariant knowledge under inter-client heterogeneity during the training phase so that the intra-client heterogeneity can be alleviated via further optimization during the test phase.

4.1 Federated Training under Inter-Client Heterogeneity

Compared with the test-time shift problem in centralized settings, there are inherent heterogeneous source data in the federated learning framework, which is of great potential to utilize to deal with the shifted test data. Although the inter-client heterogeneous data are not directly accessible due to privacy concerns [53], their underlying invariant information can be captured by each client during local invariance extraction procedure and then globally shared across all clients to further enhance the invariance extraction ability.

4.1.1 Local Invariance Extraction.

To capture invariant information from samples within the same class, we design a local invariance extraction procedure based on supervised contrastive learning. As shown in Fig. 2(a), samples are first fed into a feature encoder to generate the representation for each sample. Then, a contrastive learning scheme is applied to the representations under the guidance of their corresponding labels so that the class-level invariant information can be extracted during the local training phase.

To make the representation ability of the feature encoder more robust, we add an augmentation module during the data pre-processing stage that incorporates some perturbations to raw samples. In this way, there are various views of the same sample fed into the feature encoder to strengthen its sample-based discrimination ability. Concretely, for the i -th client, given a batch of input samples, we generate two random augmented views for each sample. The total size of the augmented local dataset is 2 |D i | . For a specific sample x , we use A ( x ) to denote the local dataset without x .

Taking advantage of the powerful representation ability brought by contrastive learning frameworks, we force the feature encoder to discriminate the augmented samples from different classes. Given an input sample x , the output of the feature encoder is

where θ i is the learnable parameter of the feature encoder. We denote z ( x ) as z x for short. Samples from the same class are forced to be closer, and samples from different classes are forced to be farther.

We perform local training with the loss function as below,

where z a is the representation of sample a ∈ A ( x ) , and τ is the temperature adjusting the tolerance for representation difference. P ( x ) is the set of samples that have the same labels with x and can be represented as

By training the local feature encoder in a contrastive manner, the local invariances can be directly learned from the augmented data and be encoded to the feature encoder. Even if there is inter-client heterogeneity among different clients, their invariance encoding abilities over heterogeneous data are universal, which makes it possible for clients to mutually boost their performance on invariance extraction in a federated framework.

4.1.2 Global Invariance Sharing.

To tame the inter-client heterogeneity during the federated training process, there should be a uniform information channel that is globally shared across clients and continuously provides meaningful knowledge to boost local invariance extraction ability. Considering the fact that the feature encoder introduced above is trained to capture and encode the invariance, we can employ the parameters in the feature encoder as the information carrier and perform global invariance sharing at the server side. These invariant robust properties captured by the feature encoder are shared within different augmented views of a specific sample and different samples from the same class.

Similar to the standard federated aggregation process, for each client i ∈ { 1 , 2 , · · · , m } , the learnable parameter θ i of its feature encoder is sent to the server. After receiving parameters from all the clients, the server aggregates them as

and sends ¯ θ back to each client as their initialized parameter in a new round. Since the feature encoder of each client is trained to extract invariance efficiently rather than classify the samples correctly, the parameter aggregation is less likely to be affected by inter-client heterogeneity. Global invariance sharing with respect to the feature encoder incorporates the invariance encoding abilities learned upon heterogeneous source data, which reversely provides sufficient homogeneous information for each client and increases the robustness of local representation learning. To employ the model trained alternately via local invariance extraction and global invariance sharing, we further train a linear classifier on top of the fixed feature encoder [54, 55]. The linear classifier can be trained in a pure local manner or trained in a federated learning manner. In this paper, we conduct multiple federated training steps at first and then personalize the classifier for each client.

4.2 Test-Time Adaptation to Handle Intra-Client Heterogeneity

In Section 4.1, we have developed components that leverage the inherent inter-client heterogeneity in FL in order to handle the test-time intra-client heterogeneity. To keep the losses of the training and test phase unified and make the whole process consistent, we adopt a self-supervised contrastive learning scheme to carry out test-time adaptation, which is also a natural and efficient way to generalize the trained model to the unlabeled and shifted test set.

Concretely, for client i , the trained feature encoder is further tuned by optimizing the following unsupervised loss function on its local test set D Q i .

where z ′ x is the representation of the augmented sample of x . Similar to the training phase, A ( x ) here also denotes the local test dataset without x . The main difference between Eq. 4 and 7 is whether the label information is accessible or not. Apart from this, Eq. 4 and 7 are consistently formulated and

can share the same optimization strategy. During the training phase, the contrastive learning process can be conducted under the guidance of label information, which encourages the feature encoder to output well-aligned representations for samples within the same class. Then, based on the robust representations learned among inter-client heterogeneity, the feature encoder can further improve the generalization ability on the test set in an unsupervised manner, which preserves the contrastive learning mode in the training phase.

The test scenario can be easily extended to the online manner where the test samples arrive one by one. First, a set of training samples is selected from the training set, replacing the A ( x ) in Eq. 7. Then, after performing the optimization for each test sample, we replace one randomly selected training sample in A ( x ) with this test sample for future test samples. This transfers the whole pipeline into an online manner and allows the model to gradually adapt to the distribution of test data.

We present our proposed FedICON in detail in Algorithm 1 in Appendix C.

5 Experiments

5.1 Experimental Setup

Datasets and test-time shift settings. We evaluate our proposed method on the following four benchmark datasets: Digit-5 [56], Office-10 [57], DomainNet dataset [58], and CIFAR10 [59]. Digit5, Office-10, and DomainNet dataset are three benchmark dataset consisting of multiple sub-datasets from different domains. We use them to simulate the covariate shift and domain shift during the test phase. We use CIFAR10 and its variant CIFAR10.1 [60, 61] to simulate the attribute shift . Details about the test-time shift settings can be found in Appendix A.2.

Covariate shift. We follow the setting in [31] by leveraging 15 common corruptions in CIFAR10-C [62], e.g. , Gaussian Noise, Brightness, etc.

Attribute shift. The training data of all clients are from CIFAR10, while the test data of all clients are from CIFAR10.1 shifted from CIFAR10, e.g. , the most common object in the airplane class can change from airliner to fighter.

Domain shift. We follow the training data setting in [18], where clients own data from different domains with heterogeneous appearances. We follow the leave-one-out principle to select one domain as the test domain and the other domains as the training domains. The training data of each client is from a single domain.

Baselines and implementation. We first compare our proposed method with two kinds of baselines: 1) Local, where each client only performs local training without any communication; FL/PFL methods including 2) FedAvg [63], 3) FedAvg-FT that fine-tunes the global model trained by FedAvg on local dataset, 4) FedProx [64], 5) FedRep [22], 6) FedBN [18], 7) FedTHE [31]. We also compare with three classical centralized test-time adaptation (TTA) methods, including 1) Tent [38], 2) EATA [40], and 3) T3A [41]. In practice, we implement these TTA methods with the FedAvg framework. Details about the baseline methods and their parameter settings can be found in Appendix A.3.

For Digit-5, Office-10, and CIFAR10, we use simple CNN architecture as the backbone model, while for DomainNet dataset, we use ResNet18 as the backbone model. Details about the model of each dataset can be found in Appendix A.4. We set the local training batch size as 32, and an SGD optimizer with weight decay rate 5e-4 and learning rate 0.01 for cross-entropy loss and 0.005 for contrastive loss. The temperature in Eq. 4 and Eq. 7 is 0.1 for Digit-5, Office-10, CIFAR10, 10 for DomainNet dataset. The size of representation used for contrastive learning is 64 for Digit-5, Office-10, and CIFAR10, and 512 for DomainNet dataset. The default number of local epochs is E = 5 and the number of communication rounds is 100. We implement all the methods using PyTorch and conduct all experiments on NVIDIA Tesla V100 GPU. Detailed settings can be found in Appendix A.

5.2 Performance Comparison

Performance on various test-time shifts. Table 1 and Table 2 report the test accuracy of FedICON and baselines on Digit-5 dataset under covariate test-time shift and Office-10 dataset under domain test-time shift, respectively. The results are reported in mean(standard error) format over three independent runs.

The results suggest that our proposed FedICON outperforms existing FL/PFL methods and test-time adaptation methods and shows a strong ability to deal with both covariate and domain test-time shift problem. It should be noticed that although FedRep and FedBN are two state-of-the-art solutions in personalized federated learning, they usually fail to perform well on the shifted test sets. The main reason for that is both of them have a special focus on local training data distribution and accordingly adjust part of their model to better fit local training data distribution, leading to a more biased model than the generic model during test phase.

Table 3 summarizes the results of test accuracy on CIFAR10 dataset under attribute test-time shift. We follow the setting of the naturally shifted test in [31]. The results in Table 3 indicate that our proposed FedICON is an effective solution to extract key invariant implicit information underlying both the training and test set, even though there are explicitly different properties conveyed by the training and test set. We also report how the performance varies on the original test set and the attribute test-time shift scenarios. The results show that FedICON has the smallest performance difference among all federated methods, which reflects the robustness of FedICON against attribute shift during test time. Additional experimental results on other datasets under various test-time shifts can be found in Appendix B.

Effect of varying numbers of participating clients. To test the performance of FedICON in a larger FL system, we increase the number of clients from 5 to 80 on Digit-5 dataset under covariate test-time shift setting. The total number of training samples is kept constant. As the number of clients becomes larger, the number of samples in each client becomes fewer. As shown in Fig. 3(a), both FedAvg and FedICON have a decreased performance when there are more clients as a result of the further divergence among participating clients. Although performance degradation is inevitable, FedICON still stably outperforms FedAvg in FL systems at different scales, demonstrating the scalability of FedICON and its ability to adapt to a large-scale FL system.

Effect of varying numbers of local epochs. To test how local epoch number affects the test performance, we vary the number of local epochs and report the results of FedAvg and FedICON on Digit-5 dataset under covariate test-time shift setting. As shown in Fig. 3(b), the performance of

5

10

20

40

80

Num of Clients

20

30

40

50

60

70

Avg Test Acc (%)

FedAvg

FedICON

(a)

1

3

5

7

9

Num of local epochs

50

52

54

56

58

60

62

64

66

68

Avg test acc

FedAvg

FedICON

(b)

FedAvg keeps dropping as the number of local epochs increases, while FedICON partially enjoys the benefit of more local training, thus more communication-efficient than FedAvg.

Extending to unsupervised scenarios. The main framework of our proposed method is based on representation learning, which makes it flexible to migrate the whole learning paradigm into the unsupervised representation learning manner. During the local invariance extraction and global invariance sharing stages, we remove the label information. The feature encoder is purely trained in a self-supervised manner and the label information is accessible only when training the classifier. As is shown in Table 4, the extended version of our method FedICON, termed as FedICONus , achieves the runner-up performance on Digit-5 and Office-10 under covariate test-time shift.

5.3 Ablation Studies

We carry out ablation studies to see the effect of each key technical component in FedICON. We alternatively remove the local invariance extraction, global invariance sharing and test-time adaptation procedure to see how these components affect the performance in handling test-time shift issues in FL. It is worth noting that when the local invariance extraction component is removed, the local training objective changes from the supervised contrastive loss to the conventional classification loss ( i.e. , cross-entropy), and the global invariance sharing component is in charge of sharing the

Table 4: Test accuracy on Digit-5 and Office-10 dataset under covariate test-time shift. The best and runner-up results are highlighted with bold and underline, respectively.

learnable parameters of the local model. In Table 5, we can observe that both global invariance sharing and test-time adaptation play an important role in improving the ability to deal with test-time shift problems. The effect of global invariance sharing is more notable because it is the core step to explore the inherent inter-client heterogeneity and broadcast the invariant encoding knowledge to benefit their local learning. The test-time adaptation brings about 1% to 5% improvement compared with the variants without this component, verifying the effectiveness of consistent representation alignment via a self-supervised learning scheme during the test phase.

6 Conclusion

In this paper, we aim to address the challenging but unexplored feature-level test-time shift problems in FL. We take the first initiative to thoroughly investigate the inherent inter-client heterogeneity in FL and propose to use it to deal with the test-time intra-client heterogeneity . We propose a novel algorithm, namely FedICON, where invariant information underlying heterogeneous data is extracted in a contrastive learning manner during the federated training phase and further used to deal with test-time shifts in a client. On our proposed benchmarks tailored for this problem, extensive experiments are conducted to show the superiority of FedICON compared with baseline methods.

This research has great potential in solving test-time shift problems in real-world FL systems. Providing theoretical guarantees for our work and discussion on privacy and security can be further explored in future studies. Considering the evident motivation, the novel setting and the effectiveness of our method, we believe this work is intriguing for future studies.

References

Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM TIST , 2019.

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.

Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv:1806.00582 , 2018.

Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. In ICLR , 2019.

Canh T Dinh, Nguyen Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. In NeruIPS , 2020.

Filip Hanzely and Peter Richtárik. Federated learning of a mixture of global and local models. arXiv:2002.05516 , 2020.

Yihan Jiang, Jakub Koneˇ cn` y, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv:1909.12488 , 2019.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. In NeruIPS , 2020.

Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In NeurIPS , 2020.

Yiqing Shen, Yuyin Zhou, and Lequan Yu. Cd2-pfed: Cyclic distillation-guided channel decoupling for model personalization in federated learning. In CVPR , 2022.

Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Yoonho Lee, Annie S Chen, Fahim Tajwar, Ananya Kumar, Huaxiu Yao, Percy Liang, and Chelsea Finn. Surgical fine-tuning improves adaptation to distribution shifts. In International Conference on Learning Representations , 2023.

Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning , pages 9229-9248. PMLR, 2020.

Yixin Liu, Kaize Ding, Jianling Wang, Vincent Lee, Huan Liu, and Shirui Pan. Learning strong graph neural networks with weak information. In ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD) , 2023.

Marvin Zhang, Sergey Levine, and Chelsea Finn. MEMO: Test time robustness via adaptation and augmentation. Advances in Neural Information Processing Systems , 35:38629-38642, 2022.

Dinh C Nguyen, Quoc-Viet Pham, Pubudu N Pathirana, Ming Ding, Aruna Seneviratne, Zihuai Lin, Octavia Dobre, and Won-Joo Hwang. Federated learning for smart healthcare: A survey. ACM Computing Surveys (CSUR) , 55(3):1-37, 2022.

Guodong Long, Tao Shen, Yue Tan, Leah Gerrard, Allison Clarke, and Jing Jiang. Federated learning for privacy-preserving open innovation future on digital health. In Humanity Driven AI: Productivity, Well-being, Sustainability and Partnership , pages 113-133. Springer, 2021.

Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. In ICLR , 2020.

Chen Zhang, Yu Xie, Hang Bai, Bin Yu, Weihong Li, and Yuan Gao. A survey on federated learning. Knowledge-Based Systems , 216:106775, 2021.

Hong-You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In ICLR , 2022.

Xin Dong, Sai Qian Zhang, Ang Li, and HT Kung. Spherefed: Hyperspherical federated learning. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXVI , pages 165-184. Springer, 2022.

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In ICML , 2021.

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

Paul Pu Liang, Terrance Liu, Liu Ziyin, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. NeurIPS , 2020.

Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang, Wenchao Xu, and Feijie Wu. Parameterized knowledge transfer for personalized federated learning. NeruIPS , 34, 2021.

Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. FedProto: Federated prototype learning across heterogeneous clients. In AAAI , 2022.

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

Filip Hanzely, Slavomír Hanzely, Samuel Horváth, and Peter Richtárik. Lower bounds and optimal algorithms for personalized federated learning. NeruIPS , 33:2304-2315, 2020.

Weiming Zhuang, Chen Chen, and Lingjuan Lyu. When foundation model meets federated learning: Motivations, challenges, and future directions. arXiv preprint arXiv:2306.15546 , 2023.

Chen Chen, Yuchen Liu, Xingjun Ma, and Lingjuan Lyu. CalFAT: Calibrated federated adversarial training with label skewness. In NeurIPS , 2022.

Liangze Jiang and Tao Lin. Test-time robust personalization for federated learning. In International Conference on Learning Representations , 2023.

Weiming Zhuang, Yonggang Wen, Lingjuan Lyu, and Shuai Zhang. MAS: Towards resourceefficient federated multiple-task learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 23414-23424, 2023.

Zhengquan Luo, Yunlong Wang, Zilei Wang, Zhenan Sun, and Tieniu Tan. Disentangled federated learning for tackling attributes skew via invariant aggregation and diversity transferring. In ICML , 2022.

Haozhe Feng, Zhaoyang You, Minghao Chen, Tianye Zhang, Minfeng Zhu, Fei Wu, Chao Wu, and Wei Chen. Kd3a: Unsupervised multi-source decentralized domain adaptation via knowledge distillation. In ICML , 2021.

Junming Chen, Meirui Jiang, Qi Dou, and Qifeng Chen. Federated domain generalization for image recognition via cross-client style transfer. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision , pages 361-370, 2023.

Wenke Huang, Mang Ye, and Bo Du. Learn from others and be yourself in heterogeneous federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10143-10153, 2022.

Honglin Yuan, Warren Richard Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? In International Conference on Learning Representations , 2022.

Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test-time adaptation by entropy minimization. In International Conference on Learning Representations , 2021.

Yixin Liu, Kaize Ding, Huan Liu, and Shirui Pan. Good-d: On unsupervised graph out-ofdistribution detection. In Proceedings of the Sixteenth ACM International Conference on Web Search and Data Mining , pages 339-347, 2023.

Shuaicheng Niu, Jiaxiang Wu, Yifan Zhang, Yaofo Chen, Shijian Zheng, Peilin Zhao, and Mingkui Tan. Efficient test-time model adaptation without forgetting. In International conference on machine learning , pages 16888-16905. PMLR, 2022.

Yusuke Iwasawa and Yutaka Matsuo. Test-time classifier adjustment module for model-agnostic domain generalization. Advances in Neural Information Processing Systems , 34:2427-2440, 2021.

Alexander Bartler, Andre Bühler, Felix Wiewel, Mario Döbler, and Bin Yang. Mt3: Meta testtime training for self-supervised test-time adaption. In International Conference on Artificial Intelligence and Statistics , pages 3080-3090. PMLR, 2022.

Yuejiang Liu, Parth Kothari, Bastien Van Delft, Baptiste Bellot-Gurlet, Taylor Mordan, and Alexandre Alahi. TTT++: When does self-supervised test-time training fail or thrive? Advances in Neural Information Processing Systems , 34:21808-21820, 2021.

Bram van Berlo, Aaqib Saeed, and Tanir Ozcelebi. Towards federated unsupervised representation learning. In Proceedings of the third ACM international workshop on edge systems, analytics and networking , pages 31-36, 2020.

Duc Bui, Kshitiz Malik, Jack Goetz, Honglei Liu, Seungwhan Moon, Anuj Kumar, and Kang G Shin. Federated user representation learning. arXiv preprint arXiv:1909.12535 , 2019.

Yue Tan, Guodong Long, Jie Ma, Lu Liu, Tianyi Zhou, and Jing Jiang. Federated learning from pre-trained models: A contrastive learning approach. Advances in Neural Information Processing Systems , 35:19332-19344, 2022.

Yue Tan, Yixin Liu, Guodong Long, Jing Jiang, Qinghua Lu, and Chengqi Zhang. Federated learning on non-iid graphs via structural knowledge sharing. In Proceedings of the AAAI conference on artificial intelligence , volume 37, pages 9953-9961, 2023.

Jingtao Li, Lingjuan Lyu, Daisuke Iso, Chaitali Chakrabarti, and Michael Spranger. MocoSFL: enabling cross-client collaborative self-supervised learning. In The Eleventh International Conference on Learning Representations , 2022.

Weiming Zhuang, Yonggang Wen, and Shuai Zhang. Divergence-aware federated selfsupervised learning. In International Conference on Learning Representations , 2020.

Jaehee Jang, Heoneok Ha, Dahuin Jung, and Sungroh Yoon. Fedclassavg: Local representation learning for personalized federated learning on heterogeneous neural networks. In Proceedings of the 51st International Conference on Parallel Processing , pages 1-10, 2022.

Weiming Zhuang, Xin Gan, Yonggang Wen, Shuai Zhang, and Shuai Yi. Collaborative unsupervised visual representation learning from decentralized data. In Proceedings of the IEEE/CVF international conference on computer vision , pages 4912-4921, 2021.

Fengda Zhang, Kun Kuang, Zhaoyang You, Tao Shen, Jun Xiao, Yin Zhang, Chao Wu, Yueting Zhuang, and Xiaolin Li. Federated unsupervised representation learning. arXiv preprint arXiv:2010.08982 , 2020.

Chen Chen, Lingjuan Lyu, Han Yu, and Gang Chen. Practical attribute reconstruction attack against federated learning. IEEE Transactions on Big Data , 2022.

Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. NeruIPS , 33, 2020.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML , pages 1597-1607. PMLR, 2020.

Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Learning to generate novel domains for domain generalization. In ECCV , 2020.

Boqing Gong, Yuan Shi, Fei Sha, and Kristen Grauman. Geodesic flow kernel for unsupervised domain adaptation. In CVPR , 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.

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

Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do cifar-10 classifiers generalize to cifar-10? 2018. https://arxiv.org/abs/1806.00451.

Antonio Torralba, Rob Fergus, and William T. Freeman. 80 million tiny images: A large data set for nonparametric object and scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence , 30(11):1958-1970, 2008.

Dan Hendrycks and Thomas G Dietterich. Benchmarking neural network robustness to common corruptions and surface variations. In ICLR , 2019.

H Brendan McMahan, Eider Moore, Daniel Ramage, et al. Communication-efficient learning of deep networks from decentralized data. In AISTATS , 2017.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. MLSys , 2020.

Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In ICLR , 2019.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a scheme to handle the test-time shift at each FL client. To achieve this, during training, a contrastive learning is adopted to extract invariant information within the same class (class-level invariant information) at each client. The updated models at the clients are also aggregated via FedAvg. During testing, at each client, the model is further fine-tuned using the test dataset and the corresponding unsupervised loss function.\n\nSTRENGTHS:\n1. The paper is generally well written and easy to understand. The figures also well describe the high-level concepts of the proposed methods.\n\n2. This paper proposes strategies for both training and testing during FL, to handle the test-time shift issue, which is practical.\n\n3. The authors consider various datasets for experiments including DomainNet.\n\nWEAKNESSES:\n1. Section 4.1.1: First of all, in 4.1.1, the authors are trying to extract invariant information within the same class, which shares the same concept with domain generalization that learns domain-invariant features during training. However, comparison with these methods are lacking. For example, one can apply any domain generalization schemes for local training to learn domain-invariant features, and then aggregate the models according to 4.1.2. Note that in domain generalization literatures, “domain shift” generally includes covariate shift, attribute shift, and domain shift the authors mentioned in this paper. Existing domain generalization methods can generally handle all these shifts.\n\n2. Related to the comment above, there are many works that focus on adopting contrastive loss for domain generalization. It is not clear what is the authors’ novelty compared to the existing works on contrastive learning loss in Section 4.1.1\n\n3. Section 4.1.2: This part is just about FedAvg. Hence, in my point of view, Section 4.1 is about performing updates locally to extract invariant information, and then aggregating the models. However, I do not find the novelty compared to applying existing works in Section 4.1\n\n4. Section 4.2: During testing, a test dataset needs to be available in each client for fine-tuning. This may not be practical when test samples arrive one-by-one. Moreover, focusing on a specific client (and given the well-trained models), what is the novelty of Section 4.2 compared to other test-time adaptation methods?\n\n5. Experiments: For experiments, as mentioned in my first comment above, the important baselines are missing. Especially, the authors’ scheme is the only method that simultaneously adopts (i) generalization during training and (ii) test-time adaptation at inference. To validate the effectiveness of the idea in Section 4.1.1, the authors need to consider schemes that apply domain generalization methods instead of Section 4.1.1 and combine it with Section 4.2, and compare it with FedICON. Similarly, to validate the effectiveness of the idea in Section 4.2, the authors may combine other test-time methods (e.g., Tent, T3A) with the idea in Section 4.1.1, and compare the scheme with FedICON.\n\n6. Experimental setup: It is not clear how data are distributed across clients. Moreover, during testing, does the proposed method assumes that full test dataset is accessible at each client to finetune the model? \n\nOverall, this paper can be viewed as a combination of domain generalization learning (or contrastive learning) and test-time adaptation in FL, but I find the novelty of the paper to be relatively weak, since there are no clear descriptions showing the difference compared to various existing works and there are no experiments comparing with the corresponding baselines.\n\nQUESTIONS:\nAnother question is that, what happens if the classifier is not personalized but shared by aggregation? \n\nMoreover, what is the experimental setup of Fig. 1(b) showing the advantage of heterogeneous data compared to homogeneous setup?\n\nRegarding other questions, please refer to the weakness above.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose FedICON, which uses inter-client heterogeneity to handle intra-client heterogeneity. During training, FedICON uses contrastive learning locally to extract invariant class-conditional information, and performs global invariance sharing under inter-client heterogeneity. During testing, the feature extractor is adapted with contrastive learning. Extensive experiments validate the effectiveness of FedICON.\n\nSTRENGTHS:\n- The paper makes clear definitions of inter-client and intra-client heterogeneities. Meanwhile, it is insightful to propose that inter-client heterogeneity in FL can be used to tackle intra-client heterogeneity. \n- FedICON has good performance on a variety of experimental settings. Especially, the author tried different types of heterogeneities.\n\nWEAKNESSES:\n- Global invariance sharing. It is unclear why sharing the model parameter of feature extractor can achieve global invariance sharing. Since FedICON does not compare representations from different clients, FedICON only guarantee local invariance on each client, but not global invariance. For example, consider a binary classification problem, a feature extractor might map positive/negative examples on client 0 to (0, 1), (0, 3), and map positive/negative examples on client 1 to (0, -2), (0, -4). In this way, although all positive samples on client 0 are mapped to the same representation (0, 1), positive samples on client 0 and 1 are mapped to different representations (0, 1) and (0, -2). \n- Since this paper focuses on how inter-client heterogeneity can help tackling intra-client heterogeneity, I believe the author may consider emphasizing the type of distribution shifts for each heterogeneity in the main text, instead of in the supplementary material.\n\nQUESTIONS:\n- Could you provide more information about whether FedICON learns local and global invariances? E.g., by numerical results of the variance of representations, or visualization (not required since it could be hard during rebuttal)\n- Other questions that I am interested: \n  - In line 190, the authors mentioned that they generate two random augmenteed views for each sample. Usually this is for constructing positive pairs for unlabelled data (e.g., in self-supervised learning). However here the setting is supervised. Also in equation (3) and (4) it seems the augmented data is not fed into the feature encoder. I am confused whether the augmentation is used during training.\n\nLIMITATIONS:\nThe limitations are clearly stated.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo deal with the feature-level test-time shift problem in federated learning, this paper proposes to leverage the inherent heterogeneity across clients based on a contrastive learning method, named FedICON. Clients acquire invariance encoding ability on heterogeneous source data and further boost the performance in the test phase. Experiments on various datasets show the effectiveness of FedICON compared with baseline methods.\n\nSTRENGTHS:\n1. It is a novel idea to leverage heterogeneous properties underlying FL systems to deal with the commonly existing test-time shift problem within a client. Compared with centralized test-time adaptation studies, the heterogeneous problem in the training phase is a characteristic of FL. \n2. Sufficient experiments are conducted and the main results show that FedICON has achieved improvement for almost all participating clients under various feature-level shifts. \n3. The paper is easy to follow and well-organized. The authors provide clear and novel definitions of inter-client and intra-client heterogeneity. The motivation for utilizing inter-client heterogeneity to deal with intra-client heterogeneity is well demonstrated.\n\nWEAKNESSES:\n1. The concepts of test-time shift and intra-client heterogeneity seem to refer to the same thing, which is not well clarified in the main paper. \n2. It is better to give a better explanation of the fundamental difference between label-level and feature-level test-time shifts, which makes it easier for the readers to understand why the shifts should be studied separately.\n\nQUESTIONS:\n1. Is it possible to deploy the proposed method in an online manner, e.g., the online test deployment in [1], to handle the dramatic shifts during the test phase?\n[1] FedTHE: Test-Time Robust Personalization for Federated Learning. In ICLR 2023. \n\n2. Given a local model, how to choose an appropriate feature space for contrastive learning?\n\nLIMITATIONS:\nYes. The authors have addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on the federated learning (FL) scenario with the test-time shift problem, which is a practical yet challenging research topic. With empirical study, the authors find that the inter-client heterogeneity in personalized FL can be further leveraged to build a robust FL framework against the test-time shift problem. Motivated by this, the authors propose a novel FL approach termed FedICON that uses train-time and test-time contrastive learning to boost the ability of local models to learn invariant information among inter-client heterogeneity caused by test-time shifts. Extensive experiments show the effectiveness of FedICON on FL scenarios with various test-time shifts.\n\nSTRENGTHS:\nStrengths\n1. Valuable research problem. This paper addresses a new research problem, namely federated learning with feature-level test-time shifts. The exploration of this research topic has significant potential for application in various industry scenarios. The authors discuss the problem from the perspective of inter- and intra-client heterogeneity, providing valuable insights for future works.\n\n\n2. Novel method. The proposed method, FedICON, is innovative in addressing the FL problem with test-time shifts. The incorporation of supervised and unsupervised contrastive learning during the training and testing phases to acquire invariant knowledge is a meaningful approach.\n\n\n3. Extensive experiments. The authors have conducted an extensive set of experiments to evaluate the effectiveness of the proposed method. They have considered multiple scenarios, such as covariate shift, domain shift, attribute shift, etc, which enhance the persuasiveness of the experiments. The experimental results demonstrate a significant performance improvement achieved by FedICON.\n\nWEAKNESSES:\nWeaknesses:\n\n1. Lack of discussion for the motivated experiment. In Sec 1, the authors present evidence that highlights the positive impact of inter-client heterogeneity in FL when addressing test-time shift problems. However, the paper lacks a thorough discussion of the underlying reasons behind this observation. It would be beneficial for the authors to provide deeper insights and further explanations regarding the reasons or deductions supporting this finding.\n\n2. Design motivation needs to be more clear. Contrastive learning plays a crucial role in the proposed method. However, the authors have not sufficiently discussed why they choose contrastive learning to provide the pivot supervision signals instead of other supervised or unsupervised objectives, such as cross-entropy or self-reconstruction. It is important for the authors to provide a clearer explanation and justification for this choice, allowing readers to understand the design motivation of using contrastive learning in the proposed method.\n\n3. There are some instances where unexpected periods (\".\") appear at the end of subheadings in certain subsections, such as Sec 3.1, 3.2, and 3.3. The authors should carefully review the paper to identify and correct any typos or format errors throughout the manuscript.\n\nQUESTIONS:\nAs I pointed out in \"Weaknesses\", more discussions are expected:\n1. The discussion for the motivated experiment.\n2. The motivation for using contrastive learning in FedICON.\n\nLIMITATIONS:\nThe authors have discussed the limitations of this work.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}