Is Heterogeneity Notorious? Taming Heterogeneity to Handle Test-Time Shift in Federated Learning
Authors: Yue Tan
Chen Chen
Weiming Zhuang
Xin Dong
Lingjuan Lyu
Guodong Long
Venue: NeurIPS 2023 poster (2023)
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
Yue Tan"
Chen Chen
University of Technology Sydney yue. tanestudent uts edu. au
Sony AI
chena. www..hhheenccccoooo com
Weiming Zhuang
Xin Dong
Sony AI
Sony AI xindong@sony. com
weiming. www.ngeson..ccooommo com
Lingjuan Lyu
Guodong Long
Sony AI 1 Lingjuan. ww@.onvvvccccooooooooooomo
University of Technology Sydney guodong. 1aigg..omm long@ut edu. au
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 perfor- mance deterioration due 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 informa- tion delivered by heterogeneous data sources, which may potentially help alleviate the latter issue. To explore the possibility of using inter-client heterogeneity handling intra-client heterogeneity, firstly propose contrastive learning based FL framework, namely FedICON, to capture invariant knowledge among heteroge- neous 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-wis 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 general- ize test data under intra-client heterogeneity. Extensive experiments validate the effectiveness of the proposed FedICON in taming heterogeneity to handle test-time shift problems.
Introduction
Federated learning (FI.) a promising machine learning paradigm that enables multiple clients or devices to collaboratively train models without sacrificing their data privacy [ Practical FL frameworks usually suffer from notorious non-III (independent and identically distributed) local data distributions, which may lead to severe performance degradation during the test phase 3
Work done during Yue Tan's internship at Sony AL. Corresponding Author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Inter Client Heterageneity
Training Data
Driginal
Distribution
90
60-
Standalone
70
Inter-Client Client Heter
Covariate Shift
60
50-
40
Attribute Shift
& 20-
30
Domain Shift
10
MNIST
SVHN
USPS
Synth
MNSTM MM
Client 1
Client m
(a)
(
Figure I: (a) Illustration of our setting where both inter-client heterogeneity and intra-client heterogeneity exist in an FL framework. (b) Compared with training on data from a single domain, the improved robustness to test-time shift is achieved when training on heterogeneous data sources, indicating that the inherent inter-client heterogene- ity in FL. has positive effect in handling test-time shift problems. Detailed settings can he found in Appendix A.
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
meta-learning 7 g. knowledge distillation [9] 10 to help the model adapt to local datasets
In existing PFL systems, it is commonly assumed that the training and test datasets within 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 largely underexplored. Take federated chest X-ray classification [16] [7] 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.
B 14 15]. the negative effects of these shifts in PFL systems
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 client, which can be defined as intra-client heterogeneity. As illustrated in Fig. I(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, natural question that arises is: what is the difference berween 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. 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. l(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 Federated Learning via InvarianCe ExtractiOn and ShariNg (FedICON), 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 helerogeneity and intra-client heterogeneity and propose corresponding experimental settings tailored for this problem.
We propose novel representation learning- -baedd federated learning algorithm, termed FedICON, that addresses the rest-time shift problem by taming inherent heterogeneity in FL for invariance extraction and sharing.
We conduct extensive experiments benchmark datasets to measure the performance of FedICON in handling feature- 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 personalized model that performs well on its local dataset [19 Existing PFL methods are developed on the basis various techniques. 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 (8 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 personalized model. Existing personalized federated learning methods can only personalize the model to local training data, lacking the ability
generalized model which can be further optimized
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-llD 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 be collected in different time period, differnt locations, and even from different photographic equipment. [33] leverages mutual information theory and diversity augmentation
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-qualit, domain consensus knowledge across shifted clients. Based on
feature-level distribution shift setting, [36] 37 aim 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 of of-diistribttonn 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
[5].
41]
presents 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. [3 42
tune the model on unlabelled test data via self- self-super learning objectives. [43] further proposes to use more informative eelf-uuper learning to solve the test-time performance degradation 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 [4]. Combining it with the federated
3

---

Training Phase
Global Invariance Sharing
Client
Client m
X1
2
X1
11
Test Phase on Client
Feature Encoder
Feature Encoder
Contrastive
Contrastive 2
Vi
4 Contrastive
Feature Encoder
Local Invoriance Extroction
(a) Federated Training under Inter-Client Heterogeneity
(b) Test-Time Time Adaptation
An overview of the proposed FedICON. (a) Foderated training phase under inter-client heterogeneity. Each client conducts representation wise local invariance extraction during local training and performs parameter- wise global invariance sharing during global aggregation. (b) Test phase on a single client. Clients carry out independent test-time adaptation to generalize to their local test sets.
Figure 2:
learning framework enables distributed clients learn powerful representation and thus benefit their downstream tasks 45
[49] introduces a generalized federated self-supervised learning framework and provides in-depth empirical studies the effect of key components in federated representation learning. 50
46
47
48
incorporates model decoupling and representation learning for better PFL. [51] 52) aim learn common representation model without supervision under the FL framework. Both of them utilize Siamese networks as their local representation module.
Notations and Preliminaries
3.1 Notations
Considering clients in FL system, P. and Qi are training and test data distributions of the i-th client (i {I, m)) respectively and D' and D/ the corresponding training dataset and test dataset. A sample-label pair is denoted as (x, v) where x is the raw input data sample and 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. and P, with respect their raw data sample and/or label y. For example, the cases where (v( * (v) while (x(x)) P; (x y), there exists an inter-client label distribution shift. The case where (x) * Py(x) while Pi x) (y(yx) can lead 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 Di across clients while overlooking the intra-client heterogeneity P: * Q; 31 examines test distribution in robust FL. benchmarks, combining label-level and feature-level shifts. To better understand their distinct impacts, propose study these shifts separately. Here, focus varying P: and Q; with respect 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 challenging task to deal with both inter-client and intra client heterogeneity in an FL system. Although there 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 a need for feasible solutions to tame inherent inter-heterogeneity to handle the test-time intra-client heterogeneity in FL.

---

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 Federated Learning via CevarianCe ExtractiOn and ShariNg (FedICON) as shown in Figure During the federated training phase, to extract invariant information each class within the local dataset, design a local invariance extraction scheme by employing contrastive learning in supervised manner (Sec. 4.I.I. 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.
Optimization Objective. Since each client has varying training and test data distribution, the global objective of the personalized FL system across clients can be formulated as a two-stage optimization process.
min 1 m [Dfl
08 ... ,0m)
I Lisup (0:; (o::
(0,.02.7 .0m) m
NP
07 min Li,us (0::D)). 6g
(2)
where Eq. and Eq. 2 are the optimization objective during the federated training phase and test phase, respectively. In Eq. Ius.sup and 0g are the local supervised loss function and learnable model parameters for client i, respectively. N= is the total number of training samples among all clients. In Eq. [ 0, is initialized 0g and further optimized on the test data of client via unsupervised learning loss function Li,us- The eventual purpose of the framework is to efficiently learn invariant know mewlledge 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 utilize to deal with the shifted test data. Although the inter-client heterogeneous data are not directly accessible due to privacy concerns [5)] 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, design local invariance extraction procedure based on supervised contrastive learning. As shown in F(). samples are first fed into a feature encoder to generate the representation for each sample. Then, 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 samples. In this way, there are various views 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 || For specific sample x, 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 the feature encoder is
=xx) = (x; 0,)
where 0 the learnable parameter of the feature encoder. We denote 2(x) as zx for short. Samples from the same class are forced to be closer, and samples from different classes are forced to be farther.
5

---

We perform local training with the loss function as below,
exp (2x 2p/T)
Li,sup (.::
Σ
log pEp(x)
Σ
P((x]
EaEA(x) exp (2x 2a/T)
(x.y)ED)
where za is the representation of sample a E 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
P(x) {p A(x) Up }x
(5)
By training the local feature encoder in contrastive manner, the local invariances 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, 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.
can employ the parameters in
Similar to the standard federated aggregation process, for each client i (1,2, .mm , the learnable parameter 0; of its feature encoder is sent to the server. After receiving parameters from all the clients, the server aggregates them as
1
m
D NP
0,,
Y11
and sends 0 back to each client as their initialized parameter in a new round. Since the feature encoder of each client trained 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 local representation learning. To employ the model trained alternately via local invariance extraction and global invariance sharing, we further train 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
we have developed components that leverage the inherent inter-client heterogeneity in FL in order 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 carry out test-time adaptation, which is also natural and efficient way to generalize the trained model to the unlabeled and shifted test set.
4.1
Concretely, for client the trained feature encoder is further tuned by optimizing the following unsupervised loss function on its local test set DDS
exp (2x 2x//T - log EaeA(s) exp (Ix Za/"
LiAus (0::D!)
Σ
XED
where = 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 (addd]j 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. Then, after performing the optimization for each test sample, we replace one randomly selected training sample Axx) 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 | 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 CIFARIO [59]. Digit- 5, Office-10, 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 10 and its variant CIFAR1O. 10.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 CIFARIO, while the test data of all clients are from CIFARI0.1 shifted from CIFARIO, 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 [8] where clients data from different domains with heterogeneous appearances. We follow the leave-one-outt principle to select one domain the test domain and the other domains 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 FT that fine-tunes the global model trained by FedAvg on local
dataset, 4) FedProx [64]. 5) FedRep [22], 6) FedBN [18, 7) FedTHE (3.. 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 CIFAR1O, we use simple CNN architecture the backbone model, while for DomainNet dataset, we use ResNet18 as the backbone model. Details about the model of each dataset can be found 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. | 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 CIFARIO. I0, and 512 for DomainNer 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 Table Teppot 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.
7

---

Table 1: Test accuracy on Digit-5 dataset under covariate test-time shift. All results reported in mean(std error) format over three independent runs.
Method
MNIST
SVHN
USPS
Synth
MNIST-M
Local
60.65(4.79)
34,43(1.87)
16.91(4.55)
39.93(3.25)
18.07(3.26)
FedAvg
86.18(0.72)
38.60(1.84)
18.19(2.78)
51.50(1.13)
40.72(1.15)
FedAvg-FT
85.92(1.14)
40.25(1.47)
54.28(1.90)
51.23(0.72)
41.08(1.21)
FedProx
86.45
38.95(2.29)
$7.80(2.52)
52.40(1.32)
40.55(0.48)
FedRop
66. 66.78(6.40)
44.17(3.29)
15.81(5.84)
44.60(2.40)
19.33(3.08)
FedBN
61.42(4.50)
41.38(0.77)
41.73(7 44)
42.77(4.44)
42.35(1.13)
FodTHE
84.52(0.21)
39.00(2221)
58.03(2.96)
51.9((1277)
40.55(1.23)
Tent
86.13(0.56)
43.87(0.81)
67.51(2.40)
50.47(0.12)
47.73(0.94)
EATA
85.37(0.08)
43.65(1.55)
67.12(1.4 (88)
50.70(0.22)
47.58(0.77)
T3A
72.58(2.50)
44.23(1.77)
49.44(5.63)
45.62(0.95)
46.13(1.53)
FedICON
89.67(1.48
45.23(0..3)
72.13
56.130063)
47.98(0.50)
The results suggest that our pro- posed FedICON outperforms exist- ing FL/PFL methods and test-time adaptation methods and shows strong ability to deal with both covari- ate and domain test-time shift prob- lem. It should be noticed that al- though FedRep and FedBN two state-of-the-art solutions in personal- ized federated learning, they usually fail to perform well on the shifted test sets. The main reason for that is both of them have special focus on local training data distribution and accord- ingly adjust part of their model to bet- ter fit local training data distribution, leading to a more biased model than the generic model during test phase.
Table 2: Test accuracy on Office- 10 dataset under domain test-time shift. All results are reported in mean(std error) format over three independent runs.
Test Domain
Method
Amazon
Caltech
DSLR
Webcam
Local
24.07(0.61)
18.62(0.31) 36.46(2.76) 28.81(2.04)
FedAvg
44.91(2.99)
26.27(1.63) 61.46(3.12)
48.21(3.76)
FedAvg-FT FT
41.78(1.23) 24.94(1.07)
55.56(2.41)
46.52(5.72) 72)
FedProx
42.53(1.76) 24.44(0.93)
56.60(1.20)
42.94(3.53)
FedRep FedBN
20.02(1.16)
17.68(1.34)
34.03(5.92)
24.86(1.49)
42.53(2.18)
28.64(0.62)
51.04(5.41)
43.69(1.63)
FedTHE
41.61(1.61)
24.79(0.95)
56.94(2.17)
44.63(1.79)
Tent
38.25(1.92) 25.93(2.57)
54.17(1.04)
42.00(1.98)
EATA T3A
38.25(2.03)
25.73(2.59)
54.86(0.60)
42.75(1.18)
39.76(0.52)
22.77(1.20)
53.82(0.60)
45.76(2.87)
FedICON
50.46(0.78) 33.28(0.67) 67.01(1.59) 48.21(2.67)
Table summarizes the results of test accuracy CIFARIO dataset under altribute test-time shift. We follow the setting of the naturally shifted test in [31. The results in mable(3indicaa 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 FedICON against attribute shift during test time. Additional experimental results on other datasets under various fest-time shifts can be found in Appendix B
Effect of varying numbers of participating clients. To test the performance FedICON in a larger FL system, increase the number of clients from to 80 on Digit-5 dataset under covariate test-time shift setting. The total number 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. (3a 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 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 FedAvg and FedICON on Digit-5 dataset under covariate test-time shift setting. As shown in Fig. 3(b) the performance of
8

---

Table 3: Test accuracy CIFARIO dataset under attribute test-time shift. All results reported in mean(std error) format over three independent runs. represents the difference between accuracy on original test set and that on aftribute test-time shift case.
Nou-IID Dir(0.1)
Nou-IID Dir(I)
Method
Original
Attribute Shift
A
Original
Attribute Shift
A
Local
89.15(1.13)
83.54(0.89)
-5.61
67.60(0.65)
61.61(1.99)
-5.99
FedAvg
64.25(1.42)
50.09(1.36)
-14.16
73.23(0.38)
57.17(0.34)
-16.06
FedAvg-FT
87.82(0.59)
80.95(2.55)
-6.87
78.21(0.68)
65.59(1.46)
-12.62
FedProx
59.79(5.91)
52.25(2.53)
-7.54
72.02(0.37)
61.72(1.29)
-10.3
FedRep
88.73(0.26)
81.72(0.63)
-7.01
79.25(0.59)
66.69(1.22)
-12.56
FedBN
83.09(4.56)
78.14(0.57)
-4.95
77.67(1.07)
66.90(1.31)
-10.77
FedTHE
90,55(0.41)
81.91(0.54)
-8.4
80.49(0.89)
66.25(0.96)
11.224
Tent
91.39(0.77)
85.72(1.24)
-5.67
(8.61(0.63)
71.50(1.29)
-7.11
EATA
91.37(0.78)
85.60(1.16)
-5.77
79.72(0.62)
70.44(1.67)
-7.28
T3A
90.79(0.73)
85.91(1.68)
-4.88
78.80(0.95)
70.93(0.68)
-7.87
FedCCON
90.24(0.52)
86.47(1.27)
-3.77
79.43(0.78)
72.41(1.29)
-7.02
70
FedAug
FedICON
66
FedAvg FediCON
Rug Test AcE (%6
60
62
50-
58
54
20
10
20
40
80
Clients
epochs
(a)
(b)
Figure 3: Test accuracy on Digit-5 under varying numbers clients (leff) and varying numbers of local epochs (right). Both are conducted the covariate test-time shift cases. In (a), the number of clients (m) ranges from 5 to 80. The total number of samples is kept constant, so the number of samples in each client is fewer as m goes larger. In (b), the number of local training epochs ranges from 9.
FedAvg keeps dropping as the number of local epochs increases, while FedICON partially enjoys the benefit of more local training, thus more communicationneefiieenn than FedAvg.
Extending to unsupervised scenarios. The main framework of our proposed method is based rep- resentation learning, which makes it flexible to migrate the whole learning paradigm into the unsupervised rep- resentation learning manner. During the local invari- ance extraction and global invariance sharing stages, we remove the label information. The feature encoder is purely trained in supervsuee manner and the label information is accessible only when training the classi- fier. As shown Table the extended version of our method FedICON, termed as FedICON-us, achieves runner-up performance on Digit-5 and Office-10 under covariate test-time shift.
Table 4:
Test accuracy Digit-5 and Office- 10 dataset under covariate test-time shift, The best and runner-up results are highlighted with bold and underline, respectively.
Method
Digit-5
Office-10
Local
34.00(3.12)
58.07(2.84)
FedAve
55.04(1.00)
57.55(1.46)
FedAvg-FT FedICON
54.55(0.83)
60.45(1.4M 40)
62.23(0.88)
64.96(1 (1.4)
FedICON-ws
61.60(1.45)
63.7 63.79(1.17)
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 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., entssyyn and the global invariance sharing component is in charge of sharing the

---

Table 5: Comparison between the FedICON and its variants. Experiments conducted Digit-5 dataset under covariate test-time shift. The number clients is 5. Each of them owns data samples from a specific domain. Corruptions are applied to their local test set to establish the covariate test-time shift setting. The number of communication rounds is 100.
Design Component
Variants
Local Invariance Extraction
Global Invariance Sharing
Test-Time Adaptation
Test Accuracy
34.00
48.60
55.04
39.92
59.39
59.54
49.67
62.23
learnable parameters of the local model. In Table 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 more notable because 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 self- supervised learning scheme during the test phase.
can observe that both global invariance
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-cliemt heterogeneiry. 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 client. On 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
[1] Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM TIST, 2019.
[2] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mchdi Bennis, Ar- jun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, al. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, (4(1-2):1-210, 2021.
[3] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv: (0t::/806.00000 80580, 2018
[4] Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang, On the convergence of fedavg on non-iid data. In ICLR, 2019.
[5] Canh T Dinh, Nguyen Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. In NeruIPS, 2020.
[6] Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. arXiv:2002.05516, 2020.
[7] Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv:1909.1224 12488, 2019.
10

---

[8] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. In NerulPS, 2020.
[9] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In NeurIPS, 2020.
[10] Yiqing Shen, Yuyin Zhou, and Lequan Yu. Cd2-pfed: Cyclic distillatin--uiieed channel decoupling for model personalization in federated learning. In CVPR, 2022.
[11] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 2022.
[12] 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.
[13] 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.
[14] 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.
[15] Marvin Zhang, Sergey Levine, and Chelsea Finn. MEMO: Test time robustness via adaptation and augmentation. Advances in Neural Information Processing Systems, 14:38629-38644 2022.
[16] Dinh C Nguyen, Quoc-Viet Pham, Pubudu 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.
[17] 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 Al: Productivity, Well-being, Sustainability and Partnership, pages 113-133. Springer, 2021.
[18] 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.
[19] Chen Zhang, Yu Xie, Hang Bai, Bin Yu, Weihong Li, and Yuan Gao. A survey on federated learning. Knowledge-Boud Based Systems, 216: 216:106775 2021.
[20] Hong-You You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In ICLR, 2022.
[21] Xin Dong, Sai Qian Zhang, Ang Li, and HT Kung. Spherefed: Hyperspherical federated learning. In Computer Vision-ECCC 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXVI, pages 165-184. Springer, 2022.
[22] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In ICML, 2021.
[23] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arxiv:191..00818. 2019.
[24] Paul Pu Liang, Terrance Liu, Liu Ziyin, Ruslan Salakhutdinov, and Louis-Philipee Morency. Think locally, act globally: Federated learning with local and global representations. NeurlPS, 2020.
[25] Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang. Wenchao Xu, and Feijie Wu. Parameterized knowledge transfer for personalized federated learning. NeruIPS, 34, 2021.
[26] 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.

---

[27] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In ICML, pages 6357-6368. PMLR, 2021.
[28] Filip Hanzely, Slavomír Hanzely, Samuel Horváth, and Peter Richtárik. Lower bounds and optimal algorithms for personalized federated learning. NerulPS, 33:2304-23155, 2020.
[29] Weiming Zhuang, Chen Chen, and Lingjuan Lyu. When foundation model meets federated learning: Motivations, challenges, and future directions. arXiv preprint arXiv:2306. 15546, 2023.
[30] Chen Chen, Yuchen Liu, Xingjun Ma, and Lingjuan Lyu. adversarial training with label skewness. In NewrlPS, 2022.
CalFAT: Calibrated federated
[31] Liangze Jiang and Tao Lin. Test-time robust personalization for federated learning. In Interna- tional Conference on Learning Representations, 2023.
[32] Weiming Zhuang, Yonggang Wen, Lingjuan Lyu, and Shuai Zhang. MAS: Towards resource- efficient federated multiple-task learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 23414-23424, 2023.
[33] 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.
[34] 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.
[35] 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.
[36] Wenke Huang, Mang Ye, and Bo Du Learn from others and be yourself in heterogeneous federated learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10143-10153, 2022.
[37] Honglin Yuan, Warren Richard Morningstar, Lin Ning. and Karan Singhal. What do mean by generalization in federated learning? In International Conference on Learning Representations. 2022.
[38] 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.
[39] Yixin Liu, Kaize Ding, Huan Liu, and Shirui Pan. Good-d: On unsupervised graph out-of- distribution detection. In Proceedings of the Sixteenth ACM International Conference Web Search and Data Mining, pages 339-347, 2023..
[40] Shuaicheng Niu, Jiaxiang Wu, Yifan Zhang, Yaofo Chen, Shijian Zheng, Peilin Zhao, and Mingkui Tan. Efficient test-time model adaptation without forgetting. In International confer- ence machine learning, pages 16888-16905. PMLR. 2022.
[41] 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.
[42] Alexander Bartler, Andre Bühler, Felix Wiewel, Mario Döbler, and Bin Yang. Mt3: Meta test- time training for self-supervised test-time adaption. In International Conference on Artificial Intelligence and Statistics, pages 3080-3090. PMLR, 2022.
[43] 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 2021.
in Neural Information Processing Systems,
12

---

[44] Bram van Berlo, Aaqib Saeed, and Tanir Ozcelebi. Towards federated unsupervised repre- sentation learning. In Proceedings of the third ACM international workshop on edge systems, analytics and networking. pages 31-36, 2020.
[45] Due Bui, Kshitiz Malik, Jack Goetz, Honglei Liu, Seungwhan Moon, Anuj Kumar, and Kang G Shin. Federated user representation learning. arXiv preprint arXiv: 12535, 2019.
[46] 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.
[47] 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-99611 2023.
[48] 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.
[49] Weiming Zhuang, Yonggang Wen, and Shuai Zhang. Divergence-aware federated self- supervised learning. In International Conference on Leurning Representations, 2020.
[50] Jaehee Jang, Heoneok Ha, Dahuin Jung, and Sungroh Yoon. Fedelassavg: Local representation learning for personalized federated learning on heterogeneous neural networks. In Proceedings of the 51st International Conference Parallel Processing, pages 1-10, 2022.
[51] Weiming Zhuang. Xin Gan, Yonggang Wen, Shuai Zhang, and Shuai Yi. Collaborative unsuper- vised visual representation learning from decentralized data. Proceedings of the IEEF/CV international conference computer vision, pages 4912-4921, 2021.
[52] 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.
[53] Chen Chen, Lingjuan Lyu, Han Yu, and Gang Chen. Practical attribute reconstruction attack against federated learning. IEEE Transactions on Big Data, 2022.
[54] Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Neru/PS, 33, 2020.
(55] Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, pages 1597-1607. PMLR, 2020.
[56) Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Learning to generate novel domains for domain generalization. In ECCV, 2020.
[57] Boqing Gong, Yuan Shi, Fei Sha, and Kristen Grauman. Geodesic flow kernel for unsupervised domain adaptation. In CVPR, 2012.
[58] Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEFE/CVF International Conference on Computer Vision, pages 1406-1415, 2019.
[59] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[60] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do cifar-10 classifiers generalize to cifar- 10? 2018 httpss//arxxv.org/abbb http//806.00444111
[61] Antonio Torralba, Rob Fergus, and William T. Freeman. 80 million tiny images: A large data set for nomparametric object and scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 19:1))):119999 2008
[62] Dan Hendrycks and Thomas G Dietterich. Benchmarking neural network robustness to common corruptions and surface variations. In ICLR, 2019.
13

---

[63] H Brendan McMahan, Eider Moore. Daniel Ramage, et al. Communication-efian learning of deep networks from decentralized data. In AISTATS, 2017.
[64] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. MLSys, 2020.
[65] Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In ICLR, 2019.
14


Reviews:
[{"review_id": "NoTRyqJKuL", "reviewer": "Reviewer_9n3S", "summary": "This 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.", "questions": "As 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", "limitations": "The authors have discussed the limitations of this work. ", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Strengths\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.", "weaknesses": "Weaknesses:\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Lack of discussion for the motivated experiment.** Thanks for the valuable comments. In the paper, we empirically prove and illustrate the fact that inter-client heterogeneity in FL can help alleviate the test-time shift problem. As for the underlying reason for that, the naturally-existing heterogeneity among FL clients potentially improves the generalization ability of the model when data of diverse feature distributions contribute to the learned model. Moreover, the inherent heterogeneity allows the model to capture invariance among heterogeneous feature distributions by exploring the common feature/representation properties. We will add more discussion and explanation regarding the reasons and insights in the revised version. \n\n**More explanation about the motivation.** Thanks for the valuable comments. We utilize contrastive learning to provide pivot supervision signals because contrastive learning is designed to learn representations that are invariant to different transformations of the same instance which is aligned with our target to extract local invariance. As for the other types of objectives or learning pipelines, some of them lack discriminative ability, e.g., self-reconstruction, which makes them less useful for distinguishing between different instances. Also, some of them lack the ability to extend to unsupervised learning scenarios, e.g., cross-entropy, making them less flexible during the inference stage. We will add more explanation about our motivation to use contrastive learning in the revised version. \n\n**Some typos.** Thanks for the valuable comments. We will correct these typos in the revised version. "}, {"review_id": "B2IJbYQzNm", "reviewer": "Reviewer_CniS", "summary": "To 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. ", "questions": "1. 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?", "limitations": "Yes. The authors have addressed the limitations.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. 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. ", "weaknesses": "1. 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Difference between test-time shift and intra-client heterogeneity.** Thanks for the valuable comments. The intra-client heterogeneity we defined in this paper refers to the test-time shift issue, especially in federated learning. The concept of test-time shift is commonly used in general machine learning, e.g., both centralized and decentralized frameworks. Since test-time shift is a more well-known challenging issue recognized by the community, we use it to pinpoint the core challenge addressed by this paper. However, to better formulate the problem of test-time shift in FL and make the notation more unified, we propose the concept of intra-client heterogeneity existing in the inference stage of a specific client, as a counterpart of inter-client heterogeneity in the training stage among all the participating clients. We will add more clarification on the concept of test-time shift and intra-client heterogeneity in the revised version. \n\n**Difference between label-level and feature-level test-time shifts.** Thanks for the valuable comments. During the federated training stage, label-level and feature-level shift across clients are two different kinds of statistical heterogeneous issues for federated learning, which occurs in the output space and input space, respectively. Usually, different strategies are designed to solve these two issues. For example, to alleviate the feature-level shift issue, the feature encoder is kept personalized or finetuned at each client; to alleviate the label-level shift issue, the classifier is tuned instead. As for the test-time shift, feature-level and label-level shift are still two different challenges that deserve specific technical design to alleviate them respectively. Experimental results in [1] show that, even though a method that is effective for most feature-level shift cases may still not work well for the label-level shift case. Hence, it is better to discuss these two shift cases respectively. We will add more discussion on the explanation of the fundamental difference between these two kinds of test-time shifts in the revised version. \n\n[1] Test-time robust personalization for federated learning. In ICLR 2023. \n\n**Deployment in an online manner.** Thanks for the valuable question. In our current method, according to Eq. 7, the size of the test set should be larger than one to guarantee that the self-supervised contrastive learning can be conducted for test-time adaptation. However, it can still be easily extended to the online manner where the test samples arrive one by one. First, a set of training samples are selected from the training set, replacing the $A(x)$ in Eq. 7. Then, after performing the optimization for each test sample, we replace one 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 will add a more detailed discussion on this part in the revised version. \n\n**How to choose feature space.** Thanks for the valuable question. Given a specific model, we can decouple the model into two parts. One is the feature encoder that maps the raw input image into a latent embedding, while the other is the classifier that further maps the latent embedding into a probability vector among classes. Usually, for vision tasks, the classifier is composed of several fully-connected layers and a softmax layer as its final layer. We can select the features fed into the classifier, which is commonly in existing works [1,2], for contrastive learning. We will add a detailed specification on this in the revised version. \n\n[1] Supervised contrastive learning. In NeruIPS 2020.\n\n[2] Model-Contrastive Federated Learning. In CVPR 2021. "}, {"review_id": "hLtKeCKSQR", "reviewer": "Reviewer_pJGu", "summary": "This 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. ", "questions": "- 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. ", "limitations": "The limitations are clearly stated. ", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- 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. ", "weaknesses": "- 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**The reason why global invariance sharing can work.** Thanks for the valuable comments. There might be some misunderstanding towards the global invariance sharing part. There is no actual global invariance extracted among participating clients. The shared global invariance in the paper refers to the invariance encoding ability underlying model parameters. We claim that sharing the model parameter of feature extractor can achieve global invariance sharing because the parameter sharing process broadcasts the invariant encoding knowledge acquired during local training and hence mutually boost invariance extraction globally. As for the binary classification example, after aligning the model parameters of client 0 and client 1, it is of great probability that the positive samples on them are mapped to the same representation, e.g., (0, -0.5) even if there is a certain degree of deviation between their data. \n\n**Emphasizing heterogeneity types in the main text instead of supplementary material.** Thanks for the valuable comments. To make it convenient for the readers to capture key information about the heterogeneity setup, we will move the specification of several distribution shifts (inter- and intra-heterogeneity) to the main text in the final version. Due to the space limit, we still leave some detailed specifications of each feature-level shift, i.e., the number of clients, the reference previous works, and the hyperparameters for data setup, to Appendix A.2 but we provide a summary of them in Sec 5.1 to increase the readability.\n\n**More information about the learned local and global invariance.** Thanks for the valuable comments and questions. In Table 5 of Sec 5.3, we show performance improvement achieved by local invariance extraction and global invariance sharing component. Furthermore, we show how the local and global invariance are learned in the following.\n\nIn the Table below, we report the average standard deviation of representations within the same class in a client and across client, respectively. The experiments are run on Digit-5 dataset under covariate test-time shift. \n\nA lower standard deviation means that the representations within a class are much more similar. It can be seen that the representation variance of FedICON is lower than those in other methods, illustrating the ability of FedICON to extract and share invariance and hence alleviate the test-shift problem. \n\n|Method|In A Client|Across Clients|\n|-|-|-|\n|Local|0.036|0.079|\n|FedAvg|0.045|0.064|\n|FedRep|0.041|0.065|\n|FedICON(Ours)|0.024|0.058|\n\nWe hope these new results can efficiently address your concerns about whether FedICON extracts local invariance and acquires global invariance encoding ability in our proposed method. We will add the complete results in the final version. Also, it would be beneficial to use some visualization techniques to present this point in a more indicative manner. We will spend some time on visualization and present the results in the final version.\n\n**Whether the augmentation is used during training.** Thanks for the valuable comments and questions. As mentioned in line 190, given an input batch of data, we ﬁrst apply data augmentation to obtain two random augmented views of each sample, which doubles the size of the input dataset. Then, for each sample $\\mathbf{x}$, we conduct Eq. 3 to obtain its feature embedding and compute the supervised contrastive loss in Eq. 4. It should be noted that both the original data and the augmented data are involved in $P(\\mathbf{x})$ and $A(\\mathbf{x})$ and fed into the feature encoder so that they contribute to the model training process.\n\nSome previous works also construct positive pairs for contrastive learning in the supervised setting [1,2,3], which ensures the existence of at least one positive sample within the input data batch when contrasting the anchor $\\mathbf{x}$. We will add more explanations to make it clearer in the revised version.  \n\n[1] Supervised contrastive learning. In NeurIPS 2020.\n\n[2] Targeted supervised contrastive learning for long-tailed recognition. In CVPR 2022. \n\n[3] Federated learning from pre-trained models: a contrastive learning approach. In NeurIPS 2022. \n"}, {"review_id": "pH8fY4SzG4", "reviewer": "Reviewer_v7Ae", "summary": "The 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", "questions": "Another 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", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. 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", "weaknesses": "1. 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for all the valuable comments and questions. \n\n**Comparison with the combination of DG and FL.** We used three DG-originated test-time adaptation (TTA) methods as baselines because they are more aligned with our setting and can be easily implemented in an FL framework. Note that not all existing DG methods can fit into FL settings and solve the test-time shift problem appropriately. \n\n**Concept confusion issue.** We would like to explain that the concept of “domain shift” in our paper refers to the more general cases of feature-level shifts compared with the covariate and attribute shifts, which can be seen as special cases of the general domain shift. In general domain shift cases, the training and test data are from different sub-datasets of the multi-domain datasets. We will add a clarification in the final version. We take the first step to comprehensively investigate this problem in FL. \n\n**Limited novelty compared with existing DG methods using contrastive loss.** We have unique novelty compared with existing contrastive learning-based DG methods in terms of **target problem** and **technical design**. *First*, we focus on the test-time shift problem in FL rather than the DG problem where the source domain and the target domain are usually clearly pinpointed. There are some unique challenges to solve in FL, e.g., how local training is conducted, how the knowledge across clients should be aligned, and how privacy is guaranteed. *Second*, existing works [1,2] usually use contrastive learning as a regularization term to the original cross-entropy loss instead of using a pure contrastive learning-based framework to optimize the model so as to resolve the test-time shift or domain shift issues. But in our method, we establish the framework of FedICON based on the representation learning scheme. We focus more on extracting and sharing the invariance knowledge by contrastive learning in FL setting. \n\n[1] Self-supervised Contrastive Regularization for Domain Generalization. In ICCV 2021.\n\n[2] Proxy-based contrastive learning for domain generalization. In CVPR 2022.\n\n**Novelty of Sec 4.1.** Compared to existing works, Sec 4.1 has the following novelty. First, the model architectures are different. We utilize the representation learning framework to solve the problem. While most FL methods train the whole model and do not take the intermediate output into consideration. Second, the loss functions are different. Most existing FL methods utilize cross-entropy as their loss functions while our method computes the loss based on the output of the feature encoder in a contrastive learning manner. Third, the communication protocols are different. In FL, some methods average all the learnable parameters at the server, while others may average part of the model parameters to achieve a specific goal, e.g., personalization. The selection and design of the communication protocol is an important component of the FL methodology. It is non-trivial to specify the concrete global update/sharing strategy in the proposed method. \n\n**Test samples arrive one-by-one.** Please refer to **C1** in the general response. \n\n**Novelty of Sec 4.2 compared to other TTA methods.** Sec 4.2 is proposed to keep the losses of the training and test phase unified and make the whole process consistent. To the best of our knowledge, we have not seen any existing works solving the test-time shift problem in this way, not to mention formulating the problem and developing a solution for federated frameworks. \n\n**Missing baselines.** We have considered three DG-originated test-time adaptation (TTA) methods and implemented them in the vanilla federated learning framework to validate the performance of our method. However, it is hard to directly combine TTA methods with Sec 4.1.1 because most of them (e.g., Tent, EATA, T3A) require the output of the classifier head which is not involved in Sec 4.1.1. Also, it is hard to directly combine Sec 4.2 with DG methods, because Sec 4.2 is kept in a unified form with the loss in Sec 4.1.1. If the training loss is not in a pure contrastive learning manner, the operation in Sec 4.2 may lead to severe performance degradation or divergence. Note that we focus on test-time shift problem in FL and compare with ten baselines that can fit into our setting. We will add more discussion on the selection of baselines in the final version. \n\n**Data setup details.** Since we followed the benchmark setting in [1,2] as mentioned in the paper, we just illustrate how the test-time shifts are simulated and the values of hyperparameters in Appendix A.2. We will more details in the final version to make the data setup clearer. \n\n[1] Federated learning on non-iid features via local batch normalization. In ICLR 2021. \n\n[2] Test-time robust personalization for federated learning. In ICLR 2023.\n\n**Full test set requirement.** Our method does not require access to the full test dataset. Please refer to **C1** in the general response for detail. \n\n**The case where classifier is not personalized but shared.** We compared several variants for classifier training on top of the frozen feature encoder and selected the case achieving the best performance. We chose not to present the results in the main text to 1) help isolate the effects of the proposed key components and the classifier part; 2) emphasize that our method is open to different kinds of additional classifier training strategies according to the specific data setting. We add the experimental results to Table 1 in the uploaded file and will provide a sufficient illustration in the final version. We hope the new results can efficiently address your concerns about the classifier alignment. \n\n**Detailed Setting of Figure 1(b).** We provided the detailed specification of the setting of Fig. 1(b) in Appendix A.1 titled “Detailed Setting of Figure 1(b)” as shown below. We will add a brief specification in the main text in the final version. "}]

Author rebuttal:
We thank the reviewers for their valuable comments. We are glad that the reviewers found that the problem we are solving is valuable and practical in federated learning (Reviewers 9n3S, CniS, v7Ae); our idea of leveraging inter-client heterogeneity to handle test-time shift problem is novel and insightful (Reviewer 9n3S, CniS, pJGu); our experiments are comprehensive and extensive (Reviewer 9n3S, CniS, pJGu). Both Reviewer CniS and Reviewer v7Ae think the paper is well-written and easy to understand. We respond to some shared concerns by more than one reviewer below. Detailed responses to each reviewer are provided in the following. We will incorporate all the feedback in the final version.

**C1. The case where test samples arrive one-by-one (Reviewer CniS and Review v7Ae).** Thanks for the your valuable comments and questions. In our current method, according to Eq. 7, the size of the test set should be larger than one to guarantee that the self-supervised contrastive learning can be conducted for test-time adaptation. *It does not require access to the full test dataset to conduct the test-time adaptation.* Moreover, it can still be easily extended to the online manner where the *test samples arrive one by one*. First, a set of training samples are selected from the training set, replacing the $A(\mathbf{x})$ in Eq. 7. Then, after performing the optimization for each test sample, we replace one training sample in $A(\mathbf{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 will add a more detailed discussion on this part in the final version to clarify the property of testing scenarios. 

Meta-review: In the submitted manuscript, the authors present a contrastive learning-based Federated Learning (FL) framework designed to address test-time shifts. Experimental results on multiple datasets substantiate the method's effectiveness.

Post-rebuttal reviews reveal mixed opinions. Concerns have been raised about the paper's handling of inter-client heterogeneity and its comparative analysis with methods that combine both test-time shift and domain generalization techniques. My own careful review indicates that the authors satisfactorily address the inter-client heterogeneity issue. Regarding the lack of certain comparative baselines, I find the authors' focus on the novel setting of test-time shift in FL frameworks to be justifiable, and consider the exploration of combining domain generalization with test-time shifts as a topic for future work.

Given these considerations, I recommend acceptance of the paper.
Decision: Accept (poster)