SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yue Tan*
yue.tan@student.uts.edu.au
Chen Chen
chena.chen@sony.com
Weiming Zhuang
weiming.zhuang@sony.com
Xin Dong
xinding@sony.com
Lingjuan Lyu†
lingjuan.lv@sony.com
Guodong Long
guodong.long@uts.edu.au
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.
Federated learning (FL) is a promising machine learning paradigm that enables multiple clients or devices to collaboratively train models without sacrificing their data privacy [1, 2]. Practical FL frameworks usually suffer from notorious non-IID (independent and identically distributed) local data distributions, which may lead to severe performance degradation during the test phase [3, 4].
*Work done during Yue Tan's internship at Sony AI.
(b)
Figure 1: (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 heterogeneity in FL has a positive effect in handling test-time shift problems. Detailed settings can be found in Appendix A.1.
To improve the performance of the trained model on local data distribution, there are a variety of personalized federated learning (PFL) methods based on different techniques, i.e., regularization [5, 6], meta-learning [7, 8], knowledge distillation [9, 10], to help the model adapt to local datasets [11].
In existing PFL systems, it is commonly assumed that the training and test datasets within a client share the same or similar distribution. However, this assumption often does not hold in real-world applications where test samples may encounter natural variations or corruptions. Although test-time distribution shifts have received increasing attention in centralized machine learning paradigms [12, 13, 14, 15], the negative effects of these shifts in PFL systems are largely underexplored. Take federated chest X-ray classification [16, 17] among multiple hospitals as an example, while most existing studies acknowledge that different hospitals can have varying data distributions due to diverse patient populations (e.g., races, ages, etc.), they overlook the possibility of variation within a single hospital. For instance, there may be multiple X-ray machines from different vendors, each with its own unique internal parameters (e.g., exposure time, spectrum, etc.). Similar issues can be found in various other domains, such as autonomous vehicles and video surveillance. Consequently, an efficient and effective solution to handle test-time shifts is necessary to achieve satisfactory performance.
The extensively studied heterogeneous distribution in FL, also known as non-IID issue, mainly refers to the inter-client heterogeneity where clients own heterogeneous training data during their federated training phase [2, 3, 18]. Differently, the aforementioned test-time distribution shift during the inference stage of PFL refers to the heterogeneity within a client, which can be defined as intra-client heterogeneity. As illustrated in Fig. 1(a), the test data of a client can differ from its training data in various aspects. This discrepancy results in unexpected features in the model optimized on the training data, ultimately leading to suboptimal performance. In this case, a natural question that arises is: what is the difference between the test-time shift in centralized settings and that in FL? Interestingly, we find that the answer lies in the unique inherent inter-client heterogeneity in FL compared with centralized settings. While inter-client heterogeneity is often considered to be a challenge for FL, in this paper, we pinpoint that heterogeneous sources can also convey a wealth of useful information that potentially helps alleviate test-time shift problems in FL systems. A benefit of this heterogeneity is that it enables the global model to capture robust features that are less sensitive to both inter-client and intra-client shifts. We firstly discover that this information has the potential to help alleviate test-time shift problems in FL systems, as it enables the global model to capture robust features that are less sensitive to both inter-client and intra-client shifts. Through our motivated experiments (Fig. 1(b)), we observe that FL under inter-client heterogeneity can improve the ability of clients to handle test-time shift problems compared to training on homogeneous data. In this way, we turn the problem of inter-client heterogeneity as a blessing in solving test-time shift.
Motivated by the above findings, we introduce Federated Learning via InvarianCe ExtractiOn and ShariNg (FedICON), a contrastive learning-based FL approach that enables clients to learn the robust features on heterogeneous data. Our theme is to fully take advantage of the benefit brought by inter-client heterogeneity during training time, and then, at the test phase, tune the local model of each client on its private test data to conquer the intra-client heterogeneity. During the training phase, clients first carry out local representation learning under the guidance of label information
so that invariant information underlying each training sample, as well as the class it belongs to, is extracted. At the global aggregation stage of federated training, the parameters of feature encoders in all clients are aggregated to further boost the ability of clients to learn invariant information among inter-client heterogeneity. Then, during the test phase, we introduce self-supervision signals to the unlabeled test samples and smoothly tune the feature encoder based on the acquired knowledge in identifying invariant properties between diverse training data and shifted test data. Notably, we leverage contrastive learning paradigms for both federated training and local test-time adaptation, enabling the trained model to well adapt to test data distribution.
- We are the first to comprehensively investigate feature-level test-time shifts in FL. We formulate the problem from the perspective of inter-client heterogeneity and intra-client heterogeneity and propose corresponding experimental settings tailored for this problem.
- We propose a novel representation learning-based federated learning algorithm, termed FedICON, that addresses the test-time shift problem by taming inherent heterogeneity in FL for invariance extraction and sharing.
- We conduct extensive experiments on benchmark datasets to measure the performance of FedICON in handling feature-level test-time shift problems. The results demonstrate that FedICON achieves significant improvement in terms of test-time adaptation ability.
Personalized Federated Learning. Personalized federated learning (PFL) aims to enable each client to learn a personalized model that performs well on its local dataset [19, 11]. Existing PFL methods are developed on the basis of various techniques. A branch of work utilizes model decoupling to separate the local model into the globally shared part and personally trained part to realize different levels of personalization [20, 21, 22, 23, 24]. [25, 26, 27, 5, 28] augment the original objective function with additional regularization terms to induce personalization via representation alignment, knowledge transfer, bi-level optimization, etc. [8, 7] incorporate meta-learning to generate the generalized model which can be further optimized as a personalized model. Existing personalized federated learning methods can only personalize the model to local training data, lacking the ability to generalize to the test dataset [29].
Federated Learning with Feature-Level Distribution Shifts. Non-IID issue is the core challenge in both FL and PFL [30]. A common type of non-IID data in real-world scenarios is known as feature-level distribution shifts, including attribute shift, covariate shift, and domain shift [31, 32]. For instance, in FL, the training and test images can be collected in different time period, differnt locations, and even from different photographic equipment. [33] leverages mutual information theory and diversity augmentation to mitigate the negative effects brought by attribute shift and domain shift. [18] proposes to use local batch normalization to alleviate the feature-level distribution shift problems. Enlightened by the idea of domain generalization methods, [34, 35] design different knowledge transfer strategies to learn high-quality domain consensus knowledge across shifted clients. Based on feature-level distribution shift setting, [36, 37] aim to alleviate the inter-domain knowledge forgetting problem and performance degradation problem on unseen clients/domains, respectively.
Test-Time Adaptation. Test-time adaptation methods aim to improve the model performance on out-of-distribution data by adapting the learned model to test data [38, 39]. Several methods propose to adjust the model by minimizing the pre-defined test-time entropy in various forms [40, 15]. [41] presents a classifier-only adjustment module for test domain adaptation, and [12] further shows that selectively fine-tuning part of the model can address different types of distribution shifts. [13, 42] tune the model on unlabelled test data via self-supervised learning objectives. [43] further proposes to use more informative self-supervised learning to solve the test-time performance degradation in the presence of severe distribution shifts. The performance of existing centralized test-time adaptation methods can be hindered due to the inherent heterogeneity in FL during the training phase, which further affects the test adaptation process.
Federated Representation Learning. Representation learning focuses on extracting higher-level representations from raw data using deep neural networks [44]. Combining it with the federated
(a) Federated Training under Inter-Client Heterogeneity
(b) Test-Time Adaptation
Figure 2: An overview of the proposed FedICON. (a) Federated training phase under inter-client heterogeneity. Each client conducts representation-wise local invariance extraction during local training and performs parameterwise 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.
learning framework enables distributed clients to learn powerful representation and thus benefit their downstream tasks [45, 46, 47, 48]. [49] introduces a generalized federated self-supervised learning framework and provides in-depth empirical studies on the effect of key components in federated representation learning. [50] incorporates model decoupling and representation learning for better PFL. [51, 52] aim to learn a common representation model without supervision under the FL framework. Both of them utilize Siamese networks as their local representation module.
Considering  $m$  clients in an FL system,  $P_{i}$  and  $Q_{i}$  are training and test data distributions of the  $i$ -th client  $(i \in \{1, \dots, m\})$  respectively and  $\mathcal{D}_i^P$  and  $\mathcal{D}_i^Q$  are the corresponding training dataset and test dataset. A sample-label pair is denoted as  $(\mathbf{x}, y)$  where  $\mathbf{x}$  is the raw input data sample and  $y$  is its corresponding label.
3.2 Inter-Client Heterogeneity
For arbitrary client  $i$  and client  $j$  in an FL system, the inter-client heterogeneity refers to varying  $P_{i}$  and  $P_{j}$  with respect to their raw data sample  $\mathbf{x}$  and/or label  $y$ . For example, in the cases where  $P_{i}(y) \neq P_{j}(y)$  while  $P_{i}(\mathbf{x}|y) = P_{j}(\mathbf{x}|y)$ , there exists an inter-client label distribution shift. The case where  $P_{i}(\mathbf{x}) \neq P_{j}(\mathbf{x})$  while  $P_{i}(y|\mathbf{x}) = P_{j}(y|\mathbf{x})$  can lead to inter-client feature distribution shift, which is common in real-world FL systems [18, 33, 36].
3.3 Intra-Client Heterogeneity
Most existing works on FL consider the heterogeneity of training datasets  $\mathcal{D}_i^P$  across clients while overlooking the intra-client heterogeneity  $P_{i}\neq Q_{i}$ . [31] examines test distribution in robust FL benchmarks, combining label-level and feature-level shifts. To better understand their distinct impacts, we propose to study these shifts separately. Here, we focus on varying  $P_{i}$  and  $Q_{i}$  with respect to  $\mathbf{x}$  rather than  $y$ , i.e., addressing the feature-level test-time shift problem in FL. Incorporating real-world scenarios, we thoroughly examine inter-client feature distribution shifts, including covariate shift, attribute shift, and domain shift. We assume that all clients have access to their own labelled training data when participating in the federated training process and test the performance of their personally tuned model on their own test set.
Challenges. It is a challenging task to deal with both inter-client and intra-client heterogeneity in an FL system. Although there is useful information underlying heterogeneous clients, how to effectively extract it and how to further leverage it to benefit the test procedure remain unexplored. Hence, there is a need for feasible solutions to tame inherent inter-heterogeneity to handle the test-time intra-client heterogeneity in FL.
4 Federated Learning via Invariance Extraction and Sharing
Both inter-client heterogeneity and intra-client heterogeneity arise from shifts in data distribution. Our core idea is to take advantage of diverse data distributions across clients so that robust features that remain invariant can be learned to handle intra-client shifts. Motivated by this, we propose Federated Learning via InvarianCe ExtractiOn and ShariNg (FedICON) as shown in Figure 2. During the federated training phase, to extract invariant information of each class within the local dataset, we design a local invariance extraction scheme by employing contrastive learning in a supervised manner (Sec. 4.1.1). Then, we perform global invariance sharing under inter-client heterogeneity to further boost the invariance extraction ability for each client (Sec. 4.1.2). During the test phase, each client tunes the personalized model to adapt to its local test set (Sec. 4.2).
**Optimization Objective.** Since each client has varying training and test data distribution, the global objective of the personalized FL system across  $m$  clients can be formulated as a two-stage optimization process,
$$
\left(\theta_ {1} ^ {\prime}, \theta_ {2} ^ {\prime}, \dots , \theta_ {m} ^ {\prime}\right) = \min  _ {\left(\theta_ {1}, \theta_ {2}, \dots , \theta_ {m}\right)} \frac {1}{m} \sum_ {i = 1} ^ {m} \frac {\left| \mathcal {D} _ {i} ^ {P} \right|}{N ^ {P}} L _ {i, \sup } \left(\theta_ {i}; \mathcal {D} _ {i} ^ {P}\right), \tag {1}
$$
$$
\theta_ {i} ^ {*} = \min  _ {\theta_ {i}} L _ {i, \mathrm {u s}} \left(\theta_ {i}; \mathcal {D} _ {i} ^ {Q}\right), \tag {2}
$$
where Eq. 1 and Eq. 2 are the optimization objective during the federated training phase and test phase, respectively. In Eq. 1,  $L_{i,\mathrm{sup}}$  and  $\theta_{i}$  are the local supervised loss function and learnable model parameters for client  $i$ , respectively.  $N^{P}$  is the total number of training samples among all clients. In Eq. 2,  $\theta_{i}$  is initialized as  $\theta_{i}'$  and further optimized on the test data of client  $i$  via unsupervised learning loss function  $L_{i,\mathrm{us}}$ . The eventual purpose of the framework is to efficiently learn invariant knowledge under inter-client heterogeneity during the training phase so that the intra-client heterogeneity can be alleviated via further optimization during the test phase.
4.1 Federated Training under Inter-Client Heterogeneity
Compared with the test-time shift problem in centralized settings, there are inherent heterogeneous source data in the federated learning framework, which is of great potential to utilize to deal with the shifted test data. Although the inter-client heterogeneous data are not directly accessible due to privacy concerns [53], their underlying invariant information can be captured by each client during local invariance extraction procedure and then globally shared across all clients to further enhance the invariance extraction ability.
4.1.1 Local Invariance Extraction.
To capture invariant information from samples within the same class, we design a local invariance extraction procedure based on supervised contrastive learning. As shown in Fig. 2(a), samples are first fed into a feature encoder to generate the representation for each sample. Then, a contrastive learning scheme is applied to the representations under the guidance of their corresponding labels so that the class-level invariant information can be extracted during the local training phase.
To make the representation ability of the feature encoder more robust, we add an augmentation module during the data pre-processing stage that incorporates some perturbations to raw samples. In this way, there are various views of the same sample fed into the feature encoder to strengthen its sample-based discrimination ability. Concretely, for the  $i$ -th client, given a batch of input samples, we generate two random augmented views for each sample. The total size of the augmented local dataset is  $2|\mathcal{D}_i|$ . For a specific sample  $\mathbf{x}$ , we use  $A(\mathbf{x})$  to denote the local dataset without  $\mathbf{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  $\mathbf{x}$ , the output of the feature encoder is
$$
z (\mathbf {x}) = h \left(\mathbf {x}; \theta_ {i}\right). \tag {3}
$$
where  $\theta_{i}$  is the learnable parameter of the feature encoder. We denote  $z(\mathbf{x})$  as  $z_{\mathbf{x}}$  for short. Samples from the same class are forced to be closer, and samples from different classes are forced to be farther.
We perform local training with the loss function as below,
$$
L _ {i, \sup } \left(\theta_ {i}; \mathcal {D} _ {i} ^ {P}\right) = \sum_ {\left(\mathbf {x}, y\right) \in \mathcal {D} _ {i} ^ {P}} \frac {- 1}{| P (\mathbf {x}) |} \sum_ {\mathbf {p} \in P (\mathbf {x})} \log \frac {\exp \left(z _ {\mathbf {x}} \cdot z _ {\mathbf {p}} / \tau\right)}{\sum_ {\mathbf {a} \in A (\mathbf {x})} \exp \left(z _ {\mathbf {x}} \cdot z _ {\mathbf {a}} / \tau\right)}, \tag {4}
$$
where  $z_{\mathbf{a}}$  is the representation of sample  $\mathbf{a} \in A(\mathbf{x})$ , and  $\tau$  is the temperature adjusting the tolerance for representation difference.  $P(\mathbf{x})$  is the set of samples that have the same labels with  $\mathbf{x}$  and can be represented as
$$
P (\mathbf {x}) = \left\{\mathbf {p} \in A (\mathbf {x}): y _ {\mathbf {p}} = y _ {\mathbf {x}} \right\}. \tag {5}
$$
By training the local feature encoder in a contrastive manner, the local invariances can be directly learned from the augmented data and be encoded to the feature encoder. Even if there is inter-client heterogeneity among different clients, their invariance encoding abilities over heterogeneous data are universal, which makes it possible for clients to mutually boost their performance on invariance extraction in a federated framework.
4.1.2 Global Invariance Sharing.
To tame the inter-client heterogeneity during the federated training process, there should be a uniform information channel that is globally shared across clients and continuously provides meaningful knowledge to boost local invariance extraction ability. Considering the fact that the feature encoder introduced above is trained to capture and encode the invariance, we can employ the parameters in the feature encoder as the information carrier and perform global invariance sharing at the server side. These invariant robust properties captured by the feature encoder are shared within different augmented views of a specific sample and different samples from the same class.
Similar to the standard federated aggregation process, for each client  $i \in \{1, 2, \dots, m\}$ , the learnable parameter  $\theta_i$  of its feature encoder is sent to the server. After receiving parameters from all the clients, the server aggregates them as
$$
\bar {\theta} = \frac {1}{m} \sum_ {i = 1} ^ {m} \frac {\left| \mathcal {D} _ {i} ^ {P} \right|}{N ^ {P}} \theta_ {i}, \tag {6}
$$
and sends  $\bar{\theta}$  back to each client as their initialized parameter in a new round. Since the feature encoder of each client is trained to extract invariance efficiently rather than classify the samples correctly, the parameter aggregation is less likely to be affected by inter-client heterogeneity. Global invariance sharing with respect to the feature encoder incorporates the invariance encoding abilities learned upon heterogeneous source data, which reversely provides sufficient homogeneous information for each client and increases the robustness of local representation learning. To employ the model trained alternately via local invariance extraction and global invariance sharing, we further train a linear classifier on top of the fixed feature encoder [54, 55]. The linear classifier can be trained in a pure local manner or trained in a federated learning manner. In this paper, we conduct multiple federated training steps at first and then personalize the classifier for each client.
4.2 Test-Time Adaptation to Handle Intra-Client Heterogeneity
In Section 4.1, we have developed components that leverage the inherent inter-client heterogeneity in FL in order to handle the test-time intra-client heterogeneity. To keep the losses of the training and test phase unified and make the whole process consistent, we adopt a self-supervised contrastive learning scheme to carry out test-time adaptation, which is also a natural and efficient way to generalize the trained model to the unlabeled and shifted test set.
Concretely, for client  $i$ , the trained feature encoder is further tuned by optimizing the following unsupervised loss function on its local test set  $\mathcal{D}_i^Q$ .
$$
L _ {i, \mathrm {u s}} \left(\theta_ {i}; \mathcal {D} _ {i} ^ {Q}\right) = \sum_ {\mathbf {x} \in \mathcal {D} _ {i} ^ {Q}} - \log \frac {\exp \left(z _ {\mathbf {x}} \cdot z _ {\mathbf {x}} ^ {\prime} / \tau\right)}{\sum_ {a \in A (\mathbf {x})} \exp \left(z _ {\mathbf {x}} \cdot z _ {\mathbf {a}} / \tau\right)}. \tag {7}
$$