SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $z_{\mathbf{x}}^{\prime}$  is the representation of the augmented sample of  $\mathbf{x}$ . Similar to the training phase,  $A(\mathbf{x})$  here also denotes the local test dataset without  $\mathbf{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(\mathbf{x})$  in Eq. 7. Then, after performing the optimization for each test sample, we replace one randomly selected 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 present our proposed FedICON in detail in Algorithm 1 in Appendix C.
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]. Digit-5, Office-10, and DomainNet dataset are three benchmark datasets 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  $5\mathrm{e - }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.
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.
Table 1: Test accuracy on Digit-5 dataset under covariate test-time shift. All results are 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)	58.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(0.83)	38.95(2.29)	57.80(2.52)	52.40(1.32)	40.55(0.48)
FedRep	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)
FedTHE	84.52(0.21)	39.00(2.21)	58.03(2.96)	51.92(1.27)	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.48)	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.35)	72.13(2.22)	56.13(0.63)	47.98(0.50)
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 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.
Method	Test Domain
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	41.78(1.23)	24.94(1.07)	55.56(2.41)	46.52(5.72)
FedProx	42.53(1.76)	24.44(0.93)	56.60(1.20)	42.94(3.53)
FedRep	20.02(1.16)	17.68(1.34)	34.03(5.92)	24.86(1.49)
FedBN	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	38.25(2.03)	25.73(2.59)	54.86(0.60)	42.75(1.18)
T3A	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 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
Table 3: Test accuracy on CIFAR10 dataset under attribute test-time shift. All results are reported in mean(std error) format over three independent runs.  $\Delta$  represents the difference between the accuracy on original test set and that on attribute test-time shift case.
Method	Non-IID Dir(0.1)	Non-IID Dir(1)
Original	Attribute Shift	Δ	Original	Attribute Shift	Δ
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.64	80.49(0.89)	66.25(0.96)	-14.24
Tent	91.39(0.77)	85.72(1.24)	-5.67	78.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
FedICON	90.24(0.52)	86.47(1.27)	-3.77	79.43(0.78)	72.41(1.29)	-7.02
(b)
Figure 3: Test accuracy on Digit-5 under varying numbers of clients (left) and varying numbers of local epochs (right). Both are conducted in 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 1 to 9.
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 FedICON-  $us$ , achieves the runner-up performance on Digit-5 and Office-10 under covariate test-time shift.
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.
Method	Digit-5	Office-10
Local	34.00(3.12)	58.07(2.84)
FedAvg	55.04(1.00)	57.55(1.46)
FedAvg-FT	54.55(0.83)	60.45(1.40)
FedICON	62.23(0.88)	64.96(1.43)
FedICON-us	61.60(1.45)	63.79(1.17)
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 5: Comparison between the FedICON and its variants. Experiments are conducted on Digit-5 dataset under covariate test-time shift. The number of 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 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.
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.
[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, 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.
[3] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv:1806.00582, 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 a 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.12488, 2019.
[63] H Brendan McMahan, Eider Moore, Daniel Ramage, et al. Communication-efficient 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.