SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wenxuan Bao $^{1*}$ , Tianxin Wei $^{1*}$ , Haohan Wang $^{1}$ , Jingrui He $^{1}$
<sup>1</sup>University of Illinois Urbana-Champaign, bao4, twei10, haohanw, jingrui}@illinois.edu
Personalized federated learning algorithms have shown promising results in adapting models to various distribution shifts. However, most of these methods require labeled data on testing clients for personalization, which is usually unavailable in real-world scenarios. In this paper, we introduce a novel setting called test-time personalized federated learning (TTPFL), where clients locally adapt a global model in an unsupervised way without relying on any labeled data during test-time. While traditional test-time adaptation (TTA) can be used in this scenario, most of them inherently assume training data come from a single domain, while they come from multiple clients (source domains) with different distributions. Overlooking these domain interrelationships can result in suboptimal generalization. Moreover, most TTA algorithms are designed for a specific kind of distribution shift and lack the flexibility to handle multiple kinds of distribution shifts in FL. In this paper, we find that this lack of flexibility partially results from their pre-defining which modules to adapt in the model. To tackle this challenge, we propose a novel algorithm called ATP to adaptively learns the adaptation rates for each module in the model from distribution shifts among source domains. Theoretical analysis proves the strong generalization of ATP. Extensive experiments demonstrate its superiority in handling various distribution shifts including label shift, image corruptions, and domain shift, outperforming existing TTA methods across multiple datasets and model architectures. Our code is available at https://github.com/baowenxuan/ATP.
Federated learning (FL) is a distributed learning system where multiple clients collaborate to train a machine learning model under the orchestration of the central server, while keeping their data decentralized [31, 18]. However, clients in FL typically exhibit distinct data distributions. For example, in the context of animal image classification, users tend to capture pictures of various animals prevalent in their respective regions, introducing label shift [51] to the local image dataset. Meanwhile, even when capturing images of the same species, the visual appearance can be influenced by the environment and camera settings, introducing feature shift [34]. It is crucial that each client can adapt the model to align with its unique data distribution [44]. Previous personalized federated learning (PFL) works have mainly focused on improving the performance on clients participating in training [41, 37, 25, 4] or generalization to new clients [8, 7, 6], assuming the availability of labeled data. However, in many real-world scenarios, clients do not have labeled data for personalization, which limits the application of PFL algorithms. For example, when employing an animal image classifier to mobile phones, their users may capture images of various animals, but without any accompanying labels indicating the species of the animal.
In this paper, we introduce a novel setting named test-time personalized federated learning (TTPFL). During the training phase, a global model is trained using source clients. During the testing phase, each target client downloads the global model and locally personalizes the model with its unlabeled data during test-time. This setting is particularly well-suited for cross-device FL, especially when generalizing to a large number of target clients that have not participated in the training phase and lack labeled data for supervised personalization. Compared to global FL, which trains a shared global model for all clients, TTPFL enables model adaptation to individual target clients facing complex distribution shifts. Compared to standard PFL, TTPFL does not necessitate additional labeled data from target clients for adaptation.
Test-time adaptation (TTA), which adapts a pretrained model from the source domain to an unlabeled target domain, could be a solution for TTPFL. However, applying current TTA methods to FL poses two challenges. First, most TTA methods assume training data are sampled from a single domain [16, 52]. In FL, where source data are distributed across multiple clients, this simplification neglects interrelationships among source domains, impacting generalization. Furthermore, the current TTA methods are usually customized for specific distribution shifts and lack the flexibility to address diverse types of distribution shifts in FL. The inflexibility of existing TTA algorithms largely results from their predefined selection of modules to adapt (e.g., feature extractor [28, 43], final linear layer [16, 36], batch normalization layers [38, 45]). However, different modules encode varying semantic information levels, and adapting specific modules may be effective for certain shifts but not others [20]. Meanwhile, although the distribution shifts among source and target clients cannot be directly inspected, the same type of distribution shifts is likely to exist among source clients. We argue that
Which modules to adapt should depend on the type of distribution shifts among clients, which can be inferred from source clients.
Motivated by this, we propose a new Adaptive Test-time Personalization algorithm called ATP to learn the adaptation rates from distribution shifts among source clients. During training, each source client simulates unsupervised adaptation and refine the adaptation rates of each module to maximize the effect of unsupervised adaptation. The server aggregates local adaptation rates periodically to improve generalization. During testing, each target client leverages learned adaptation rates to locally adapt the global model, and cumulatively averages adapted models from previous batches to enhance the performance for online TTA. Theoretical analysis confirms ATP's robust generalization due to its utilization of multiple sources and low-dimensional adaptation rates. Extensive experiments demonstrate its superiority in addressing various distribution shifts scenarios, including label shift, image corruptions, and domain shift, consistently outperforming existing TTA methods across multiple datasets and model architectures. We summarize our contributions as follows.
- We consider TTPFL, a new learning setting in FL, addressing the challenge of generalizing to new unlabeled clients under complex distribution shifts. (Section 3)
- We introduce ATP, which adaptively learns the adaptation rate for each module, enabling it to handle different types of distribution shifts. (Section 4)
- We provide theoretical analysis confirming ATP's robust generalization. (Section 5)
- We empirically evaluate ATP over various distribution shifts scenarios, using a wide range of datasets and models. (Section 6)
2 Related works
Federated learning (FL) is a distributed learning system where multiple clients collaborate to train a machine learning model under a central server's orchestration while keeping data decentralized [18].
Personalized federated learning (PFL) extends this framework by allowing each client to personalize the model to its own local data. The most straightforward PFL method is fine-tuning the global model with a few steps of gradient descent [48, 8, 7]. Similarly, another line of works use the global model as a regularizer [22] during local training. FedTHE [17] focuses on evolving local testing set, and proposes a test-time adaptation algorithm for FL that adaptively combines global and personalized models. However, all these methods require labeled data to construct personalized models. Fed-RoD [6] uses hypernetworks to generate personalized model, relaxing the requirements for labeled data. But it still requires the label distribution of the client. FedUL [30] trains a global model with only
unlabeled clients. However, it is limited to label shift where each client shares the label-conditional feature distribution  $p(\pmb{x}|\pmb{y})$ . Our setting is mostly similar to OD-PFL [2], which also focuses on generalization to new unlabeled client. It uses an unsupervised client encoder and a hypernetwork [39] to generate personalized model. However, OD-PFL requires re-training a large hypernetwork, while our TTPFL setting focuses on adapting an existing global model.
Test-time adaptation (TTA) aims to adapt a machine learning model to a testing set with dataset shift during test-time without re-accessing training data. Most of the TTA methods focus on either feature shift or label shift. For feature shift (same  $p(\pmb{y}|\pmb{x})$ , different  $p(\pmb{x})$ ), entropy minimization is frequently used to adapt the model in the unsupervised fashion. Tent [45] minimizes the average prediction entropy by adapting the batch normalization layers [15]. MEMO [52] minimizes the marginal entropy over different augmentations of the sample input image by adjusting all model parameters. SHOT [28] exploits information maximization and pseudo-labeling to achieve target-specific feature extraction. Differently, T3A [16] adjusts the final classification layer, but it is also shown to implicitly reduce the entropy. It is important to notice that all these methods pre-define which modules to be adapted in the network. For label shift (same  $p(\pmb{x}|\pmb{y})$ , different  $p(\pmb{y})$ ), most of the previous works focus on estimating the shifted label distribution. EM [36, 1] iteratively uses model predictions to estimate the label prior distribution and uses label prior distribution to adjust model predictions. BBSE [29, 3] constructs a confusion matrix on the validation dataset, and uses the prediction distribution to estimate the ground-truth label distribution. The estimated label distribution is used for re-training a model with importance sampling. [49] generalizes these methods to the online dataset shift setting where the label distribution for testing data is evolving over time. However, all these methods heavily rely on the assumption of the same  $p(\pmb{x}|\pmb{y})$ , which can be violated in real applications.
Comparison with FedTHE [17] Recently, FedTHE also explored TTA in FL. However, FedTHE focus on the test-time distribution shift for clients that participate in FL training, while we focus on improving the performance on novel clients. Moreover, FedTHE fuses global head and personalized head to get robust prediction. It cannot be easily generalized to target clients which does not have labeled data to train the personalized head.
Our paper is also related to partial fine-tuning and hyperparameter optimization. We discuss these works in Appendix A.1 in detail.
3 Motivation
In this section, we first introduce the setting of test-time personalized federated learning, and then show that current TTA methods lack the flexibility to various types of distribution shifts in TTPFL.
3.1 Test-time personalized federated learning
Preliminary We consider a standard setting for cross-device FL [46] and domain generalization [47]. Considering an FL system with  $N$  source clients  $\{\mathcal{S}_i\}_{i=1}^N$  and  $M$  target clients  $\{\mathcal{T}_j\}_{j=1}^M$ . Each source client  $\mathcal{S}_i$  has its own labeled dataset  $\mathbb{D}^{S_i}$  with  $n_i$  samples  $\{(\boldsymbol{x}_1^{S_i}, \boldsymbol{y}_1^{S_i}), \dots, (\boldsymbol{x}_{n_i}^{S_i}, \boldsymbol{y}_{n_i}^{S_i})\}$  i.i.d. drawn from its distribution  $P^{S_i}(\boldsymbol{x}, \boldsymbol{y})$ , where  $\boldsymbol{x}$  is the input and  $\boldsymbol{y}$  is its corresponding label. Each target client  $\mathcal{T}_j$  has its own unlabeled dataset  $\mathbb{X}^{\mathcal{T}_j} = \{\boldsymbol{x}_1^{\mathcal{T}_j}, \dots, \boldsymbol{x}_{m_j}^{\mathcal{T}_j}\}$  i.i.d. drawn from its distribution  $P^{\mathcal{T}_j}(\boldsymbol{x}, \boldsymbol{y})$ , while the corresponding labels  $\{\boldsymbol{y}_1^{\mathcal{T}_j}, \dots, \boldsymbol{y}_{m_j}^{\mathcal{T}_j}\}$  cannot be accessed. The distributions for different source/target clients are different, sampled from a meta-distribution  $\mathcal{Q}$ , i.e., distribution of distributions. Global federated learning (GFL) aims to find a single global model minimizing the expected loss over client population [46]:
$$
\mathcal {L} \left(\boldsymbol {w} _ {G}\right) = \mathbb {E} _ {P \sim \mathcal {Q}} \mathcal {L} _ {P} \left(\boldsymbol {w} _ {G}\right), \text {w h e r e} \mathcal {L} _ {P} \left(\boldsymbol {w} _ {G}\right) = \mathbb {E} _ {\left(\boldsymbol {x}, \boldsymbol {y}\right) \in P} \ell \left(f \left(\boldsymbol {x}; \boldsymbol {w} _ {G}\right); \boldsymbol {y}\right) \tag {1}
$$
where  $\ell$  represents the loss function and  $f$  represents model. GFL enforces that each client uses the same global model for prediction, which does not allow for adaptation to each client's unique data distribution. In contrast, personalized federated learning (PFL) personalizes the global model  $\boldsymbol{w}_{G}$  using its labeled data, and uses the personalized model for prediction, replacing the  $\boldsymbol{w}_{G}$  in Eq. (1). However, most of the PFL algorithms [8, 7, 22] require the assumption that the target client also possesses additional labeled data, which is a stronger assumption compared to GFL.
Test-time personalized federated learning In this paper, we introduce a novel setting named test-time personalized federated learning (TTPFL), and compare it with the standard GFL and PFL
Figure 1: Comparison between the testing phase of GFL, PFL, and TTPFL. TTPFL enables model personalization without requiring labeled data.
in Figure 1. TTPFL focuses on how to adapt a trained global model to each target client's data distributions during test-time, with an adaptation rule  $\mathcal{A}$  only using unlabeled data. The objective function can be formulated as
$$
\mathcal {L} \left(\boldsymbol {w} _ {G}, \mathcal {A}\right) = \mathbb {E} _ {P \sim Q} \mathcal {L} _ {P} \left(\boldsymbol {w} _ {G}, \mathcal {A}\right), \text {w h e r e} \mathcal {L} _ {P} \left(\boldsymbol {w} _ {G}, \mathcal {A}\right) = \mathbb {E} _ {\left(\boldsymbol {x}, \boldsymbol {y}\right) \in P} \ell \left(f (\boldsymbol {x}; \mathcal {A} \left(\boldsymbol {w} _ {G}, \boldsymbol {X}\right)); \boldsymbol {y}\right) \tag {2}
$$
which can be unbiasedly estimated by the average loss over  $M$  target clients unseen during training
$$
\hat {\mathcal {L}} \left(\boldsymbol {w} _ {G}, \mathcal {A}\right) = \frac {1}{M} \sum_ {j = 1} ^ {M} \hat {\mathcal {L}} _ {P} ^ {\tau_ {j}} \left(\boldsymbol {w} _ {G}, \mathcal {A}\right), \text {w h e r e} \hat {\mathcal {L}} _ {P} ^ {\tau_ {j}} \left(\boldsymbol {w} _ {G}, \mathcal {A}\right) = \frac {1}{m _ {j}} \sum_ {r = 1} ^ {m _ {j}} \ell \left(f \left(\boldsymbol {x} _ {r} ^ {\tau_ {j}}; \mathcal {A} \left(\boldsymbol {w} _ {G}, \boldsymbol {X} _ {r} ^ {\tau_ {j}}\right)\right); \boldsymbol {y} _ {r} ^ {\tau_ {j}}\right) \tag {3}
$$
The adaptation rule  $\mathcal{A}$  adapts a global model with unlabeled samples  $X_{r}^{\mathcal{T}_{j}}$ . We consider two standard settings: test-time batch adaptation (TTBA) and online test-time adaptation (OTTA) [27]. TTBA individually adapts the global model to each batch of unlabeled samples, where  $X_{r}^{\mathcal{T}_{j}}$  is the data batch that  $\boldsymbol{x}_{r}^{\mathcal{T}_{j}}$  belongs. OTTA adapts the global model in an online manner, where  $X_{r}^{\mathcal{T}_{j}}$  contains all the data batches arriving before or together with  $\boldsymbol{x}_{r}^{\mathcal{T}_{j}}$ .
3.2 Limitation of test-time adaptation
As the precursor to TTPFL, TTA [45, 52, 29] studies how to adapt a trained model to target dataset under certain types of dataset shifts. Since TTA methods only require unlabeled target data for adaptation, they can be applied in TTPFL. We test state-of-the-art TTA methods with ResNet-18 on CIFAR-10 under two types of distribution shifts: label shift and feature shift, with results presented in Figure 2. As expected, each algorithm can boost the model's accuracy under the distribution shift it is designed for. However, most algorithms improve their performance in one scenario while simultaneously impairing it in another scenario, demonstrating a trade-off in their performance on feature shift and label shift. Moreover, when facing a more complex hybrid of distribution shifts, most TTA methods fail to introduce satisfactory performance gain (Table 1). Therefore, TTA methods are not suitable for TTPFL given the variety of distribution shifts in FL client.
The inflexibility of TTA algorithms largely results from their predefined selection of modules to adapt, e.g., batch normalization (BN) layers [38, 45], the feature extractor [28, 43], or the last linear layer [16, 36]. However, which modules to adapt is closely related to the type of distribution shift. For example, adapting the last linear layer can encode the label shift (Proposition 3.1), while it may fail when the extracted features are already corrupted due to feature shift. Similarly, adapting the BN layers can improve the performance under feature shift by distribution alignment (Proposition 3.2), while distribution alignment can harm the performance under label shift [53].
Proposition 3.1 (Adapting the last layer to handle label shift). Consider two distribution  $p, q$  with  $p(\boldsymbol{x}|\boldsymbol{y}) = q(\boldsymbol{x}|\boldsymbol{y})$  and  $p(\boldsymbol{y}) \neq q(\boldsymbol{y})$ . When a neural network is calibrated on  $p$ , i.e.,  $f(\boldsymbol{x};\boldsymbol{w}) = p(\cdot |\boldsymbol{x})$ , it is calibrated on  $q$  after adding  $\log \frac{q(\boldsymbol{y})}{p(\boldsymbol{y})}$  to the bias term of the final last layer.
Proposition 3.2 (Adapting the BN layer to handle feature shift [38]). When the feature shift only causes differences in the first and second order moments of the feature activations  $\mathbf{z} = g(\mathbf{x})$  where  $g$
Figure 2: Performance trade-off of existing TTA methods under two distribution shifts.
Figure 3: Performance trade-off of entropy minimization when adapting different modules.
is the combination of layers before the BN layer, the feature shift can be removed by adapting running mean and variance of the BN layer.
To verify the connection between distribution shift and the selection of modules for adaptation, we experiment with adapting different subsets of modules within the network to minimize the entropy loss [45]. In Figure 3, we observe a similar performance trade-off between feature shift and label shift: while adapting certain modules can boost the accuracy under one distribution shift, it is less likely to succeed under the other shift. To break the performance trade-off, it is essential to adaptively choose which modules to adapt according to the present type of distribution shift. Moreover, while [20] suggests adapting different blocks in the network, we find it more important to decide (1) which module type to adapt and (2) what is the adaptation rate (i.e., learning rate for adaptation). For example, adapting all BN running means significantly outperforming adapting any one block under feature shift. Meanwhile, employing positive or negative adaptation rates for running means yields contrasting outcomes, favoring adaptation in the presence of label shift or feature shift while impairing the other. These observations motivate us to choose which module to adapt (instead of blocks) while optimizing the adaptation rates for each module.
4 ATP: adaptive test-time personalization
In this section, we propose ATP that automatically learns the adaptation rates for each module. We introduce the training and testing phase of ATP in subsection 4.1 and 4.2, respectively.
4.1 Training phase: learn to adapt with source clients
In this part, we introduce how ATP learns adaptation rates from source clients without sharing local data. ATP uses the communication protocol of FedAvg [31] to optimize adaptation rates. In each communication round, each source client first simulates unsupervised adaptation with the current adaptation rates, and then refines the adaptation rates to maximize the effect of adaptation. After local computation, the local adaptation rates are then aggregated on the server to ensure better generalization to target clients. Algorithm 1 gives the overview of the training phase of ATP. We then explain each step in detail.
Unsupervised adaptation We consider a neural network model  $f(\cdot; \boldsymbol{w}_G)$  with global model parameter  $\boldsymbol{w}_G \in \mathbb{R}^D$ . Similar to previous works [45, 38], we consider the model processes a data batch  $X_k^{S_i} = \{x_{k,b}^{S_i}\}_{b=1}^B$  at a time where  $B$  is the batch size,  $i$  is the client index and  $k$  is the batch index. In the following, we omit the superscript  $S_i$  for clarity, e.g.  $X_k^{S_i} \to X_k$ , as unsupervised adaptation and supervised refinement operate identically across all source clients. The network has  $d$  modules, with corresponding parameters  $\boldsymbol{w}^{[1]}, \dots, \boldsymbol{w}^{[d]}$ . Typically we have  $d \ll D$ . During unsupervised adaptation, we allow each module  $\boldsymbol{w}^{[l]}$  to have a different adaptation rate  $\alpha^{[l]}$ . ATP learns to adapt both trainable parameters and running statistics for batch normalization (BN) [15] layers. To achieve more precise control of adaptation, the 'module' in ATP is slightly more fine-grained than the 'layer'. For example, each BN layer has four modules: running mean, running variance, weight, and bias.
Update trainable parameters A common strategy for updating trainable parameters is performing one step of gradient descent to minimize the cross-entropy loss. Since label information are unavailable for computing cross-entropy, we instead minimize the entropy loss  $\ell_{H}(\hat{Y}) = \frac{1}{B}\sum_{b=1}^{B}(-\sum_{c}\hat{y}_{b,c}\log \hat{y}_{b,c})$ , where  $\hat{Y}$  is the prediction probabilities over the label space of a data batch. Entropy quantifies the uncertainty of the model prediction, and is frequently used in previous TTA algorithms [45, 52, 28]. For each trainable parameter module  $\boldsymbol{w}^{[l]}$ , the corresponding unsupervised update direction for each client is the negative gradient direction, i.e.,
$$
\boldsymbol {h} _ {k} ^ {[ l ]} = - \nabla_ {\boldsymbol {w} ^ {[ l ]}} \ell_ {H} \left(f \left(\boldsymbol {X} _ {k}; \boldsymbol {w} _ {G}\right)\right) \tag {4}
$$
Update running statistics The running statistics (mean/variance) in BN layers are not updated by gradient descent. Instead, they are updated by running average.
$$
\boldsymbol {w} _ {k} ^ {[ l ]} \leftarrow (1 - m) \boldsymbol {w} _ {G} ^ {[ l ]} + m \hat {\boldsymbol {w}} _ {k} ^ {[ l ]} = \boldsymbol {w} _ {G} ^ {[ l ]} + m (\hat {\boldsymbol {w}} _ {k} ^ {[ l ]} - \boldsymbol {w} _ {G} ^ {[ l ]})
$$
where  $\boldsymbol{w}_G^{[l]}$  is the running statistics and  $\hat{\boldsymbol{w}}_k^{[l]}$  is the statistic for the current batch of inputs. In previous works, the momentum  $m$  is usually a fixed hyperparameter in [0, 1]. In ATP, we consider the momentum for each module as an adaptation rate  $(\alpha^{[l]} \in \mathbb{R})$  to be learned. We define the corresponding update direction as
$$
\boldsymbol {h} _ {k} ^ {[ l ]} = \hat {\boldsymbol {w}} _ {k} ^ {[ l ]} - \boldsymbol {w} _ {G} ^ {[ l ]} \tag {5}
$$
After computing the update direction, each module will be updated along the update direction with its corresponding adaptation rate, i.e.,  $\pmb{w}_k^{[l]} \gets \pmb{w}_G^{[l]} + \alpha^{[l]}\pmb{h}_k^{[l]}$ . Expressed in a compact form,
$$
\boldsymbol {w} _ {k} \leftarrow \boldsymbol {w} _ {G} + (\boldsymbol {A} \boldsymbol {\alpha}) \odot \boldsymbol {h} _ {k} \tag {6}
$$
where  $\odot$  is the element-wise product,  $h_k \in \mathbb{R}^D$  is the concatenation of  $\{h_k^{[l]}\}_{l=1}^d$ ,  $\alpha = [\alpha^{[1]}, \dots, \alpha^{[d]}\]^\top$  and  $\pmb{A} \in \mathbb{R}^{D \times d}$  is a 0-1 assignment matrix that maps each adaptation rate  $\alpha^{[l]}$  to the indices of  $l$ -th module's parameters in  $\pmb{w}_G$ .