Adaptive Test-Time Personalization for Federated Learning
Authors: Wenxuan Bao
Tianxin Wei
Haohan Wang
Jingrui He
Venue: NeurIPS 2023 poster (2023)
Keywords: federated learning
personalized federated learning
test-time adaptation
TL;DR: We propose a federated learning algorithm that personalize a global model to new client without labeled data under complexdistribution shifts.

Abstract: 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.

Adaptive Test-Time Personalization for Federated Learning
Wenxuan Bao' Tianxin Wei!*, Haohan Wang', Jingrui He
University of Illinois Urbana-Champaign (wbao4, twei10, hachanw, wingrnn)eiiiiiiiiiis...o.mo
Abstract
Personalized federated learning algorithms have shown promising results in adapt- ing models to various distribution shifts. However, most these methods require labeled data on testing clients for personalization, which is usually unavailable in real-world scenarios. this paper, we introduce a novel setting called test-time per- sonalized federated learning (TTPFL), where clients locally adapt 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
designed for specific kind 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 do- main shift, outperforming existing TTA methods across multiple datasets and model architectures. Our code is available at pittps:://// //gi ghg......... com/ harm//awwenxxan/ /AP
ATP
Introduction
Federated learning (FL) a distributed learning system where multiple clients collaborate to train 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 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 align with its unique data distribution learning (PFL) works have mainly focused on improving the performance on clients participating in generalization to new clients 88 (). 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.
Previous personalized federated
training [1 37 25
"Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 FI., especially when generalizing to 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, TTPEL 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 In FI., where source data are distributed across multiple clients, this simplification neglects
[6 (2.. interrelationshiis 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 [6 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] Mcanwhile, 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 follows.
We consider TTPFL., a new learning setting in FI., addressing the challenge of generalizing to new unlabeled clients under complex distribution shifts. (Section B)
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 ()
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 few steps of gradient descent [8 77. Similarly, another line of works use the global model
as regularizer [22] during local training. FedTHE [17] focuses on evolving local testing set, and proposes 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 66] 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 global model with only

---

unlabeled clients. However, it is limited to label shift where each client shares the conditianl feature distribution p(zly). Our setting is mostly similar 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 adapt a machine learning model 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(y æ). different p(z) p(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 imput image by adjusting all model parameters. SHOT [28]
exploits information maximization and aseudo-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 important to notice that all these methods pre-define which modules to be adapted in the network. For label shift (same p( x y), different p(y)), most of the previous works focus on estimating the shifted label distribution. EM [36] I 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- re-training 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( y), which can be violated in real applications.
Comparison with FeE 17 Recently, Fe 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
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 also related partial fine-tuning and hyperparameter optimization. We discuss these works in Appendix A.I 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 {8. and M target clients (T, J.M Each source client S, has its own labeled dataset D=S: with n, samples {== yf ... ( v i..dd drawn from its distribution (2, where x is the input and is its corresponding label. Each target client Tj has its own unlabeled dataset XT, {m= , .mii) i.i.d. drawn from its distribution PTs (x, while the corresponding labels {y] ..... Um cannot be accessed. The distributions for different source/target clients are different, sampled from a meta-distti Q, i.e., distribution of distributions. Global federated learning (GFL) aims to find single global model minimizing the expected loss over client population (6)
((wa) = EpweCr(wec) where Cp(wa) = E v)
(1)
where l represents the loss function and ff represents model. GFL enforces that each client uses the same global model for prediction, which does not allow for adaptation to each client's 's unique data distribution. In contrast, personalized federated learning (PFL) personalizes the global model we using its labeled data, and uses the personalized model for prediction, replacing the wa in Eq. (I). However, most of the PFL algorithms [8] 7 (2 require the assumption that the target client also possesses additional labeled data, which is stronger assumption compared to GFL
Test-time personalized federated learning this paper, we introduce a novel setting named test-time personalized federated learning (TTPFL), and compare it with the standard GFL and PFL
3

---

GEORIA GOOML
Lubilted 20RA
Blooa
Rersonstiiaaion
Too Frenocal Eratiio
PersenalLedd modied
Posdictiat
Prodlctiinn
PrRdSCT
Global FL
Personal
Test-time Personalized FL (Ours)
Figure I: Comparison between the testing phase GFL, PFL, and TTPFL. TTPEL enables model personalization without requiring labeled data.
in Figure 1 TTPFL focuses on how adapt trained global model each target client's data distributions during test-time, with an adaptation rule A only using unlabeled data. The objective function can be formulated
L(wa, A) Ep~eLp(wg: A), where Lp(wG, B(a.))e Alwa, x));
(2)
which can be unbiasedly estimated by the average loss over M target clients unseen during training
L(wa, A) M 5=1
1
> EnT (Wa,A), where E,7, (wc,
: A(wc, x ). (3)
The adaptation rule A adapts the global model with unlabeled samples X We consider two standard settings: test-time batch adaptation (TTBA) and online test-time adaptation (OTTA) [27]. TTBA individually adapts the global model each batch of unlabeled samples, where X the data batch that x belongs. OTTA adapts the global model in an online manner, where X To contains all the data batches arriving before together with at?²
3.2 Limitation of test-time adaptation
As the precursor to TTPFL, TTA 45 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 applied in TTPFL. We test state-of-the of-the-art TTA methods with ResNet-18 on CIFAR-I1 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 ''s accuracy under the distribution shift it is designed for. However, most algorithms improve their performance in scenario while simultaneously impairing feature shift and label shift. Moreover, when facing most TTA methods fail to introduce satisfactory performance gain (Table i Therefore, TTA methods suitable for TTPFL given the variety of distribution shifts in FL client.
another scenario, demonstrating a trade-off in their performance on
complex hybrid of distribution shifts,
The inflexibility of TTA algorithms largely results from their predefined selection modules to adapt, e.g., batch normalization (BN) layers [38 45), the feature extractor 28 43 or the last linear layer [16] (3). However, which modules to adapt is closely related to the type of distribution shift. For example, adapting the last linear layer 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 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. with p(æly)= = q( y) and p(y) # q(y). When neural network is calibrated on i.e., f(x; w) = -(-12),
g(y) the bias it is calibrated q after adding log p(y)
the final last layer.
Proposition (Adapting the BN layer handle feature shift (38). When the feature shift only causes differences in the first and second order moments of the feature activations z g(x) where

---

0,70
Parkin
Againtnn
TEA
SHOT
StellAcunecy
Figure 2: Performance trade-off of existing TTA methods under two distribution shifts.
Figure 3: Performance trade-off of entropy min- imization 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 we observe 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 succeed under the other shift. To break the performance trade-off, is essential 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 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 BN running means significantly outperforming adapting any one block under feature shift. Meanwhile, employing positive negative adaptation rates for running means yields contrasting outcomes, favoring adaptation in the presence label shift or feature shift while impairing the other. These observations motivate to choose which module to adapt (instead of blocks) while optimizing the adaptation rates for each module.
loss [4]] In Figure 3
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 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 [3] 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 adaptation. After local computation, the local adaptation rates then aggregated the server to ensure better generalization to target clients. Algorithm gives the overview the training phase of ATP. We then explain each step in detail.
Unsupervised adaptation We consider a neural network model f(: wG) with global model parameter we RD_ Similar to previous works 45 38), we consider the model processes data batch X 1B at time where B is the batch size, i is the client index and l: is the batch index. In the following, we omit the superscript S, for clarity, e.g. X + X as unsupervised adaptation and supervised refinement operate identically across all source clients. The network has modules, with corresponding parameters wIll| wil Typically we have d < D. During unsupervised adaptation, we allow each module w to have different adaptation rate a!l), ATP learns 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.
5

---

Update trainable parameters A common strategy for updating trainable parameters is perform- ing 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 é (Y) = 1 E>= (- .b.c log j/b.c), where Y is the prediction probabilities over the label space of data B
batch. Entropy quantifies the uncertainty of the model prediction, and is frequently used in pre- vious TTA algorithms 45
12 28]. For each trainable parameter module w!ll, the corresponding unsupervised update direction for each client is the negative gradient direction, i.e.,
k = V wi & (f wa))
(4)
Update running statistics gradient descent. Instead, they are updated by running average.
The running statistics (mean/variance) in BN layers are not updated by
w + (1 m0- + m@x = w + mm - w
where w is the running statisties and :h is the statistic for the current batch of inputs. In previous works, the momentun m is usually a fixed hyperparameter in [0, 1]. In ATP, we consider the momentum for each module as an adaptation rate (al4) e IR) be learned. We define the corresponding update direction as
11 - :h w
(5)
After computing the update direction, each module will be updated along the update direction with its corresponding adaptation rate, i.e., w + w + a Expressed in a compact form,
Wk wa + (Ax) h
(6)
where C is the element-wise product, h2 E RD the concatenation of (h ad Ji a
[a(2) .... , T and A E Rxd is a 0-1 assignment matrix that maps each adaptation rate a!!! to the indices of i-th module's parameters in wt.
Supervised refinement After unsupervised adaptation, we refine the adaptation rates on each source client with label information minimize fce(f(X, ko wx), Yx), where tc is the cross-entropy loss. We use gradient descent to optimize a, i.e.,
a + 1 aaacce(((X wa). Yy)
where 7) is the learning rate of adaptation rates. Notice that the gradient of a can be computed
Fafce(f(Xx; wa), Ya) (Xx; Ya) dws
A (ha fwwfcelf(Xe w))) Y&))
Own
da
To estimate the gradient of ax, each training client only needs to adjacently compute the unsupervised and supervised gradient, and compute their module-wise inner products. Different from many meta-learning algorithms 9 (66.. ATP is computationally very efficient since it requires no second- order derivatives. In the practical implementation, since each module in the model has significantly different number of parameters, the raw gradient for each al usually has different scales. Therefore normalize the gradient with the square root of the number of parameters in the corresponding module.
Server aggregation To incorporate adaptation knowledge from multiple source clients and enhance generalization to the clients' population, ATP use standard federated aggregation [31 to periodically aggregates the local adaptation rates. In each communication rounds, after each client locally update a for a few iterations, the local adaptation rates are uploaded to the server for averaging (as shown in line 6 of Algorithm[I] I). and then sent to source clients for the next round of training. With server aggregation, ATP learn the adaptation rates that enables successful adaptation to all source clients in average.
Communication cost Notice that ATP only optimizes the adaptation rates a without changing the global model wG- Therefore, only the adaptation rates are kept transmitted between the server and each client, while the global model parameter is only broadcasted once the start of the ATP training. Such design significantly reduces the communication cost from 2T D (for standard FedAvg) to D 2Td.
Some literatures consider (1 m) as the momentum. Here we follow the definition in Py Torch.

---

Algorithm ATP Training
ServerTrainwwc, 0% 0)
I: Broadcast we to all source clients
2: for communication round = 1 to I do
Algorithm ATP Testing
3: S' (random set of C source clients)
ClientTest([j; wG. ()) # Run on target client Tj
4: for source elient S; E & parallel do
1: B + (split x into KT batches of size B)
5: a f- ClientTrain(S, a'))
2: hhistrr 0 # Cumulative moving average
6: ag Es,es of
3: for batch :: = to K do
return ag
4: Estimate update direction h with unlabeled
X according to Eq. and
Cl ClientTrainSS, at) # Run on source client S
5: if TTBA then
8: for local epoch
E do
6: wg +- wG + (Axx)
c h
18: 6- (split 1 into K batches of size B)
7: else if OTTA then
10:
11:
for batch h to KS
(X)) Y + (k-th laheled batch in B
8:
Aniistory
k-1
K Iehistory
E 1
9:
Ta
we (Aa) haiscry
12:
Estimate update direction h with unla-
beled x according to Eq. and 5
10: Make prediction:
T3 Y
f(X B
;
wx
wx f- wc + (Ax) @
14:
a /Valce(f(XX
).
15:
return a
4.2
Testing phase: exploit adaptation rates on target clients
During testing, each target client downloads both the global model and the adaptation rates. We propose two versions of ATP: ATP-batch for test-time batch adaptation (TTBA) and ATP-online for online test-time adaptation (OTTA). We summarize the testing phase in Algorithm 2
ATP-batch For TTBA, each target client makes independent predictions each batch. For each batch of target data, ATP-batch first conducts the unsupervised adaptation identical to source clients, and then makes prediction.
ATP-online For OTTA, data comes in stream of batches ,x?" Previous works 43
usually keep updating the model batch after batch. However, such accumulative adaptation can introduce severe batch dependency problem, i.e., each batch evaluated when the model takes different number of update steps [54]. For the first few batches, the model has not adapted to the local distribution well; while for the last few batches, the model may over-minimize the entropy but increase the cross-entropy loss. To avoid batch dependency. we propose an averaged adaptation mechanism for online adaptation, whose scale of adaptation stable during online adaptation.
For each batch X k in the data stream, we always compute the update direction h starting with the fixed global model we according Eq. (4) and (5). Subsequently, instead of using only the current update direction to adapt the model, we average all the stored update direction to update the model, i.e.
To u <-
wG WG + (Ao) o
Σ h
(8)
By using the average of previous updates, simulate updating with larger batch size to utilize historical data, while controlling the number of update steps to one. In the practical implementation, use cumulative moving average (as shown in line 8 of Algorithm 2 whose space complexity does not increase with the increment of step k.
5 Theoretical analysis
In this section, we show that ATP enjoys good generalization guarantees because of the low dimension- ality of adaptation rates. Formal definitions, assumptions and full proofs are provided in Appendix
B3
We also show in B.2 that ATP has convergence guarantce similar to FedAvg [3I 46).
Theorem 5.1 (Generalization). Let H {a ||a|]| & R} be the hypothesis space (space of adaptation rates), N be the number of source clients, and K be the number of data batches on each
7

---

source client. Assuming (1) L- L-Lipschitz model, and (2) H- update. For any fixed global model wG and any E O, we have
2-norms for each module's s
Pr(sup (e(a)) - 8(a:) ((o))22)) 2 12LHRR HEH
NK (2(VK + 1)3
exp
(9)
where ε(() is the average post-adaptation error rate on source clients, and E(C) is the expected post-adaptation error rate on clients' population.
Theorem 5.1 shows that, although ATP improves the model expressiveness by adapting the model to each client's ' distribution, ATP can still provably generalize well to the clients' population. Especially, this generalization benefit from low dimensionality of adaptation rates, since the bound get looser when d increases. Moreover, this bound shows the importance of learning adaptation rates from multiple source clients: if we merge all N source domains with batches into one domain with NK batches, then the bound will be much looser.
6 Experiments
In this section, we design experiments to answer the following research questions:
RQQ:
Can ATP handle different distribution shift and outperform prior TTA methods?
. RQQ2: Does ATP learn adaptation rates specific to distribution shift?
Setup We evaluate ATP on variety of models, datasets and distribution shifts. We first evaluate on CIFAR-10(-C) with a standard three-way split [50]: we randomly split the dataset to 300 clients: 240 source clients and 60 target clients. Each source client has 160 training samples and 40 validation samples, while each target client has 200 unlabeled testing samples. We simulate three kinds of distribution shifts: feature shift, label shift, and hybrid shift. For feature shift, we follow [1] [7], randomly apply 15 different kinds of corruptions to the source clients, and 4 new kinds of corruptions to the target clients to test the generalization of ATP. For label shift, we use the step partition [5], where each client has 8 minor classes with 5 images per class, and major classes with 80 images per class. For the hybrid shift, we apply both step partition and fcature perturbations. To test ATP domain generalization datasets: under more challenging domain shifts, then evaluate ATP evaluation protocol [1O), i.e., one Digits-5 (25) and PACS [21]. We adopt the taave-one-domain-e domain is chosen to construct target clients, and the remaining domains used construct source clients. We follow similar data preprocessing in [25), while additionally applying step partition to inject label shift. Each domain divided into 10 clients, leading to 40/10 source/target clients for Digits-5 and 30/10 source/target clients for PACS. For the experiments above, we use ResNet- 18 [II] We also test ATP with two different architectures:
common choice in FL experiments 42 14 33
a five-layer CNN on CIFAR-10(-C) and ResNet-50 on CIFAR-100(-C), Detailed experiment settings are given in Appendix C.I
6.1 RQI: Can ATP handle different distribution shift?
We compare ATP with three kinds of base- line TTA methods. For feature shift methods,
Table 1: Accuracy (mean + s.d. %) on target clients under various distribution shifts on CIFAR-10
we compare to BN-Adapt [38] and Tent (45) which adjusts the batch normalization layers, SHOT (28) which adjusts the feature extrac- tor, T3A [66 which adjusts the final classifier, and MEMO [52] which uses augmentation
Method
Feature shift
Label shift
Hybrid shift
Avg. Bank
No adupttation
69.42
1: 0.13
72.96
1 0.24
68
0.24
7.7
BN- BN-Avdapi
73.52
0.22
$6.54
1: 0.10
50:42
0.39
7,0
SHOT
71.76
44
to adjust the whole network. For label shift,, we compare to EM [36] which adjusts the
Tem
71.76 0.09
50.13 021
46.05
0.25
8.3
T3A
69.53
0.08
21.70
62.17
0.17
8.0
MEMO
72.43 0.22
77.30
1: 0.15
68.07
0.28
label priori unsupervisedly with expectation- maximization, and BBSE [2 which uses the validation data to construct a confusion matrix to estimate the label priori. Since re-training a model with different label weights for each
0.12
69.85
BBSE
63.98
0.17
79.30
0.37
61.96
0.45
6.7
Surgical
69.85
0.22
76.00
1 0.17
66.94
0.43
6.3
ATP-
0.10
0.15
75.37
client is not realistic in FL. We use the estimated label distribution to adjust the output of classifier.
8

---

Table Accuracy (mean s.d. %) on target clients under hybrid shift on Digits-5 and PACS
Method
Digits-5
PACS
MNIST
SVHN
USPS
SynthDigiis
MNIST-M
Art
Cartoon
Platoo
Sketch
No odupeation
95.47
0.22
52.28
1.45
89.62
0.44
19.75
0.69
55.62
0.80
71.57
+: 1.16
74.71
0.70
90.25
0.75
74.20
0.72
BN- Adapt
94,90 1: 0.29
$7.57
1 0.53
89.51
0.39
75.34
1 0:48
59.68
0.44
73.55
1: 0.51
71.54
0.55
92.07
1: 0.26
70.92
1 0.53
SHOT
94.69
0.31
57.91
0:23
89.55
0.69
76.43
1 0.34
60.19
0.69
69.32
1 0.67
67.77
1 0.40
86.97
1:060
59:40
0.91
91.65
62.49
0.24
T3A
91.63
0.61
89.00
1.10
88.46
0.75
75.47
1.14
51.25
1.55
72.15
0.72
75.02
02 0.78
91.51
1:062
70.14
1.21
MEMO
95.92 - 0.19
$2.85 1 109
89.84
0.44
80.12
1 0.90
53.48
1.13
71.47
1: 1.29
75.57
00 0.98
90.65
1:090
76
0.65
EM
96.64
0.31
57.21
92.29
0.32
85.89
0.46
62.08
73.96
188
78.91
92.30
0:99
1: 1.52
BBSE
91.34
Surgical
97.35
0.13
59.33
2.11
94.19
0.40
86.06
0.44
05.87
0.78
74.29
2.09
77.48
0.64
92.34
0.78
50.50
3.42
ATP butch
90.81
0.27
62.18
1.71
95.41
0.26
87.91
1 0.45
09.98
1.96
82.92
1- 0.96
79.64
0.75
95.40
8 0.41
$2.26
1.57
ATP-culine
0.23
62.64
95.56
0.23
70.78
83.51
0.44
95.52
1.69
0.25
0.25
Dataset shift
Dataset shift
0.20
feature
feature
label
label
hybrid
hybrid
0.30
Adaptation
0.05
0.05
0.00
0.00
0.05
0.05
-0.10
-0.10
Inputt cany
bleck1
block2
block3
block
last Uevear
conv weivweeight
bn. bn.rumting, mean
bn.runningg. var
beiweight
in bias
linear weightt
linear Jila
Block
Module type
Figure 4: Adaptation rates learned by ATP with different distribution shifts on CIFAR-10
We also compare to Surgical 20 which uses the validation data decide which blocks adapt. For all baselines we use the validation data to select hyperparameters.
ATP can handle different types of distribution shifts Table [ shows the results CIFAR-IC 10. Under feature and label shifts, most TTA methods suffer from performance trade-off as they improve the performance utilizes data augmentation to robustify the model prediction. However, it also introduces significant computational cost during inference. As adaptive framework simpler than ours, Surgical also introduces accuracy gain across all distribution shifts. However, its coarse-grained adaptation rule prevents further improvement on the accuracy. ATP reaches great performance comparable the strongest baseline TTA method under both feature and label shifted. Under the more complex hybrid shift, ATP achieves the highest performance gain with significant margin. Meanwhile, ATP-online can further improve the performance of ATP-batch by using information from previous batches.
distribution shift while harm the other. The only exception is MEMO, which
ATP can handle more challenging domain shifts Table shows the results two domain generalization datasets with a hybrid of domain and label shifts. Compared to baselines, ATP consistently achieves higher accuracy across all domains.
ATP compatible to multiple model architectures Finally, we evaluate ATP on more model architectures: Shallow-CNN as smaller model and ResNet-50 as larger model. As shown in Table in Appendix C.2 ATP has uniformly good performance on two new models.
6.2 RQ2: Does ATP learn adaptation rates specific to distribution shift?
Besides ATP's good performance, we are also interested in whether ATP successfully learns adaptation rates specific the type distribution shift. To explore this, we group the adaptation rates by their corresponding block and module type under three kinds of distribution shifts. As shown in Figure 4 ATP learns significantly different adaptation rates under different distribution shifts. In Figure 4 (left). ATP learns to adapt the last linear layer under label shift, while mainly adapt the former layers under feature shift. More interestingly, we notice Figure (right) that the adaptation rates batch norm running statistics are positive under feature shift, but negative under label shift. Negatíve adaptation rate is usually counter-intuitive, since it disaligns the training and testing distributions. However, benefits the model under label shifts because it explicitly adapts the label prior distribution towards

---

0.775
0.715
0.750
0.725
0.725
0.700
0.075
0.675
algerkhm No a adaptation ATP bach ATP-online
algorithm ATP batch ATPenlins
0.625
0.625
60
120
240
1.50
Batch
Figure 5: Effect of cohort size and batch size
the prediction distribution. We use toy example in Appendix rate can improve the model performance under label distribution.
C.5
to show why negative adaptation
Moreover, we examine whether the learn adaptation rates are specific to the type of distribution shift by training on one distribution shift, but testing on another. We observe in Table that, ATP performs the best when trained and tested with the same type of distribution shifts. However, the adaptation rates trained on feature/label shift fails boost the performance on the other distribution shift. The adaptation rates trained on hybrid shift can generalize to feature shift and label shift, but still worse than the adaptation rates trained with same type of distribution shifts. These results show that the learn adaptation rates are specific to the type of distribution shift.
Table 3: Train and test adaptation rates with different distribution shifts, accu- racy (mean + s.d. %)
Text
Train
Feature shift
Hybrid
No aduptation
72.98
0.24
Feature shift
73.68 0.10
65.05
1.22
60.64
1.43
Label shift
67.99 +: 0.28
79.90
0.22
69.50 1: 0.52
78.92
73,05
0.35
Further discussion
Ablation study We present two variants of ATP to study how trainable parameters and running statistics contribute to the adaptability of ATP. ATP-params only learns to adapt trainable parameters, while ATP-stats focuses solely on adapting the running statistics. As shown in Table adapting trainable parameters and running statistics both play critical roles in achieving successful adaptation. More specifically, ATP-params primarily facilitate adaptation to label shift, whereas ATP-stats essentially aid in adapting to feature shift.
Table 4: Ablation study, accuracy (mean s.d. %)
Methed
Feature
Label shift
Hybrid shift
adaptation
69.42 0.13
72,98 1: 0.24 63.68 0.24
purams
ATP stats
71.27
0.17
74.0 + 0.18
64.78 0.27
ATP-batch
73.71
0.14
79,90 + 0.22
73J05 0.35
Hyperparameter sensitivity Figure Jiowss shows the effects of cohort size and batch size with CIFAR- 10 under the hybrid shift, where cohort size refers to the number of clients sampled at each round. ATP demonstrates remarkable consistency in accuracy across different cohort sizes, indicating its robustness. For batch size, we optimize the adaptation rates with the algorithm with different batch sizes. We find that ATP consistently improves the model's 's accuracy across different batch sizes, with larger batch sizes yielding greater benefits for the model. ATP-online more robust to batch size than ATP-batch since it can utilizes information from previous batches.
and subsequently evaluate
7 Conclusion
In this paper, we propose ATP that unsupervisedly learns the adaptation rate for each module to handle various types of distribution shifts encountered in test-time personalized federated learning. As potential future direction, incorporating the training of the global model could offer advantages in terms of facilitating easier and better personalization.
10

---

Acknowledgments and Disclosure of Funding
This work is supported by National Science Foundation under Award No. I1S-1947203, IIS-21 IIS-2137468, IIS- IIS-2002540, Agriculture and Food Research Initiative (AFRI) grant no. 2020-67021- 32799/project accession no. 1024178 from the USDA National Institute of Food and Agriculture, the U.S. Department of Homeland Security under Grant Award Number, I 17STQAC0000) 100-000,, and IBM-IIlinois Discovery Accelerator Institute a new model of an academie-industo partnership designed to increase access to technology education and skill development to spur breakthroughs in emerging areas of technology. The views and conclusions are those of the authors and should not be interpreted representing the official policies of the funding agencies or the government.
References
[1] Amr Alexandari, Anshul Kundaje, and Avanti Shrikumar. Maximum likelihood with bias- corrected calibration is hard-to-beat at label shift adaptation. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 222-232. PMLR, 2020.
[2] Ohad Amosy, Gal Eyal, and Gal Chechik. On-demand unlabeled personalized federated learning, 2022.
[3] Kamyar Avizzadenesheli, Anqi Liu, Fanny Yang, and Animashree Anandkumar. Regularized learning for domain adaptation under label shifts. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA. USA. May 6-9, 2019. OpenReview.nek, 2019.
[4] Wenxuan Bao, Haohan Wang, Jun Wu, and Jingrui He. Optimizing the collaboration structure in cross-silo federated learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 1718 1736. PMLR, 2023.
[5] Hong- You Chen and Wei-Lun Chao. Fedbe: Making bayesian model ensemble applicable federated learning. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net 2021.
[6] Hong- You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In The Tenth International Conference on Learning Representations. ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[7] Canh T. Dinh, Nguyen Hoang Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems, 2020.
[8] Alireza Fallah, Aryan Mokhtari, and Asuman F. Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learnnng approach. In Advances in Neural Information Processing Systems, 2020.
[9] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh, editors, Proceedings the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 1126 1135. PMLR, 2017.
[10] Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In 9th International Con- ference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net. 2021.
[11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 770-778 IEEE Computer Society, 2016.
[12] Dan Hendrycks and Thomas G. Dietterich. Benchmarking neural network robustness common corrup- tions and perturbations. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.ned. net, 2019.
[13] Wassily Hoeffding, Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association, (8(30)):1333)) 1963.
[14] Samuel Horváth, Stefanos Laskaridis, Mário Almeida, Ilias Leontiadis, Stylianos I. Venieris, and Nicholas D. Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with or- dered dropout. In Marc' Aurelio Ranzato, Alina Beygelzimer. Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Anmual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6 14, 2021, virtual, pages 12876 12889, 2021.

---

[15] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis R. Bach and David M. Blei, editors. Proceedings the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pages 448 456. JoLR.org. 2015.
[16] Yusuke Iwasawa and Yutaka Matsuo. Test-time classifier adjustment module for model-agnostic domain generalization. In Mare' Aurelio Ranzato, Alina Beygelzimer. Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 1427-2440. 2021.
[17] Liangze Jiang and Tao Lin. Test-time robust personalization for federated learning. In The Eleventh International Conference on Leurning Representations, 2023.
[18] Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji. Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D' O'Oliviira Hubert Eichner. Salim El Rouayheh, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi. Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konečný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Pranecth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Found. Trends Mach. Learn.. 14(1-2)11-20, 2021.
[19] Mikhail Khodak, Renbo Tu, Tian Li, Liam Li, Maria- Balcan, Virginia Smith, and Amect Tal- walkar. Federated hyperparameter tuning: Challenges, baselines, and connections to weight-sharing. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 19184-1 19197, 2021.
[20] Yoonho Lee, Annic S Chen, Fahim Tajwar, Ananya Kumar, Huaxiu Yao, Percy Liang, and Chelsea Finn. Surgical fine-tuning improves adaptation to distribution shifts. In The Eleventh International Conference on Learning Representations, 2023.
[21] Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M. Hospedales. Deeper, broader and artier domain generalization. In IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 2017, pages 5543-5551. IEEE Computer Society, 2017.
[22] Tian Li. Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICMI 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 6357-6366 PMLRR, 2021.
[23] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi. Ameet Talwalkar, and Virginia Smith. Federated optimization heterogencous networks. In Inderjit S. Dhillon, Dimitris Papailiopoulos, and Vivienne Sze, editors, Proceedings of Machine Learning and Systems 2020, MLSys 2020, Austin, TX, USA, March 2-4, 2020. misys.org. 2020.
[24] Tian Li. Maziar Sanjabi, Ahmad Beirami, and Virginia Smith. Fair resource allocation in federated learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[25] Xiaoxiao Li. Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non- iid features via local batch normalization. In 9th International Conference on Learning Representations. ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[26] Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li. Meta-sgd: Learning to learn quickly for few shot learning. CoRR, abs//177.000 700.0093........ 2017.
[27] Jian Liang, Ran He, and Tieniu Tan. A comprehensive survey on test-time adaptation under distribution shifts. CORR, abx/2303. 15361, 2023.
[28] Jian Liang, Dapeng Hu, and Jiashi Feng. Do we really need to access the source data? source hypothesis transfer for unsupervised domain adaptation. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 6028-6039 6039. PMLR, 2020.
[29] Zachary C. Lipton. Yu- Xiang Wang, and Alexander Smola. Detecting and correcting for label shift with black box predictors. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan. Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages
PMLR, 2018
12

---

[30] Nan Lu, Zhao Wang, Xiaoxiao Li, Gang Niu, Qi Dou, and Masashi Sugiyama. Federated learning from only unlabeled data with Blas-conditional--aainng clients. In The Tenth International Conference on Learning Reprexentaions. ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[31] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiicee learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, pages 1273-128%, PMLR, 2017.
[32] Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of Machine Learning. Adaptive computation and machine learning. MIT Press, 2012.
[33] Jachoon Oh, Sangmook Kim, and Se- Young Yun. Fedbabu: Toward enhanced representation for federated image classification. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[34] Amirhossein Reisizadeh, Farzan Farnia, Ramtin Pedarsani, and Ali Jadbabaic. Robust federated learning: The case of affine distribution shifts. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020. December 6-12, 2020. virtual, 2020.
[35] Youngmin Ro and Jin Young Choi. Autolr: Layer-wise pruning and auto-tuning of learning rates in fine ine-tuning of deep networks. In Thirty-Fifth AAAI Conference Artificial Intelligence, AAAI 2021, Thirty- Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9. 2021, pages 2486-2494, AAAT Press, 2021.
[36] Marco Saerens, Patrice Latinne, and Christine Decaestecker. Adjusting the outputs of classifier to new a priori probabilities: A simple procedure. Neural Comput., 14 14(1):21-41, 2002.
[37] Felix Sattler, Klaus-Robert Müller, and Wojeicch Samek. Clustered federated learning: Madel-agnostic distributed multitask optimization under privacy constraints. IEEE Trans. Neural Networks Learn. Syst., 32(8):3710-3722, 2021.
[38] Steffen Schneider, Evgenia Rusak, Luisa Eek, Oliver Bringmann, Wieland Brendel, and Matthias Bethge. Improving robustness against common corruptions by covariate shift adaptation. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Elsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtuul, 2020.
[39] Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernetwrrss. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 1021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 95022 PMLR, 2021.
[40] Zhiqiang Shen, Zechun Liu, Jie Qin, Marios Savvides, and Kwang- Ting Cheng, Partial is better than all: Revisiting fine-tuning strategy for few-shot learning. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Thine Conference Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances Artificial Intelligence, EAAI 2021, Viriual Event, February 2-9, 2021, pages 9594-9602, AAAI Press, 2021.
[41] Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio. Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan. and Roman Gameett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 4424-4434, 2017.
[42] Benyuan Sun, Hongxing Huo, Yi Yang, and Bo Bai. Partialfed: Cross-domain personalized federated learning via partial initialization. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 23309-23320, 2021.
[43] Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei A. Efros, and Moritz Hardt. Test-time training with self self-supervision for generalization under distribution shifts. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 9229 -248. PMLR, 2020.
[44] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. CoRR, abs/2103.007710. 2021.
[45] Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno A. Olshausen, and Trevor Darrell. Tent: Fully time-time adaptation by entropy minimization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
13

---

[46] Jianyu Wang, Zachary Charles, Zheng Xu. Cauri Joshi, H. Brendan McMahan, Blaise Agüera Arcas. Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, Suhas N. Diggavi, Hubert Eichner, Advait Gadhikar, Zachary Garrett, Antonious M. Girgis, Filip Hanzely, Andrew Hard, Chaoyang Hc, Samuel Horváth, Zhouyuan Huo, Alex Ingerman, Martin Jaggi, Tara Javidi, Peter Kairouz, Satyen Kalc, Sai Praneeth Karimireddy, Jakub Konečný, Sanmi Koyejo, Tian Li, Luyang Liu, Mchryar Mohri, Hang Qi Sashank J. Reddi, Peter Richtárik, Karan Singhal, Virginia Smith, Mahdi Soltanolkotabi, Weikang Song. Ananda Theertha Suresh, Sehastian U. Stich, Ameet Talwalkar, Hongyi Wang. Blake E. Woodworth, Shanshan Wu, Felix X. Yu, Honglin Yuan, Manzil Zaheer, Mi Zhang, Tong Zhang. Chunxiang Zheng, Chen Zhu, and Wennan Zhu. A field guide to federated optimization. CoRR, ahs/2107.06911 2021.
[47] Jindong Wang. Cuiling Lan, Chang Liu, Yidong Ouyang. Tao Qin. Wang Lu. Yigiang Chen, Wenjun Zeng. and Philip S. Yu. Generalizing to unseen domains: A survey on domain generalization. IEEE Trans. Knowl. 2023.
Data Eng., 35(8):8052-80)
[48] Kangkang Wang, Rajiv Mathews, Chloé Kiddon, Hubert Eichner, Françoise Beaufays, and Daniel Ramage. Federated evaluation of on- on-device personalization. CoRR, abs/1910.100255 2019.
[49] Ruihan Wu, Chuan Guo, Yi Su, and Kilian Q. Weinberger. Online adaptation to label distribution shift. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 11340-11351, 2021.
[50] Honglin Yuan, Warren Richard Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[51] Jie Zhang. Zhiqi Li. Bo Li, Jianghe Xu, Shuang Wu, Shouhong Ding, and Chao Wu. Federated learning with label distribution skew via logits calibration. Kamalika Chaudhuri, Stefanie Jegelka, Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 26311-26633 162622.. PMLR, 2022.
[52] Marvin Zhang. Sergey Levine, and Chelsea Finn. MEMO: test time robustness via adaptation and augmentation. In NeurIPS, 2022.
[53] Han Zhao and Geoffrey J. Gordon. Inherent tradeoff's in learning fair representations. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc Emily Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 15649- 15659, 2019.
[54] Hao Zhao, Yuejiang Liu, Alexandre Alahi, and Tao Lin. On pitfalls of test-time adaptation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 42058 42080. PMLR, 2023.
[55] Yi Zhou, Parikshit Ram, Theodoros Salonidis, Nathalic Baracaldo, Horst Samulowitz, and Heiko Ludwig. Single-shot general hyper-parameter optimization for federated learning. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
14

---

A More discussions
A.1 More related works
Partial fine-tuning, i... updating subset of modules of a pretrained network on a new dataset, has In FL. PartialFed 42] adaptively decides whether each parameter is shared or personalized. However, it cannot generalize to testing clients that do not participate in the training. Recently, surgical fine-tuning (20) selectively fine-tunes subset of blocks with similar intuition that the type of distribution shift influences which part of the network to be adapted. Different from their method, we focus on the unsupervised setting and propose refine the adaptation rate for each module.
been studied in supervised settings [26 35 40
Hyperparameter optimization is also related to algorithm if considering adaptation rates as set of hyperparameters. [19] first investigates the problem of federated hyperparameter tuning and proposed FedEX that leverages
from neural architecture search to efficiently tune hyperparameters. (55) introduces FloRA that addresses use cases of tabular data and enables single-shot federated hyperparameter tuning. While these methods focus on improving the efficiency of hyperparameter optimization, our paper focuses on finding the optimal adaptation rates that benefit test-time personalization.
A.2 Broader impacts and limitations
Broader impacts We are not aware of any potential negative societal impacts regarding our work the best of our knowledge. For all the used data sets, there is no private personally identifiable information or offensive content.
Linitations One possible limitation is that consider fix global model for lower communication cost and better generalization, while might be beneficial to also train a global model for easier personalization, which could be promising future direction.
15

---

B Theoretical analysis
In this section, we give theoretical proofs of convergence, and generalization of ATP.
B.1 Approximation analysis
In this subsection, we give detailed proofs of Proposition 3.1 3.1 and 3.2 in Section 3 of the main text.
These propositions show why certain types of distribution shifts can be handled by adapting certain layers in a neural network.
B.1.1 Proof of Proposition 3.1
Proposition 3.1 (Adapting the last layer to handle label shift). Consider two distribution p- q with p(æly) = q y) and p(y) # q(y). When a neural network is calibrated on i.e., f(æ; (x) = =((+))) it is calibrated on q after adding log giv) to the bias term of the final last layer. p(v)
Proof. W.L.og... assuming the last layer of the neural network is a linear layer. Denoting g w as the input of the last layer, where is the input and wy is the model parameters for the feature extractor (i.e., all layers except for the last classification layer). Denote w1.- w k as the weights of the last layer and ".. .b the bias terms of the last layer, assuming K classes. Then we have
exp(w (x: wg) + be)
f(x; w)
Ee= exp(w]gll w. be)
Since the neural network is calibrated on p. for all class index = 1,.,. K, we have
f(x, w) = p(y = ex(2))
where eg an one-hot vector with its e-th element as one. For distribution g with the same conditional distribution and different priori, by Bayes' theorem, Væ,
g( ( Ey ((
p(aly))(())
p az))
2(v) p(y)
Σy p (x)
givl
p(P)
Therefore, we can calibrate the neural network on distribution q simply by adding log p(v) terms, i.e.,
2(v)
to the bias
exp(w, g(:; w.) + be log g(ca) p(e.)
Seal (a;; War
=1= exp( w [ wg) ) + b + log g(e.) ple.))
exp(w, gl x: wg + be) ple.)
g(ec)
D'= exp(px [ mgg be) gle.)
ple,)
p(e.(a)) g(e.)
ple.)
Σk Ed=1 p (x) g(eo)
p(e))
g(y = eo x))
16

---

B.1.2 Proof of Proposition 3.2
Proposition 3.2 (Adapting the BN layer to handle feature shift [8]). When the feature shift only causes differences the first and second order moments of the feature activations 2 g(z) where the combination of layers before the BN layer, assuming independent activations, the feature shift can be removed by adapting running mean and variance the BN layer.
Proof. Denote the source and target feature (marginal) distributions to be p(z) and q(x). Given independent, activations, we only need to test the marginal distribution of each z E 2 = g(x). For each 2, since the feature shift only introduces differences in the first and second order moments, there exists A and r > O, s.t., V2, E
21
A r
Pr (r 21) 2m-
Pr wmpp
which indicates that the distribution of is first shifted by A and then scaled by Such distribution shift in the feature activation can be removed by adapting the running mean /4p and variance of
My r Mp A
0g 0p"
As result, for all & E R
Pg 2 1)
Pr
Pr (= 2 14 . and
0g
(2 My + og A mmp
Pr (2 2 Up + Op t)
amp
Up
Pr
a~p
Op
which indicates that the feature shift is removed after normalization with running statistics Pa: Tg- 1.1
17

---

B.2
Convergence analysis
We first show In this part, we show that ATP has the same convergence guarantee as FedAvg 31 in Lemma B.5 and [B.00 15.00that that ATP preserves convexity and smoothness, which are two important conditions the analysis of convergence. Then we formally prove the convergence of ATP in Theorem B.11
B.2.1 Definitions: local and global objective
For clarity, we first formally define the data generation process, and local/global objectives for optimization.
Meta-distribution
Q
(distribution of distributions)
distribution
p%
p'x
(each one is client)
Data batches
X- v
5 y.
x y
Figure Data generation process
Data generation
We consider two-stage sampling process illustrated in Figure
There are N source clients' distributions ps., ps. ........ pSx and M target clients' distri- bution pTi, pTs..... pTar i.i.d. drawn from meta-distribution Q.
For each source client i's distribution pS., there are K data batches (X Y (X ,Y, .., Sx K Y drawn i.i.d. from pS.
Each batch consists of B samples, (X Y { k,b vi where B is the batch size.
For simplicity, we assume that all source client has the same number of batches K and batch size B.
Definition B.1 (Batch objective). Define the batch objective of the k-tth batch client S be
1 E (CElffEE wh v
where S. we (Ax) o S and h is the update direction computed with x (x 1B with Eq. and 5
km
Definition B.2 (Local objective). Define the local objective of client to be
F.(a:) 1
Fu((a)
K
Definition B.3 (Global objective). Define the global objective to be
F(a)
F((a)
18

---

B.2.2 ATP preserves convexity and smoothness
In this part, we show that ATP preserves convexity and smoothness, which are two important conditions in the analysis of convergence.
Definition B.4 (Convexity). A function f RP + R is convex if for all 0 E IR and À E [0, 1]
ʃ()æ) + (1 - A)m²) S \f(r)) + (1 - ()2(2))
Lemma B.5 (Convexity preserving). (] w) y) is convex w.rl. w given any data sample (x, y). then (.(a) is convex w.r.t
Proof. Noticing that / wG + (Aa) @ h is linear to a, linear transformation preserves convexity. For any update direction h and data sample (3r V..)) find that
lce(f(xi we + (A(Aa1 + (1-A)a2) @ h y
-cF(s(< We + (Axi) @ h + (1 >) [wo + (Ao2) @ n -vh.)) McCE(f( we (Aai) : (1 A)ece(f( we (Aoz) o h S vk.)
ie., (f we + (Aa) h v is convex w.r.t.
Finally, since
Fx(a)= KB k=1b=1
ΣΣ
loels(a (a we + (Aa) o h y,y
which is the average of K B convex functions, have that Pi(ax) also convex a. Definition B.6 (3-smoothness), A function f R + R is L-smoothness with / > if for all 21, 22 E R.D.
[V/(())) (=22)) B||a0] -a2n--allll
Definition B.7 (H bounded for data batch S if
update direction). The update direction is H 7-moodule-wise
S. 12 H, 11 .
where ( )0 the update direction corresonding the I-th module and d the number of modules in the neural network.
Lemma B.8 (Lipschitz parameter). the update direction is H soodul--wise-boonneee for a data batch S Given two adaptation rates Or1, 022 and the global model wG. have
S (ax) ) - (a2))l2 S H | 0:212
where w (an) = WG + (Aai) @ S is the personalized model updated with on as the adaptation rate.
Proof.
(l11)(a)) - w (a2)ll2 = [[(wa + (Aci) C h S.) - (wa + (Axz) @ he )ll
[I(A(a) o2)) 3 1t
1| (A(a) -2z))
110 1/2
2 ( a!" !!!
1=1
d
Σ H2 (a - a!!¹
<
1=1
H
||o
a2212
19

---

Remark B.9. Lemma B.S indicates that when the adaptation rate is perturbed by little, the personal- ized model parameter " is also only perturbed by a little.
Lemma B.10 (Smoothness preserving). (f(1) ( w)) y) is B-smooth w.r.t. w given any data sample (æ, y). and (2) the update direction hog is --mouule-wiss-boundeee for all data batches X then F.((x) is ( -moothness w.r.l. Ox.
Proof. We first give an upper bound of || diag(h )ll when h is H-module- wise-bounded. The update direction h E R2 is the concatenation of update directions for each module {(h)) ) J/=1"
i.e.
h1 ( = = .... (1
where (h |1 is column vector representing the update direction of the 1-th module in the model. Similarly, any other vector E RD can be correspondingly expressed as
e² = (1l1)) ,.. (olu) '
Then,
|| diag( ( Jull2
|A T diag(h;")
sup v
||0
||
(h O 1
sup vERD
|
Σd
(h (1)))))
bol
sup vR
Ei=1
|01111
Σ
S:
)0
k
sup vRDD
Σ
14]
Σ t=l
H
|
sup VERD
(Definition B.7
Σ
1
2
H
We then prove that for any H H-module- bise-buneed update direction h and data sample (x vh.)) we have (cE(f(Z)) u (a)), v is H2 Bsmoothhsss w.r.t. a.
Va ( w (an)
V lon k, w (az2), D
||AT (
w (a)) (CEls i (a1)) S A T V " (as) tas lan(s
R we (az2), d
< A diag( h
of (an) lce (f ( w (an)) v V ( lce ( w (az)), yk.b)
(an)
A T diag( h
1| B S (an) x (02)1l2
(Definition B.6
A diag(h
$ ||2 B H. ||on 02/12
(Lemma B.8
H². 3. ||01 0:2 2
20

---

B.2.3 Convergence of ATP under FedAvg framework
Finally, we show that with preservation of convexity and smoothness, ATP shares the same conver- gence guarantee as FedAvg [1. We apply the proof in [46]
Theorem B.11 (Convergence of ATP). Assume that
1. At any round t, each client takes T SGD steps with learning rate 1μ.
2. Full participation, i.e, each source client participates every round
3. (x; w). y) is convex and B-smooth w.r.t. w given any data sample (æ, y).
4. The update direction h is fmodule-wise-boumdee for all i.
5. Bounded inner variance: for any and client i,
-VaF
6. Bounded outer variance: for any a and client i,
[VFi((a) var(a)
the client learning rate satisfies 1) 2 aH13. then one has
-(a -(a")) S 11a8 x
n1²
+ #2 18
+ N
TT
where a' arg ming F()) and azk communication rounds and k local epochs.
ΣN ii=l a. a;
1kk is the local adaptation rates after
Proof. The optimization process of ATP is similar as FedAvg [31, where the difference is that ATP adapts the adaptation rates instead of model parameter. Lemma B.3 and B.10 that ATP preserves convexity and smoothness, i.e., each client ((())) is convex and (H2 2/2 smmothhnsss.. Therefore, we can apply Theorem in 46 to complete the proof.
Remark B. The convergence rate of ATP is O 77
21

---

B.3 Generalization analysis
In this part, we studied how an adaptation rate a learned by ATP that performs well on source clients can generalize to target clients. More specifically, we are interested in how many different source clients are required to ensure a certain generalization error. Similar most the other generalization analysis, we (1) derive generalization bound for any fixed hypothesis (ox,, and (2) quantify the size of hypothesis space.
B.3.1 Definitions: data generation and error rates
We first formally define the error rates.
Definition B.13 (Error rate for one data sample). Let f(; (n) : - A be the neural network with model parameters w that takes one data sample E input and outputs probability
distribution over the label space, i.e., ((aa,,bi 2 0 and 1 1 (ax ) 1. Given adapted model parameters " define the error rate on one data sample (ar Dk,b) to be
euus(w)
= := (w) f (x w
Remark B.14. Definition B.13 equivalent to the expected misclassification rate if when making random decision based the output probability f( (
Definition B.15 (Error rate for data batch). Given global model parameter wG, adaptation rate a, and batch of data X 1 Y {yxb) define the error rate for one data batch
B
(sx( = - = B
Cakt ( St)
where
w wa + (Ax) @ h
and h is the update direction computed with X
Definition B.16 (Error rates for one client). Given global model parameter wa, adaptation rate a. and a source client S; with data batches {(x Y > define the empirical error rate for source client S;
K
(a() = K
(ar)
k=1
Also, define the expected error rate for source client S
ex(x)) == E (x yf [c)(a)]
Remark B.17. (a:) quantifies the error rate on client Sj's 's finite dataset D {(X Y ) E((a)) quantifies the expected error rate on a new data batch from client S, Notice that the same definition applies to target clients.
Definition B.18 (Source error rate and expected target error rate), Given global model parameter wG. adaptation rate or, and N source client SI.... . ,S. each with K data batches ((x)) S: define the training rate
& Y k
É(a:) = K 1=1
(()
Also, define the expected testing error rate
=(ax) Epp:ng (e:((r) ps:] = Epsnng Fxx y [E):(a))
Remark 19 ε(ax quantifies the averaged error rate across source clients' finite samples. s(a:) quantifies the expected error rate on new data batch from new client (target client). Noting that both emorr rates are defined with respect to the personalized model after adaptation.
22

---

B.3.2 Generalization bound for one hypothesis
Next, we derive generalization bounds for one fixed adaptation rate a. Since we consider fixed a. for clarity, we denote
7 = Esk (ax)
K
Z. =
Zik = Ei(a)
K k=
A4 = ( @
[Zix] = (a)
N
Z. =
Zz = E(a)
N ==1
1
14
N
== Eps: e((x)
Intuitively, with enough number of source clients and number of batches, have Z., ft. = μ. Lemma B.20 (Hoeffding's ss inequality). Let XI,... X be independent random variables such that a, S X, S b, almost surely. Consider the sum of these random variables Sn = X1 + ... + Xn. For all > 0,
2
Pr(S, - [[S,] 2) S exp
Li=l (be a11)
Proof. Please refer to
Lemma B.21 (Concentration of averaged client expected error rates). For any c 0, we have Pr(ū. μ e) < exp(-2Ne2
Proof. Notice that "i.-. , are independent given Q. For all 1,... N,, Epg and 0 < p4 < 1. Therefore,
Pr(ū. μ = Pr
Σ
NN 2 Ne
N
Pr
Σ Mi E
Σ H4
2 Ne
2 (Ne)²
eexp
(Hoeffding's 's inequality)
N. (1-1-- (1-0)² 0)2
= exp( 2N
Lemma B.22 (Concentration of client empirical error rate). For any e >> 0,
Pr Z... - f.. 2)) I Sop(-2NKet)
Proof. Given distributions pSi .... pSv, we have Z11 LIk. Z21. ZnK are independent. For any i = 1,... N and = =,=1--. ... K, we have EE Therefore.
xx S ps Zip = ai and 0 S Zia S 1.
N Σ Zik
Pr (Z.. n. c psi
pSN) Pr
Kp > NKc ps. ps....., PSN
23

---

N
K
N K
Pr
Zin
E
Zik
2 NKke
psi psi.... psv
k=1
k=1
2 (N)))
exp
(Hoeffding's 's inequality)
NK (1 0)
exp( 2
Then, the tower property,
Pr(Z.. j.. 2 c) = p..,... PSN w Pr(Z.. p.. 2 ,
< sup
Pr(Z.. jr. 2 pôi ,..,,,,pp.
psi,.., pgy
S expl (xx((22Nee))
Proposition B.23 (Generalization for one hypothesis). For any fixed global model we and adaptation rate a, for any e > 0, we have
2N
2 s(a))
4exp
(VK +13
Proof. For any > 0, we have
Pr(Z.. p 2 c) = Pr((Z. - p.) (T. -1) -4) 2)
< inf Pr(((.. - 2 d') v (i. - μ E E))]
< inf Pr(Z.. - n.. 2 c') + Pr(j. -2
< inflexp(-2N - exp( 22N(e
(Lemma B.21 and B.22
To make the bound clear (although not optimal), we choose e' Then the bound becomes,
VK+1 and thus - é' VK
VK+1 e.
2NK?
Pr(Z.. - μ 2) < 2exp
(VK + 11²
Similarly we can show that
2N (VK+ + +12
Pr(Z.. - 11 -e) 2exp
Therefore,
(-
2N
(((((a))((a))) 2 = Pr([Z. 1l S 4eep
[[KK 112
Remark B.24. The RHS is function of both (1) N, the number of source clients and (2) K, the number of data batches on each client.
When N + oo, given any fixed K the RHS + 0, indicating that é(a) & e(a:)
(convergence in probability).
However, given fixed finite N, when K + co, the RHS does not limit to zero. Intuitively, sampling more batches on finite source clients only help the algorithm learn finite distribution pS..... pSx_ However, more data batches on existing clients does not help further exploration of the meta-distribution and generalization novel target clients. Actually, given a fixed finite N, when + co, ē(a) E N EN ((a)) # c(a.).
If we put data from N sources (each with batches) together as one source with NK batches. The generalization bound is looser.

---

B.3.3 Generalization bound for hypothesis space (proof of Theorem 5.1)
Finally, we derive the generalization bound for the hypothesis space. We first show in Lemma B.27 that Esy(ar (a) is LLisssipsciizzz a, then we apply standard generalization analysis in Theorem 5.1 based on covering number [32].
Definition B.25 (L-Lipschitz). The neural network f(;; w) is L-Lipschitz w.r.t. w, Vo and 11; W2-
(x: wi) f(;; w2)ll2 < L. | W2ll2
Definition B.26 (H wi--module-wie-bbooeeeee wise-bounded for data batch S if k
update direction). The update direction is H H-module-
[ S 1 12 - H, 4I 1,...,
where ( 1t1 is the update direction corresonding to the (-th module and d is the number of modules the neural network.
Lemma B.27 (Lipschitz error rate). Given a data batch (x Y if the update direction is I (1-module-wise--ouundd given any two adaptation rates 01, 022 and the global model WG, we have
(ei((r)) (0()) S LH -||011 02212
Proof.
[((((r)) - Eig (a2)
B
w (f (ci)))
B
1 -
S.
Σ
1 -
f
w (
B
B
b=1
b=1
(f ux (ai)) -fn--f(asz f( w (a2))
B
B
1
1 ll |[ w (ai)) -
( fax ux (a12)
B
B
1
1 L. (ai)
(a1) - w (az)
(L-Lipschitz model)
B
b=1
1
<
Σ
|yr.bl1 ll L. 11 ||a0 02212
(Lemma B.8
B
I.H.
||
02l12
Remark B.28. Intuitively, Lemma B.27 shows that small change in will result in bounded change on (c)( Corollary B.29. ε(ai and c(ai are Lips-iipschitz w.n.t. a.
Proof. É(()) and e(aa are expectations of Éij given the empirical and expected distribution of SX., Y&.). Lipschitz property is preserved.
Theorem 5.1 (Generalization for hypothesis space). Let H. = {a : |||/| < R} be the hypothesis space (space of adaptation rates), N be the number of source clients, and K be the number of data batches on each source client. Assuming (1) L-Lipschitz model, and (2) H w--module-wise-ooumee update direction. For any fixed global model wG and any e O, we have
[12LHR)
d
NK²2 2(VK + 1)2
Pr(sup |=(a) - =(a)) 2 2 acH
.4exp
(10)
where É((x) is the average post-adaptation error rate on source clients, and E(a:) is the expected post-adaptation emror rate on clients' population.
25

---

Proof. We use covering number to derive the generalization bound [32 Define estimation error
A.(a) = =(a) - E((x)
Then,
Aa((n)) Ac(a2)) [[{(0)) {(()) )] [5( {(22)
< (((()) - e(0r2) + (8((n) {(02)
2LH -|a 22112
(Corollary [B29
= (a: ||0|| < R. E IRd} can be covered by (2, L2 balls with radius बा./7 Lemma 6.27 in [32] shows that
3R
d
12LHR d
S = N2, -
1
c
Denote these 1.2 balls be B Bs.
Pr sup (e(a) 5(cx)) 26 aeA
Pr
sup aEB
(s(()) - =(a)) 2 20)
For each ball B., s 1,... S, denote the center be ag. For any a E Bs, we have || a,l s 4TH therefore
A. A.(a)) S 2LH laa a, a,15
Intuitively, every Bg has similar error rate. Therefore, the error rate for the whole ball is upper bounded, as long as the center a, has small error rate
Pr sup |=(a) - E(a:) 2c) Pr sup AEB,
Ac(a))
aEB.
< Pr
sup aEB,
| ) (1..(a.) - A.(a)
A (o.)) + E =
Pr
sup aEB,
Pr (le(ar)) {(ar)) 2 2
Finally, by Proposition B.23 for each ax
Pr (1((((s)) - E(ax,) 2 2) S 4exp
NK²
(<(VK +)2
Put all together
12LHR) d
NK
Pr sup (e(a)) - E(()) 20) 5
.4exp
2(VK +
acA
26

---

C Additional experiments
C.1 Detailed experiment settings
C.1.1 CIFAR-10 experiments
Data preparation We use a benchmarking three-way split [50]: we randomly split the dataset to 300 clients, 240 of them source clients and 60 are target clients. Each source client has 160 training samples and 40 validation samples, while each target client has 200 testing samples. We simulate three kinds of distribution shifts: feature shift, label shift, and hybrid shift. For feature shift, we follow 12 17 randomly apply 15 different kinds corruptions to the source clients (Figure 7a and 4 new kinds of corriptions to the target clients (Figure (b) test the generalization of ATP. The corruption severity is randomly selected from { 2, 3, 4, 5}. For label shift, we use the step partition [5], where each client has 8 minor classes with 5 images per class, and major classes with 80 images per class. For the hybrid shift, we apply both step partition and feature corruptions.
Imguiee
Motton Blur
Zoom Btur
Snow
Frost
Fog
Speokle Noise
Goussion Blur
Spotter
Saturate
(b) 4 corruptions for testing clients
(a) corruptions for training clients
Figure 7: 15 + 4 different corruptions use construct feature shift
Global model training We first train a global model with FedAvg over the training sets of source clients.
ResNet-18: The global model is ResNet-18 18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for = 200 communication rounds with full participation (cohort size C 240), local epochs E = 1, learning rate 1) 0.01 and batch size 20.
Shallow CNN: The global model is randomly initialized 5-layer CNN. We train the global model for T 200 communication rounds with full participation (cohort size C = 240), local epochs É 1, learning rate 7 0.1 and batch size B 20.
ATP training We initialize the adaptation rates as all-zero vector, and optimize over the validation sets of source clients. We optimize the adaptation rates for T = 200 (for ResNet- 18) 400 (for Shallow CNN) communication rounds with partial participation (cohort size = 60), learning rate 0.1 and batch size B 20.
ATP testing We test the optimized adaptation rates on each target client. We use batch size B 20 by default, and test different batch size in Subsection 6.3
C.1.2 CIFAR-100 experiments
Data preparation The data preparation is similar to CIFAR-10 experiments. The only difference is for label shift, each client has 98 minor classes with image per class, and 2 major classes with 51 images per class. Same partition is applied to hybrid shift.
Global model training source clients. The global model is ResNet- 18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for 200 communication rounds with full participation (cohort size C' = 240), local epochs E learning rate 7 0.01 and batch size B = 20.
first train a global model with FedAvg over the training sets of
27

---

ATP training We initialize the adaptation rates as a all-zero vector, and optimize it over the validation sets of source clients. We optimize the adaptation rates for T 200 communication rounds with partial participation (cohort size C' = 60), learning rate v = 0.1 and batch size B = 20.
ATP testing
We test the optimized adaptation rates each target client. We use batch size B 20.
C.1.3 Digits-5 experiments
Data preparation Digits-5 dataset contains five domains: MNIST, SVHN, USPS, SynthDigits, and MNIST-M. We adopt the cave-ooe-domain-out evaluation protocol 10], i.e., one domain is chosen as the held-out testing domain, and the remaining domains are regarded as source training domains. We follow the data preprocessing in [25], while additionally applying step partition to inject label shift. Each domain is divided into 10 clients, leading to a total of 40 source clients and 10 target clients. Consequently, each client ends up with approximately 743 images spread across 10 classes. Each source client has 80% of its samples as training set and the remained 20% as testing set. Each client has major classes and minor class, where the ratio of images per class is approximately 16: 1 (the same as our CIFAR- 10 experiments). Since there is already domain shift, we do not add corruptions.
Global model training We first train a global model with FedAvg 31 over the training sets of source clients. The global model is ResNet-18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for T = 200 communication rounds with full participation (cohort size C' = 50), local epochs E 1, learning rate rj = 0.01 and batch size B 20.
ATP training We initialize the adaptation rates as a all-zero vector, and optimize over the validation sets of source clients. We optimize the adaptation rates for T 200 communication rounds with partial participation (cohort size C' = 10), learning rate v7 = 0.5 and batch size = 200.
ATP testing
We test the optimized adaptation rates on each target client. We use batch size B 200.
C.1.4 PACS experiments
Data preparation PACS dataset contains four domains: art, cartoon, photo, and sketch. We adopt the cave-one-domain-ee evaluation protocol [10], i.c., one domain is chosen as the held-out testing domain, and the remaining domains are regarded as source training domains. We follow the data preprocessing in [10) while additionally applying step partition to inject label shift. Each domain is divided into clients, leading to a total of 21 source clients and target clients. Each source client has 80% of its samples as training set and the remained 20% as testing set. Each client has 2 major classes and minor class, where the ratio of images per class is approximately 16 (the same as our CIFAR-10 experiments). Since there is already domain shift, we do not add corruptions.
Global model training We first train a global model with FedAvg 31 over the training sets of source clients. The global model is ResNet- 18 with ImageNet pretrained parameter (provided by torchvision). We train the global model for T 200 communication rounds with full participation (cohort size C = 21), local epochs E 1, learning rate rj 0.05 and batch size B 20.
ATP training We initialize the adaptation rates as all-zero vector, and optimize it over the validation sets of source clients. We optimize the adaptation rates for = 500 communication rounds with full participation (cohort size C 21), learning rate 7) = 0.5 and batch size B 200.
ATP testing
We test the optimized adaptation rates on each target client. We use batch size B = 200.
C.1.5 Algorithm details
Assignment matrix A In the main test, we mentioned that A E Dxd is 0- - 1 assignment matrix that maps each adaptation rate all to the indices of the 1-th module's parameters in w. Mathematically.
1, if the k-th parameter w belongs to the I-th module
Asi =
0, otherwise
28

---

If there are d = modules, each with 1, 2, and 3 parameters, so D = 1+2+3 6, the corresponding
assignment matrix will be
0
0
0
1
0
0
1
A
0
0
0
0
0
0
1
Computation
We did our experiments with single NVIDIA Tesla V100 GPU. However, our experiment should only require less than 2GB of GPU memory.
29

---

C.2
Compatibility to model architecture (RQ1)
In this part, we evaluate ATP with two more model architectures: 5-layer Shallow CNN as smaller model and ResNet-50 as larger model.
Table 5: ATP with different model architectures, accuracy (mean + s.d. %) target clients
Shallow CNN on CIFAR-10
ResNet-50 on CIFAR-I00
Method
Feature shill
Label shift
Hybrid shift
Avg. Rank
Feature shift
Label shift
Hybrid shift
Avg. Rank
No udaptation
64.39 0.18
69.33
0.37
61.99
0.47
7.3
45.31
+ 0.30
51.63
0.15
40.01
0.17
7.3
BN-Adapt
66.46 + 0.22
$4,99
0.38
50,40
0.43
7.0
47.75
= 0.29
34,85
= 0.26
30.31
= 0.09
7,3
SHOT
0.29
0.47
45.42
31.06
27.44
Tent
65.61 1: 0.24
50.12 + 0.25
45.91
0.49
8.7
45,91
0.46
31.34 = 0.11
27.93 = 0.31
8.3
MEMO
65.89 = 0.31
71.95
= 0.25
64.17 0.47
5.3
48.42 = 0.14
55.19
0.28
42.53
0.20
3.7
EM
61.74 1: 0.25
76.28
1 0.29
67.54
0.41
5.0
43.00
1 0.31
59.34
0.15
44,82
0.27
BBSF
56,92 0.53
75.99
0.44
66.64
0.53
6.3
37.26
0.64
56,97
0.20
40.09
Surgical
64.45 1: 0.12
73.75
1: 0.42
65.67
0.44
5.7
45.18
0.38
$4,83
1 0.26
42,50
0.33
6.7
hTt-hatch
2.3
46.82
ATP-online
67.13 = 0.17
78.56
0.32
71.52
0.51
1.0
49.08
0.26
61.86
0.25
49.51
0.23
1.0
From Table 5 observe that under the model architecture (and the new dataset), the perfor- mance of ATP is highly similar to the results the ResNet-18 + CIFAR1O experiment Table iin Subsection 6.1 ATP, in all three scenarios, handle various types of distribution shifts and surpass baseline methods. This suggests that ATP is compatible with multiple model architectures.
30

---

C.3 Robustness to global model
In this subsection, we design experiments to answer the following question: ATP robust to the choice of global model? Specifically, we have three sub-questions:
Is ATP robust to the parameter of global model? (C.31
Is ATP robust to the algorithm train global model? (C.3.2)
C.3.1 Robustness to the parameter of global model (online updated global model)
In the main text, we primarily focused on the scenario where the global model remains fixed. However, in practical FL systems, the global model may also undergo continuous online updates. Therefore, after obtaining the adaptation rates through ATP training, the global model might have been further updated for several rounds. This raises a question: Are the "outdated" adaptation rates still effective after several rounds of updates to the global model?
Table 6: Accuracy (%), ATP can learn adaptation rates that generalize global models with different numbers of communication rounds under hybrid shift on CIFAR-10
Method
200 + Rounds
+-10 Rounds
+020 Rounds
+550 Rounds
+-100 Rounds
No adaptation
63.68 0.24
63.88 0.20
64.03 0.13
64.30 + 0.08
64.56 0.11
ATP-batch
73.05 + 0.35
73.20 0.40
73.25 + 0.37
73.47 + 0.48
73.61 0,28
ATP-online
75.37 + 0.22
75.61 0.23
75.69 + 0.20
75.80 + 0.15
75.83 0.28
We design experiment apply the "outdated" adaptation rates to the global model that has undergone additional updates for several rounds, to see if they can still improve the test-time accuracy of the global model. Specifically, optimize the adaptation rates with w where
200, but the adaptation rates with w T+AT with AT = 10, 20, 50, 100 rounds. We use the same setting of hybrid shift on CIFAR-10 experiments. As shown in Table O while further optimizing the global model can marginally improve the accuracy, both ATP-batch and ATP-online can effectively enhance the test-time accuracy through personalization, even when is trained using outdated version of the global model.
C.3.2 Robustness to the algorithm to train global model
In the main text, we used FedAvg 31 to train the global model. However, in real-world FL systems, other Fl. algrrithms may be employed for training the global model, considering stability optimization or fairness. Therefore, aim to investigate whether ATP also be applied to other commonly used FL algorithms.
Table 7: Accuracy (%), ATP enhances different global models under hybrid shift on CIFAR-10
Method
FedAvg
FedProx (μ = 0.01)
q-FFL (q 1)
No adaptation
63.68 + 0.24
63.77 + 0.25
63.87 + 0.23
ATP-batch
73.05 1 0.35
72.95 1 0.33
73.15 + 0,21
ATP-online
75.37 + 0.22
75.51 1: 0.19
75.79 + 0.15
In particular, we use FedProx 23 an FL algorithm designed to handle heterogeneous setting, and
4-FFL [24], an FL algorithm enhancing performance fairness among participating clients. For all global model, we use same setting of hybrid shift on CIFAR- 10 experiments. As shown in Table 7 both ATP-batch and ATP-online can consistently improve the test-time accuracy across different FL algorithms to train global models.
31

---

C.4 Convergence and generalization
In Section 5 Appendix B.2 and we theoretically show that ATP has good convergence and
generalization guarantees. In this section, we visualize the training and testing loss curves to verify the fast convergence and superior generalization of ATP under different cohort size C. The results are shown in Figure
15
15
13
13-
1.3
1:
0.8
0.
0.9
200
Roand
Roand
Round
(a) C' = 240
(b) C' 120
(c) C 60
8
11
4
25
50
75
200 tand
325
150
175
100
125
faadd
(d) C = 30
(e) C' 15
Figure 8: Loss curves of ATP under different cohort size C
Convergence Under full participation (C 240), both the training and testing loss converge stably and fast, indicating the reliable convergence of ATP. With partial participation, as the cohort size decreases (C 120, 60, 30, 15), the training loss curve exhibits greater fluctuations, primarily due to sampling different subsets clients in each communication round. However, the testing loss curve still converge stably with similar speed, indicating that ATP is robust partial participation.
Generalization Under full participation (C = 240), the training and testing loss curves decrease synchronously without any overfitting. This implies that our algorithm exhibits excellent generaliza- tion. Similar observations can made partial participation (C = 120, 60, 30, 15). Additionally, it is worth noting that the test loss is lower than the train loss, which may seem counterintuitive. This is primarily due to the use of diffcrent corruptions between the testing and source clients. The accuracy of clients varies significantly under different corruptions, evidenced by the fluctuations in the training curve when comparing the trends of the two curves.
15. However, we can still analyze the generalization performance by
32

---

C.5
Toy example for negative adaptation rate (RQ2)
In Section 6.2 we notice that ATP learns negative adaptation rates for running means and variance under label shift. In this subsection, we use toy example to show why negative adaptation rate can improve the model performance under label distribution shift.
We consider binary classification problem with input x ε I. and binary output E [-1,+1}, where -1 is the negative class and +1 is the positive class. Let the feature for negative samples (,-1,
(xly = -1) ~
and for positive samples (æly = +1) ~ M(+1, 0.8²). Let the label
distribution Pr(y = 1) for training set, and Pr(y = 1) = à for testing set. Therefore, for the training distribution, we have
Ex = Pr(y = (2) = -1) + Pr(y = +1)E(zly = +1) =0
Var(#) = B[Var (x|)) + Var (B([ = 1.64
We consider a simple network with only one BN layer, with both normalization and affine transfor- mation (as a linear classifier). There are four modules, each is a scalar: running mean μt, running variance x², weight y, bias B3
Alter
Nlt SN
Nlw
120-
120
120-
60-
60-
(a) Train, Acc=0.89
(b) a = 1. Acc=0.73
(c) 0.5, Acc=0.83
BN
240
220
1e0
00
(d) a = O, Acc=0.89
(c) a = 0.5, Acc=0.92
Figure 9: Adapting batch norm running statisties under label shift.
Training During training, given enough training data, we have Merain = Ex 0 and orain 1.64. Figure Dalshows the histogram of - i... the intermediate feature after normalization before the transformation. By comparing the histograms of z of two classes, we notice that the optimal decision boundary is = 0, which indicate that Strain Merrain Prain "ttrin, Strain, and only update running statistics ptrain, Trran during testing.
and "trrin O. We store the corresponding
Testing without updating running statistics (a 0) Figure Pd shows the testing result when we do not update the running statistics, i.e., a = 0. Since two conditional feature distributions are symmetric, the accuracy will not change.
Testing with a > Positive adaptation rates align the intermediate feature distribution. When we use a = 1, the distribution of z will be centralized. As shown in Figure 9b] such alignment greatly reduces the accuracy. Similar result is also observed with any positive a, e.g., a 0.5 in Figure 9c
Testing with a < 0 While a = O has stable accuracy under label shift, by comparing the histograms of of two classes in Figure [d] we notice that : = 0 is not the optimal decision boundary anymore, because there are less negative samples than positive samples. By using negative adaptation rate < the normalization layer can further "disalign" the intermediate feature, which can further improve the accuracy, as shown in Figure De
33


Reviews:
[{"review_id": "Y5JcSx9NjC", "reviewer": "Reviewer_A9BH", "summary": "This paper introduces a novel setting where personalized FL during the test procedure is considered and multiple distribution shifts are involved. A method termed ATP is proposed to solve the challenges posed in this setting. Adaptive learning rates are learned for the model. Both theoretical and empirical studies are carried out to demonstrate the effectiveness of ATP.", "questions": "In Section 4.2, the authors have mentioned the smaller communication costs achieved by ATP, are there any corresponding experimental results that demonstrate this in a quantitive manner?", "limitations": "This paper does not adequately address the limitations in terms of privacy issues, efficiency, etc.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-organized and easy to follow. \n\n2. The main claims, e.g., the capability of dealing with multiple distribution shifts and the effectiveness of adaptive learning schemes, are well supported by empirical studies. \n\n3. Theoretical analyses on both the convergence and generalization ability of ATP are provided.", "weaknesses": "1. One of the key factors of TTPFL is confusing. Usually, we assume test data are unlabeled and the target of the classification task is to predict the label of the test data. However, in the summarization of TTPFL, it is emphasized that each testing client only has unlabeled data for personalization. It is better to further clarify this factor.  \n\n2. The relation between FL and test-time shift is weak. It seems that the proposed adapting trainable parameters, adapting running statistics, and adapting rates can also benefit centralized test-time shift problems. For FL systems, the unique challenges brought by test-time shift issues and how they motivate these adapting solutions are not explicitly demonstrated.\n\n3. Lack of detailed discussion on the difference between this work and the previous study[15]. [15] also considers feature shift, label shift, and a mixture of these shifts in their recent paper. There are various feature shifts in [15], which can also be considered as part of the experimental setting in this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer A9BH, \n\nWe would like to express our sincere gratitude for the thoughtful review and constructive feedback provided. We are grateful for your positive feedback on the paper's organization, empirical support, and theoretical analyses. We have carefully considered your comments and will address each concern in this rebuttal. \n\n## W1. TTPFL setting\n\n> 1. One of the key factors of TTPFL is confusing. Usually, we assume test data are unlabeled and the target of the classification task is to predict the label of the test data. However, in the summarization of TTPFL, it is emphasized that each testing client only has unlabeled data for personalization. It is better to further clarify this factor.\n\nThank you for pointing out this confusion. We emphasized that each testing client only has unlabeled data for personalization, in comparison to traditional PFL algorithms like Per-FedAvg [7], pFedMe [6], and FedTHE [15]. In these algorithms, testing clients have both labeled data (for personalization) and unlabeled data (for evaluation/prediction), which is a stronger requirement than ours. However, in the test-time of TTPFL, our ATP algorithm only requires unlabeled data for local personalization, utilizing the distribution of unlabeled data to adapt the global model and achieve better performance. We revised the description of TTPFL according to your suggestion in [our response to all reviewers, part 1](https://openreview.net/forum?id=rbw9xCU6Ci&noteId=PEt62AClVa). \n\n## W2. Relation between FL and test-time shift\n\n> 2. The relation between FL and test-time shift is weak. It seems that the proposed adapting trainable parameters, adapting running statistics, and adapting rates can also benefit centralized test-time shift problems. For FL systems, the unique challenges brought by test-time shift issues and how they motivate these adapting solutions are not explicitly demonstrated.\n\nThank you for pointing out this issue. We summarize the unique challenges of TTPFL in comparison with centralized TTA as follows: \n1. **Various distribution shifts**: Since each FL client collects its data in a distributed manner, the data can exhibit multiple distribution shifts, e.g., a complex combination of feature and label shifts. However, the complexity of distribution shifts have been overlooked in most centralized TTA works. As shown in Subsection 3.2, *most TTA algorithms cannot tackle feature and label shifts simultaneously*. However, our ATP algorithm can learn to tackle different types of distribution shifts as shown in Subsection 6.1. \n2. **Exploiting multiple data sources**: *centralized TTA considers adaptation from one source domain to one target domain, while our TTPFL includes multiple clients as data sources*. In TTPFL, centralized TTA algorithms adapt a global model which can only exploit training clients’ data as a mixed distribution $\\mathcal{P}_G = \\sum_i p_i \\mathcal{P}_i$, ignoring how they are different from each other. However, our ATP algorithm learns how each client’s distribution $\\mathcal{P}_i$ is different from the mixed distribution $\\mathcal{P}_G$, and optimizes the adaptation rates accordingly. \n\n## W3. Discussion of FedTHE\n\n> 3. Lack of detailed discussion on the difference between this work and the previous study[15]. [15] also considers feature shift, label shift, and a mixture of these shifts in their recent paper. There are various feature shifts in [15], which can also be considered as part of the experimental setting in this work. \n\nThanks for your suggestion for a detailed comparison with FedTHE [15]. Our paper is substantially different from FedTHE [15] in research problems, algorithms, and experiment design. We discuss our differences in detail in [our response to all reviewers, part 3](https://openreview.net/forum?id=rbw9xCU6Ci&noteId=PEt62AClVa). Specifically, regarding the experiment design, while FedTHE and our paper use similar techniques to construct feature and label shifts, we consider a stronger fusion of them. In FedTHE, the author simply mixed samples from different shifted distributions, which mitigates the severity of distribution shifts. For example, a smaller portion of testing samples are corrupted. As a result, “mixture of test” is less challenging than “corrupted local test” as shown in Table 1 of [15]. Noticing this deficiency, we improved the way of mixing feature and label shifts in our paper, making “hybrid shift” more challenging than both feature and label shifts (as shown in Table 1 of our paper). \n\n## Q1. Communication cost\n\n> In Section 4.2, the authors have mentioned the smaller communication costs achieved by ATP, are there any corresponding experimental results that demonstrate this in a quantitive manner?\n\nThanks for mentioning the advantages of ATP in communication efficiency. Unlike FedAvg [25] which transmits the full model parameters in every communication round, training clients in ATP only download the full model once, and only transmit the adaptation rates during the training process. Therefore, for $T$ communication rounds, ATP only transmits $D+2Td$ floating numbers, much fewer than $2TD$ for FedAvg. Notice that $d \\ll D$: For our ResNet-18 experiments on CIFAR-10, $d=102$ while $D=11,181,642$. For our ResNet-50 experiments on CIFAR-100, $d=267$ while $D=23,581,642$. \n\n## Limitations\n\n> This paper does not adequately address the limitations in terms of privacy issues, efficiency, etc.\n\nWe appreciate your comment. Regarding privacy, while our ATP framework shares the same communication protocol as FedAvg while transmitting fewer parameters, we did not further study the privacy of our framework. We believe differential privacy techniques can further protect the privacy of clients in ATP systems.  \n"}, {"review_id": "dHf5A28vkJ", "reviewer": "Reviewer_h3BM", "summary": "This paper proposes a new setting called test-time personalized federated learning (TTPFL) and proposes an Adaptive Test-time Personalization algorithm. The authors show effectiveness of proposed method over other test-time adaptation methods.", "questions": "See weakness.", "limitations": "Generally, I am really confused about the setting in this paper after reading. Please provide comprehensive explanations to correct me if I am wrong,  and I would consider re-rating.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper proposes an Adaptive Test-time Personalization algorithm and shows its effectiveness over other test time adaptation methods.", "weaknesses": "1. The proposed setting is strange and not self-consistent. The authors claim that in this setting 'clients adapt a trained global model in an unsupervised manner without requiring any labeled data.' However, the proposed method involves labeled clients to learn the learning rates for different modules (Figure 3, left). At least, real-world scenarios as examples should be provided.\n2. Missing representative federated learning baselines. This paper does not compare with any federated learning method. If the proposed method uses the labeled datasets for the training stage, then many federated learning methods can be seen as baselines that are trained on the labeled datasets. Baseslines include FedAvg, FedAvg with Fine-tuning, FedProx, FedProx with finetuning, Ditto, pFedMe.\n3. Since the proposed method focuses on different operations on different modules of a model, model architecture is important to this paper, thus experiments on different model architectures are required.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer h3BM, \n\nThank you for your insightful review and valuable feedback on our paper. We sincerely appreciate your recognition of the effectiveness of our algorithm. Regarding the weaknesses, we address them point by point as follows. \n\n## W1. TTPFL setting\n\n> 1. The proposed setting is strange and not self-consistent. The authors claim that in this setting 'clients adapt a trained global model in an unsupervised manner without requiring any labeled data.' However, the proposed method involves labeled clients to learn the learning rates for different modules (Figure 3, left). At least, real-world scenarios as examples should be provided.\n\nThanks for pointing out this confusion of labeled and unlabeled clients. We would like to clarify that our TTPFL setting has the same data requirement as standard FedAvg [25]. In FedAvg, the global model is trained on training clients with labeled data, and then tested on testing clients with only unlabeled data. Similarly in TTPFL, during training, the global model and the adaptation rule are optimized over labeled training clients; while during testing, each testing client downloads the global model and the adaptation rule, and personalizes the global model with only its unlabeled data. We revised the description of TTPFL according to your suggestion in [our response to all reviewers, part 1](https://openreview.net/forum?id=rbw9xCU6Ci&noteId=PEt62AClVa). \n\n## W2. More baselines\n\n> 2. Missing representative federated learning baselines. This paper does not compare with any federated learning method. If the proposed method uses the labeled datasets for the training stage, then many federated learning methods can be seen as baselines that are trained on the labeled datasets. Baseslines include FedAvg, FedAvg with Fine-tuning, FedProx, FedProx with finetuning, Ditto, pFedMe.\n\nThanks for your suggestion of comparing to FL baselines. We compared our algorithm to all the baselines you mentioned and summarized the experiment results in [the rebuttal pdf, part A](https://openreview.net/attachment?id=PEt62AClVa&name=pdf). \n\nWe would like to clarify that although our TTPFL setting requires labeled data for training clients (same as FedAvg), it does not require any labeled data for testing clients. On the contrary, most of the existing PFL algorithms either focus on the training clients exclusively (e.g., Ditto) or require labeled data for personalization on the testing clients (e.g., Fine-tuning, pFedMe). These algorithms introduce stronger data requirements to FL systems, and cannot be used for our TTPFL. \n\nTo compare to these baselines, we keep the training phase of these algorithms, while using “pseudo-labels” for personalization on testing clients according to [R1]. The pseudo-label is the prediction of the global model. Experiment results in [Table A](https://openreview.net/attachment?id=PEt62AClVa&name=pdf) shows that most FL and PFL baselines can only introduce limited improvement to the testing clients. Meanwhile, our proposed ATP-Episodic outperforms all FL baselines in the setting of TTPFL across three types of distribution shifts. \n\n[R1] Dong-Hyun Lee. Pseudo-label: The simple and efficient semi-supervised learning method for deep neural networks. Workshop on challenges in representation learning, ICML 2013. \n\n## W3. Different model architectures\n\n> 3. Since the proposed method focuses on different operations on different modules of a model, model architecture is important to this paper, thus experiments on different model architectures are required.\n\nThanks for your suggestion to include different model architectures. In our paper, we tried two architectures, ResNet-18 and ResNet-50, which are representative in FL [36,43,44] and TTA [14,39,45]. Experiment results show that ATP is compatible and effective with different model architectures. Please refer to Appendix C.2 for our experiments with ResNet-50. \n"}, {"review_id": "BSKQhzWVll", "reviewer": "Reviewer_2b9x", "summary": "The paper studies test-time personalization in a federated learning setting --- after training on participating clients, the goal is to locally adapt the  global model given unlabeled test data. The paper's main idea is by pointing out that label non-IID and domain non-IID require adaptation on different layers of DNNs, and propose a novel way to learn the adaptation learning rates of each layer automatically in a data-driven fashion.  A simple learning method that alternatively do SGD on the DNN parameters and the layer-wise learning rates shows effective improvements in test-time personalization.  ", "questions": "Section 3.2 is the limitation to TTA but not about PFL. Can the author discuss why the problem or the solution is dedicated to FL? I suggest adding a discussion around L134. In my opinion, the key is that in PFL we can figure out the ideal adaptive rates by leveraging the training clients who have diverse distributions. ", "limitations": "See weakness.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The method proposes a novel aspect that to tailor different types of non-IIDness, the degrees of adaptation are different across layers. it is an interesting point. The solution is simple, sound, and effective. \n\n- Empirical results are overall satisfying. The experiments provide enough comparisons to centralized TTA algorithms; see some suggestions below. \n", "weaknesses": "Although I am positive on this paper, I observe several important concerns. If the authors could address them, my score can be higher. \n\n[Major 1] Overclaims: considering TTA in PFL setting was first introduced in [15]. This is a natural extension of centralized TTA. It is unnecessary and imprecise for this paper to make \"test-time personalized federated learning (TTPFL)\" a new setting. The definition in L38 about the combination of distribution shifts of labels and styles itself is nothing to do with federated setting; centralized TTA can have both label and style shifts.\n\n[Major 2] Misleading section 4.2. I cannot understand why the proposed refinement has a connection with meta-learning. ClientTrain in Alg 1 is simply a coordinate gradient style optimization. Alternatively, a discussion about hyperparameter optimization should be more relevant.\n \n[Major 3] Theorem 5.1 is not informative. It's simply extending the classic FedAvg generalization bound [18] for the adaptive parameters. It will be more informative and fit to the context to discuss the generalization to the new clients of different distributions ''after test-time personalization'' in TTA sense. Otherwise, I would recommend avoid this laundry theorem. \n\n[Minor 1] The datasets are rather small scales and synthetic. it will be better to include natural federated datasets like FEMNIST or iNaturalist-GEO or large-scale common datasets in TTA (like ImageNet).\n\n[Minor 2] Why FedTHE or FedTHE+ in [15] are not compared? [15] seems to be the closest work about TTA in FL. Is the discussion in L76 faithful? It seems to me FedTHE does not need labeled data.  \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 2b9x, \n\nThank you for your detailed and insightful review of our paper. We appreciate your positive feedback on the novelty and empirical results of our work. We have carefully considered your concerns and suggestions, and will address them point by point as follows. \n\n## W1. Overclaims\n\n> [Major 1] Overclaims: considering TTA in PFL setting was first introduced in [15]. This is a natural extension of centralized TTA. It is unnecessary and imprecise for this paper to make \"test-time personalized federated learning (TTPFL)\" a new setting. \n\nThanks for your question regarding the comparison of our work and FedTHE [15]. We would like to clarify that although our setting has a similar name to theirs, our problem is substantially different. FedTHE focuses on improving the model robustness against test-time shift on seen labeled training clients (i.e., clients that participate in FL training with labeled data), while our algorithm focuses on better generalization to unseen unlabeled testing clients (i.e., clients that neither participate in training nor have labeled data) with different distributions. We compare our work with FedTHE in detail in [our response to all reviewers, part 1](https://openreview.net/forum?id=rbw9xCU6Ci&noteId=PEt62AClVa). \n\n## W2. Connection to meta-learning & hyperparameter optimization\n\n> [Major 2] Misleading section 4.2. I cannot understand why the proposed refinement has a connection with meta-learning. ClientTrain in Alg 1 is simply a coordinate gradient style optimization. \n\nThanks for raising this confusion of the connection between supervised refinement and meta-optimization. The ClientTrain in Alg 1 shares similarity with important meta-learning algorithms, e.g., MAML [8]. \n- During unsupervised adaptation, each client conducts coordinate gradient style optimization, which resembles adapting the meta-model on a specific task in MAML (line 6 in their algorithm 1). \n- During supervised refinement, each client updates the adaptation rates to minimize the post-TTA loss, which resembles updating the meta-model in MAML (line 7 in their algorithm 1). It is important to notice that in line 16 of ClientTrain, we compute the gradient w.r.t. adaptation rates $\\boldsymbol{\\alpha}$, not the personalized parameter $\\boldsymbol{w}_{ij}$. \n\n> Alternatively, a discussion about hyperparameter optimization should be more relevant.\n\nWe agree with you that a discussion about hyperparameter optimization can also be relevant. [R1] first investigated the problem of federated hyperparameter tuning and proposed FedEX that leverages weight-sharing from neural architecture search to efficiently tune hyperparameters. [R2] introduced FloRA that addresses use cases of tabular data and enables single-shot federated hyperparameter tuning. While these methods focus on improving the efficiency of hyperparameter optimization, our paper focuses on finding the optimal adaptation rates that benefit test-time personalization. We will include this discussion in the revision of our paper. \n\n[R1] Mikhail Khodak, et al. Federated Hyperparameter Tuning: Challenges, Baselines, and Connections to Weight-Sharing. NeurIPS 2021. \n\n[R2] Yi Zhou, et al. Single-shot General Hyper-parameter Optimization for Federated Learning. ICLR 2023. \n\n## W3. Theorem 5.1\n\n> [Major 3] Theorem 5.1 is not informative. It's simply extending the classic FedAvg generalization bound [18] for the adaptive parameters. It will be more informative and fit to the context to discuss the generalization to the new clients of different distributions ''after test-time personalization'' in TTA sense. Otherwise, I would recommend avoid this laundry theorem.\n\nThanks for your attention regarding our generalization analysis. We would like to clarify that our Theorem 5.1 actually considers the error rate **after test-time personalization**, instead of before it. As defined in Definition B.12 in Appendix B, we consider the error rate of the adapted model $\\boldsymbol{w}_{ij}$, instead of the global model $\\boldsymbol{w}_G$. \n\nTheorem 5.1 aims to show that if ATP achieves low post-TTA classification error over training clients after refining the adaptation rates, we are expected to get a similar low post-TTA classification error on testing clients. \n\n## W4. More datasets\n\n> [Minor 1] The datasets are rather small scales and synthetic. it will be better to include natural federated datasets like FEMNIST or iNaturalist-GEO or large-scale common datasets in TTA (like ImageNet). \n\nThanks for your suggestion on the improvement of our experimental setup. We further conduct experiments on PACS, a large-scale common datasets in TTA with resolution $224\\times 224$. The result is shown in [the rebuttal pdf, part B](https://openreview.net/attachment?id=PEt62AClVa&name=pdf), where ATP consistently outperforms all baseline algorithms. \n\n## W5. Comparison to FedTHE\n\n> [Minor 2] Why FedTHE or FedTHE+ in [15] are not compared? [15] seems to be the closest work about TTA in FL. Is the discussion in L76 faithful? It seems to me FedTHE does not need labeled data.\n\nThanks for your question. We would like to clarify that we did not compare to FedTHE [15] because it requires labeled data. **FedTHE trains a model with two heads (global head and personalized head) during FL training in a supervised manner**, and fuses two heads during test-time in an unsupervised manner. When generalizing to unseen unlabeled testing clients, the client can only download the global head but cannot generate its personalized head due to the lack of labeled data. Therefore, FedTHE cannot be used for unseen unlabeled testing clients, which is the target of our paper. \n\nIn our [rebuttal pdf, part A](https://openreview.net/attachment?id=PEt62AClVa&name=pdf), we design a variant of FedTHE using pseudo-labels to train the personalized head. Our algorithm outperforms this FedTHE variant across different types of distribution shifts. \n\n\n"}, {"review_id": "Uo48sM4PCk", "reviewer": "Reviewer_yJjv", "summary": "This paper considers the federated learning setting of adaptive test-time personalization. Traditional test-time adaptation (TTA) can only handle specific target domain distributions, while federated learning requires flexible handling of multiple target domains. Existing TTA methods pre-define which modules to adapt, which limits the application of TTA in federated learning. Therefore, this paper proposes the Adaptive Test-time Personalization algorithm called ATP to automatically decide which modules to adapt and how much to adapt.", "questions": "1. In Figure 2, why does setting \"bn.running_mean (m=-0.1)\" improve the accuracy of TTA in the presence of label shift, while \"bn.running_mean (m=0.1)\" significantly harms the accuracy? This seems counterintuitive.\n2. The task addressed in this paper is more akin to source-free unsupervised domain adaptation rather than Test-Time adaptation. In my opinion, the focus of most TTA works is on an online setting where the entire test set cannot be obtained at once, and operations are performed on individual test samples to be predicted (as the details in Section 4.3).\n3. See Weaknesses 2, how important is the step of \"supervised refinement\"? Can the authors provide any ablation experiments?\n", "limitations": "There is no limitations or potential negative societal impact discussed in this paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The experiment in Section 3.2 is crucial as it effectively illustrate the inherent challenges encountered by current TTA methods when applied to federated learning, thus offering valuable insights and guiding directions to potential enhancements.\n2. The proposed method is simple and easy to implement, and it has achieved impressive performance in the current experiments.", "weaknesses": "1. The paper claims that it is the first to propose test-time personalized federated learning. However, this claim is questionable because previous works, such as [1], had already explored test-time personalized tasks in the context of federated learning. \n[1] Jiang, Liangze, and Tao Lin. \"Test-Time Robust Personalization for Federated Learning.\" ICLR2023 (preprint arXiv:2205.10920 (2022))\n2. I am concerned about the significance of the “supervised refinement” step. If labeled data is available in this step, using the labeled data itself already leaks the distribution about the test samples in TTA tasks (while the test distribution is not known in the TTA setting), which obviously reduces the difficulty of TTA. If labeled data is not available, the proposed method for adaptively learning alpha seems unworkable.\n3. The datasets used in the experiments are small-scale. It would be more convincing if experiments were conducted on real-world images, such as DomainNet, with at least a resolution of 224x224.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer yJjv, \n\nWe sincerely thank you for your comprehensive review and insightful feedback on our paper. We appreciate your positive feedback on the crucial experiment and effective algorithm of our work. We would like to address your concerns as follows.\n## W1. Comparison to FedTHE\n> 1. The paper claims that it is the first to propose test-time personalized federated learning. However, this claim is questionable because previous works, such as [15], had already explored test-time personalized tasks in the context of federated learning. \nWe would like to clarify that our proposed TTPFL, is substantially different from the setting in FedTHE [15]. FedTHE focuses on improve the model robustness to test-time shifts on seen labeled training clients (i.e., clients that participate in FL training with labeled data), while our algorithm focuses on better generalization to unseen unlabeled testing clients (i.e., clients that neither participate in training nor have labeled data) with different distributions. Our research problems are orthogonal. We discuss our differences in detail in [our response to all reviewers, part 3](https://openreview.net/forum?id=rbw9xCU6Ci&noteId=PEt62AClVa).\n## W2. Significance of supervised refinement\n> 2. I am concerned about the significance of the “supervised refinement” step. If labeled data is available in this step, using the labeled data itself already leaks the distribution about the test samples in TTA tasks (while the test distribution is not known in the TTA setting), which obviously reduces the difficulty of TTA. \n\nThanks for raising this concern. We would like to clarify that labeled data is only available for training clients, which participate in FL training. Neither the image nor the label for testing clients are used for supervised refinement. After training, each unseen testing client personalizes the global model in an unsupervised manner with the trained adaptation rates and test the adapted model’s performance. We report the average accuracy on testing clients.\n\n[Our response to all reviewers, part 1](https://openreview.net/forum?id=rbw9xCU6Ci&noteId=PEt62AClVa) could be helpful in further clarifying this issue.\n## W3. Large-scale dataset\n>3. The datasets used in the experiments are small-scale. It would be more convincing if experiments were conducted on real-world images, such as DomainNet, with at least a resolution of 224x224.\n\nThanks for your valuable advice. We further conduct experiments on PACS, a large-scale common datasets in TTA with resolution $224\\times 224$. The result in the [rebuttal pdf, part B](https://openreview.net/attachment?id=PEt62AClVa&name=pdf) shows that ATP consistently outperforms all baseline.\n## Q1. Figure 2\n> 1. In Figure 2, why does setting \"bn.running_mean (m=-0.1)\" improve the accuracy of TTA in the presence of label shift, while \"bn.running_mean (m=0.1)\" significantly harms the accuracy? This seems counterintuitive.\n\nThanks for your insightful question. We used a visual example in Appendix C.4 to explain why different momentum m has opposite effects when adapting to label shift (m in Figure 2 refers to $\\alpha$ in Appendix C.4). In brief, adapting the running mean can be seen as feature aligner/disaligner.\n- Positive m aligns the intermediate feature distributions for training and testing data, resulting in aligned distribution of predictions. However, since the label distribution is shifted, such alignment introduces negative effects.\n- Meanwhile, negative m disaligns the intermediate feature distribution, which follows the change in label distribution and improves the accuracy.\n## Q2. Source-free unsupervised domain adaptation\n> 2. The task addressed in this paper is more akin to source-free unsupervised domain adaptation rather than Test-Time adaptation. In my opinion, the focus of most TTA works is on an online setting where the entire test set cannot be obtained at once, and operations are performed on individual test samples to be predicted (as the details in Section 4.3).\n\nThanks for your comment. We would like to clarify that our proposed ATP does not obtain the entire test set at once. Instead, \n- For ATP-Episodic, testing clients process each test batch independently, resembling test-time batch adaptation according to [R1].\n- For ATP-Online, testing clients process test batches sequentially like Tent [39], resembling online test-time adaptation.\n\nMoreover, we would like to emphasize that our ATP is substantially different from previous works in TTA since ATP learns from multiple FL clients to tackle different types of distribution shifts.\n\n[R1] Jian Liang, Ran He, Tieniu Tan. A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts.\n## Q3. Supervised refinement\n> 3. See Weaknesses 2, how important is the step of \"supervised refinement\"? Can the authors provide any ablation experiments?\n\nThanks for your question. Supervised refinement is important in ATP since it learns a set of adaptation rates specific to the type of distribution shifts among FL clients. Without supervised refinement, the adaptation rates will be initialized as zero, thus ATP becomes identical to “no adaptation”. For ablation studies, we try (1) using constant adaptation rates for all modules, and (2) fitting adaptation rates with a different meta-distribution, e.g., fitted on clients with feature shift but tested on clients with label shift. As shown in [our rebuttal pdf, Table C](https://openreview.net/attachment?id=PEt62AClVa&name=pdf), \n- When fitted on the *same type* of distribution shift, ATP significantly improves the performance on testing clients and outperforms constant adaptation rates. It verifies the importance of supervised refinement. \n- When fitted with a *different type* of distribution shift, the effectiveness of ATP noticeably decreases. It proves that supervised refinement can learn distribution-shift-specific adaptation rates, which further validate the importance of supervised refinement. "}]

Author rebuttal:
Dear Reviewers, 

We would like to express our sincerest gratitude for your thoughtful and insightful reviews of our paper. We are particularly grateful for the recognitions bestowed upon our work, including
- novel and effective algorithm (from all reviewers), 
- satisfying experiments (from reviewer A9BH, 2b9x, yJjv), and
- comprehensive convergence and generalization analysis (from reviewer A9BH).

Your insights have been immensely helpful in refining our work. In this response, we address the common questions raised by your reviews. 

## 1. Revised description of test-time personalized federated learning (TTPFL)
We consider an FL system with $N$ training clients (who participate in FL training, i.e., "seen" during training) and $M$ testing clients (who do not participate, i.e., "unseen"). Each training client $i$ has its labeled dataset $S_i$ with $m_i$ i.i.d. data points drawn from its own underlying distribution $\mathcal{P}_i$, while each testing client $i$ has only unlabeled data drawn from its distribution $\mathcal{P}_i$. The data distribution $\mathcal{P}_i$ is usually different across the clients, and is sampled from a meta-distribution $\mathcal{Q}$, i.e., a distribution of distributions. **Our data requirement is identical to FedAvg [25] and FedSR [26].**

During FL training, the global model $w_G$ and the adaptation rule $\mathcal{A}$ (i.e., adaptation rates in ATP) are optimized only over labeled training clients. Testing clients’ data is not used during training. During FL testing, each testing client downloads the global model and the adaptation rule, and personalizes the global model with only its unlabeled data. The testing clients’ labels are never used in FL training or testing.

It is important to notice that **our TTPFL setting does not require any labeled data for testing clients**, unlike the setting in traditional personalized federated learning (PFL). Most of the existing PFL algorithms either focus on the training clients exclusively [35,19] or require labeled data for personalization on the testing clients [7,5]. These algorithms introduce stronger data requirements to FL systems, and cannot be used for our TTPFL.

Considering the enormous number of testing clients (compared to training clients) in real FL applications [16] and the challenge of generalization to testing clients [43], we pay special attention to the performance on the testing clients, and **report the average accuracy on testing clients**, instead of training clients (as in traditional PFL algorithms).

## 2. Relation between FL and test-time adaptation (difference between TTPFL and TTA)
Different from centralized TTA, TTPFL has its own challenges, which TTA algorithms do not fully tackle.
1. **Various distribution shifts**: Since each FL client collects its data in a distributed manner, the data can exhibit multiple distribution shifts, e.g., a complex combination of feature and label shifts. However, the complexity of distribution shifts have been overlooked in most centralized TTA works. As shown in Subsection 3.2, *most TTA algorithms cannot tackle feature and label shifts simultaneously*. However, our ATP algorithm can learn to tackle different types of distribution shifts as shown in Subsection 6.1.
2. **Exploiting multiple data sources**: *centralized TTA considers adaptation from one source domain to one target domain, while our TTPFL includes multiple clients as data sources*. In TTPFL, centralized TTA algorithms adapt a global model which can only exploit training clients’ data as a mixed distribution $\mathcal{P}_G = \sum_i p_i \mathcal{P}_i$, ignoring how they are different from each other. However, our ATP algorithm learns how each client’s distribution $\mathcal{P}_i$ is different from the mixed distribution $\mathcal{P}_G$, and optimizes the adaptation rates accordingly.

## 3. Difference to FedTHE/FedTHE+ [15]
Despite having a similar name, our paper tackles a problem substantially different from FedTHE.
- **Problem**: FedTHE focuses on improving the model robustness to test-time shift on seen labeled training clients (i.e., clients that participate in FL training with labeled data), while our algorithm focuses on better generalization to unseen unlabeled testing clients (i.e., clients that neither participate in training nor have labeled data) with different distributions. Our research problems are orthogonal.
- **Algorithm**: FedTHE trains a model with two heads (global head and personalized head) during FL training in a supervised manner, and fuses two heads during test-time in an unsupervised manner. When generalizing to unseen unlabeled testing clients, the client can only download the global head but cannot generate its personalized head due to the lack of labeled data. Therefore, FedTHE algorithm cannot be used for unseen unlabeled testing clients, which is the target of our paper. Different from FedTHE, testing clients in ATP download the global model as well as the adaptation rates, and adapt the global model locally with its unlabeled data. Labels are not required in ATP testing.
- **Experiment**: While FedTHE and our paper use similar techniques to construct feature and label shifts, we consider a stronger fusion of them. In FedTHE, the author simply mixed samples from different shifted distributions, which mitigates the severity of distribution shifts. For example, a smaller portion of testing samples are corrupted. As a result, “mixture of test” is less challenging than “corrupted local test” as shown in Table 1 of [15]. Noticing this deficiency, we improved the way of mixing feature and label shifts in our paper, making “hybrid shift” more challenging than both feature and label shifts (as shown in Table 1 of our paper).

## 4. New experiments
In the rebuttal pdf, we conduct experiments regarding
1. comparison to FL baselines (including FedTHE)
2. new experiments on PACS, a large-scale dataset, and
3. additional ablation study. 

Details are provided in the pdf. 

Meta-review: Initially, this submission received mixed scores. Several reviewers were confused about the problem setting and the proposed solution. The authors addressed this, as well as concerns about the similarity to prior work, runtime and data usage. On the positive side, the submission provides a new solution to a relevant problem, which is backed up by extensive experimental evaluation as well as some theoretical results about generalization as well as convergence. Ultimately, a majority of the reviewers recommended accepting the work.
Decision: Accept (poster)