Test-Time Robust Personalization for Federated Learning
Authors: Liangze Jiang
Tao Lin
Venue: ICLR 2023 poster (2023)
Keywords: Federated Learning
Personalized Federated Learning
Test-time Robustness

Abstract: Federated Learning (FL) is a machine learning paradigm where many clients collaboratively learn a shared global model with decentralized training data. Personalization on FL models additionally adapts the global model to different clients, achieving promising results on consistent local training & test distributions. However, for real-world personalized FL applications, it is crucial to go one step further: robustifying FL models under the evolving local test set during deployment, where various types of distribution shifts can arise. In this work, we identify the pitfalls of existing works under test-time distribution shifts and propose Federated Test-time Head Ensemble plus tuning (FedTHE+), which personalizes FL models with robustness to various test-time distribution shifts. We illustrate the advancement of FedTHE+ (and its degraded computationally efficient variant FedTHE) over strong competitors, for training various neural architectures (CNN, ResNet, and Transformer) on CIFAR10 and ImageNet and evaluating on diverse test distributions. Along with this, we build a benchmark for assessing the performance and robustness of personalized FL methods during deployment. Code: \url{https://github.com/LINs-lab/FedTHE}.


Published as a conference paper at ICLR 2023
TEST-TIME ROBUST PERSONALIZATION FOR FEDER- ATED LEARNING
Liangze Jiang?., Tao Lin
liangze. @ww.ggianneei.. ch; lint lintao@eestlake edu. en RReeaarch Center for Industries of the Future, Westlake University
2EPFL
ABSTRACT
Federated Learning (FL) is machine learning paradigm where many clients collaboratively learn shared global model with decentralized training data. Per- sonalization on FL model additionally adapts the global model to different clients, achieving promising results on consistent local training & test distributions. How- ever, for real-world personalized FL applications, it is crucial further: robustifying FL models under the evolving local test set during deployment, where various types of distribution shifts can arise. In this work, we identify the pitfalls of existing works under test-time distribution shifts and propose Feder- ated Test-time Head Ensemble plus tuning (FedTHE+), which personalizes FL models with robustness to various test-time distribution shifts. We illustrate the advancement of FedTHE+ (and its degraded computationally efficient variant FedTHE) over strong competitors, for training various neural architectures (CNN, ResNet, and Transformer) on CIFARIO and ImageNet and evaluating on diverse test distributions. Along with this, we build a benchmark for assessing the per- formance and robustness of personalized FL. methods during deployment. Code: https: /ttt://gitttt com/ LINs lab/ /FedTHE.
go one step
INTRODUCTION
Federated Learning (FL) is an emerging ML paradigm that many clients collaboratively learn shared global model while preserving privacy (McMahan et al., 2017: Lin et al., 2020b; Kairouz et al., 2021: Li et al., 2020a). As variant, Personalized FL (PFL) adapts the global model to personalized model for each client, showing promising results on In-Distribution (ID).
However, such successes of PFL may not persist during the deployment for FL, as the incoming local test samples are evolving and various types of Diisti-Distiiuttion (OOD) shifts can occur (compared local training distribution). Figure showcases some potential test-time distribution shift scenarios, e.g., label distribution shift (c.f. (i) & (ii) and co-variate shift (c.f. (iii):: (i) clients may encounter new local test samples in unseen classes; (ii) even no unseen classes emerge, the local test class distribution may become different; (iii) as another real-world distribution shift, local new test samples may suffer from common corruptions (Hendrycks & Dietterich, 2018) (also called synthetic distribution shift) natural distribution shifts (Recht et al., 2018). More crucially, the distribution shifts can appear in a mixed manner, i.e. new test samples undergo distinct distribution shifts, making robust FL deployment more challenging.
Making FL more practical requires generalizing FL models to both ID & OOD and properly estimating their deployment robustness. To this end, we first construct new PFL benchmark that mimics ID & OOD cases that clients would encounter during deployment, since previous benchmarks cannot measure robustness for FL. Surprisingly, there is a significant discrepancy between the mainstream PFL. works and the requirements for real-world FL deployment: existing works (McMahan et al., 2017; Collins et al., 2021; Li et al., 2021b; Chen & Chao, 2022: Deng et al., 2020a) suffer from severe accuracy drop under various distribution shifts, as shown in Figure 2(a).
Test Time Adaptation (TTA)-being orthogonal to training strategies for OOD generalization has shown great potential in alleviating the test-time distribution shifts. However, they were designed for homogeneous and non-distributed settings, current TTA methods offer limited performance gains in FL-specific distribution shift scenarios (as shown in Figure 2(b)). especially in the label distribution shift case that is more common and crucial in non-i.i.d. setting.
"Equal contribution. Correspondence to: Tao Lin.

---

Published as a conference paper at ICLR 2023
Co-veriate shift [Common corruptions)
Co-variate shift (Natural distribution shift)
Chai Labels
dleg ass distribution has changed
Class distribution remains unchanged
Commio sorruptionss / Natural distribution shift
30 11
12 13 14
Cass that unseen during training
Different scenarios levalving test distribution) for FL during test time
Figure 1: Potential distribution shift scenarios for FL during deployment, e.g., (1) new test samples with unseen labels; (2) class distribution changes within seen labels; (3) test samples suffer from co-variate shifts, either from common corruptions or noturally shifted datasets. In summary, Car & Boat: label distribution shift; Dog: unchanged; Bird: co-variate shift.
Common comuptore
Acouracy gan
Natural di shitt
Label disoitution shit
Assurany
60
10
40
FedNwg+FT FedRep
APFL
FedRep
KNN
APFL
MEMO
FeOTHE-
(a) The limitations of PFL & naively adapted TTA
(b) The accuracy gain of baselines and FedTHE+ (ours)
methods test-time distributions shifts in FL.
compared FedAvg+FT ID & OOD cases.
Figure 2: Neither existing PFL methods nor applying TTA methods PFL is sufficient to handle the issues.
As remedy and our key contribution, we propose to personalize and robustify the FL model by our computationally efficient Federated Test-time Head Ensemble plus tuning (FedTHE+)): we unsuper- visedly and adaptively ensemble a global generic and local personalized classifiers of two-head FL model for single test-sample and then conduct an unsupervised test-time tne-tuning. We show that our method significantly improves accuracy and robustness on ID 4 OOD test distributions, via extensive numerical investigation on strong baselines (including FL., PFL., and TTA methods), models (CNN, ResNet, and Transformer), and datasets (CIFARIO 10 and ImageNet). Our main contributions are:
We revisit the evaluation of PFL and identify the crucial test-time cest-time distribution shift issues: a severe gap exists between the current PFL methods in academia and real-workde deployment needs.
We propose a novel test-time robust personalization framework (FedTHE/FedTHE+) with superior ID & OOD accuracy (throughout baselines, neural networks, datasets, and test distribution shifts).
As a side product to the community, we provide the first PFL benchmark considering comprehen- sive list of test distribution shifts and offer the potential to develop realistic and robust PFL methods.
RELATED WORK
We give a compact related work here due to space limitations. A complete discussion is in Appendix Federated Learning (FL). Most FL works focus on facilitating learning under non-i.i.d. client training distribution, leaving the crucial test-time distribution shift issue (our focus) unexplored. Personalized FL (PFL). The most straightforward PFL method is locally fine-tune the global model (Wang et al., 2019; Yu et al., 2020). Apart from these, Deng et al. (2020a); Mansour et al. (2020); Hanzely & Richtárik (2020); Gasanov et al. (2021) linearly interpolate the locally learned model and the global model for evaluation. This idea is extended in Huang et al. (2021); Zhang et al. (2021b) by considering better weight combination strategies over different local models. Inspired by representation learning, Arivazhagan et al. (2019); Collins et al. (2021); Chen & Chao (2022); Tan al. (2021) suggest decomposing the neural network into shared feature extractor and personalized head. FedRoD (Chen & Chao, 2022) uses a similar two-headed network as ours, and explicitly decouples the local and global optimization objectives. pFedHN (Shamsian et al., 2021) and ITU-PFL (Amosy et al., 2021) both use hypernetworks, whereas ITU-PFL focuses on the issue of unlabeled new clients and is orthogonal to our test-time distribution shift setting.
Note that the distribution, lacking the resilience to test-time distribution shift issues (our contributions herein).
PFL methods only focus on better generalizing the local training
2

---

Published as a conference paper at ICLR 2023
OOD generalization in FL. Investigating OOD generalization for FL is timely topic. Distribution shifts might occur either geographically or temporally (Koh et al., 2021; Shen et al., 2021). Existing works merely optimize a distributionally robust global model and cannot achieve impressive test-time accuracy. For example, GMA (Tenison et al., 2022) uses a masking strategy on client updates to learn the invariant mechanism across clients. A similar concept can be found in fairness- FL (Mohri et al., 2019; Li et al., 2020c; 202 Ib;a; Du et al., 2021; Shi et al., 2021; Deng et al., 2020b; Wu et al., 2022a), in which accuracy stability for the global model is enforced across local training distributions. However, the real-world common corruttions (Hendrycks & Dietterich, 2019b), natural distribution shifts (Recht et al., 2018; 2019; Taori et al., 2020; Hendrycks et al., 2021b;a) or label distribution shift in FI., is underexplored. We fill this gap and propose FedTHE and FedTHE+ as remedies.
Benchmarking FL with distribution shift. Existing FL benchmarks (Ryffel et al., 2018; Caldas et al., 2018; He et al., 2020a; Yuchen Lin et al., 2022: Lai et al., 2022: Wu et al., 2022b) do not properly measure the test accuracy under various ID & OOD cases for PFL. Other OOD benchmarks, Wilds (Koh et al., 2021; Sagawa et al., 2021) and DomainBed (Gulrajani & Lopez-Paz, 2020), primarily concern the generalization over different domains, and hence are not applicable for PFL. Our proposed benchmark for robust FL (BRFL) in section 5 fills this gap, by covering scenarios specific to FL made up of the co-variate and label distribution shift. Note that while the concurrent work (Wu al., 2022b) similarly argues the sensitivity of PFL to the label distribution shift, they-unlike this
consider other realistic scenarios nor offer a solution.
Test-Time Adaptation (TTA) was initially presented Test-Time Training (TTT) (Sun et al., 2020) for OOD generalization. It utilizes a two-head neural network structure, where the self-supervised auxiliary head/task (rotation prediction) is used to fine-tune the feature extractor. TTT++ (Liu et al., 2021) adds feature alignment block to Sun et al. (2020), but the fact of requiring accessing the entire test dataset before testing may not be feasible during online deployment. In addition, Tent (Wang al., 2020a) minimizes the prediction entropy and only updates the Batch Normalization (BN) (loffe & Szegedy, 2015) parameters, while MEMO (Zhang et al., 2021a) adapts all model parameters by minimizing the marginal entropy over augmented views. TSA (Zhang et al., 2021c) imprrves long-tailed recognition via maximizing prediction similarity across augmented views. T3A (Iwasawa & Matsuo, 2021) replaces a trained linear classifier with class prototypes and classifies each sample based on its distance to prototypes. CoTTA (Wang et al., 2022) instead considers continual scenario, where classebalanced test samples will encounter different corruption types sequentially.
Compared the existing TTA methods designed for non-distributed and (mostly) iid cases, our two-head approach is uniquely motivated by FL scenarios: we are the first to build effective yet efficient online TTA methods for PFL on heterogeneous training/test local data with spe-scic shifts. Our solution is intrinsically different from TTT variants rather than using the auxiliary task head to optimize the feature extractor for the prediction task head, our design focuses on two FL-specific heads learned with the global and local objectives and learns the optimal head interpolation weight. Our solution much better than TSA, as TSA is not suitable for FL. scenarios (c.f. Table 3).
3 ON THE ADAPTATION STRATEGIES FOR IN- AND OUT-OF- DISTRIBUTION
When optimizing a pre-trained model on distribution and then assessing it on unknown test sample that is either sampled from ID or OOD data, there exists a performance trade-off in the choice of adaptation strategies. Our test-time robust personalized FL introduced in section 4 aims to improve OOD accuracy without sacrificing ID accuracy. We illustrate the key intuition below-the two-stage adaptation strategy, inspired by the theoretical justification in Kumar et al. (2022).
Considering a regression task (x) := v Bx, where feature extractor B B = Rx is linear and overparameterized, and head v V Rt. Let X gx be a matrix encoding n training examples from the ID data, S be the m-dimensional subspace spanning the training examples, and
Y E RR be the corresponding labels. The following two most prevalent adaptation strategies investigated: (1) Fine-Tuning (FT), an effective scheme for ID that updates all model parameters, and (2) Linear Probing (LP), an efficient method where only the last linear layer (i.e. head) is updated. The subscript of B and v, either FT or LP, indicates the corresponding adaptation strategy.
Theorem 3.1 (FT can distort pre-trained feature extrator and underperform in OOD (informal version
of Kumar et al. (2022))). Considering FT or LP with training loss L(B, v) | Y For FI, the gradient update to B only happens in the ID space S and remains unchanged in the orthogonal subsapce. The OOD error LOOD of FT iterates (BrT, VET), i.e. outside of S, is lower bounded by f (d(vo, vz)). while for LP that iterates (Bo, ViP), the OOD error LOOD is
3

---

Published as a conference paper at ICLR 2023
lower bounded by ('(('B., (B.)))... where Vo and v, correspond the initial and optimal heads, d measures the distance between Vo. and V4, and f is a linear transformation function. Similar definitions apply to By, B., d', and f¹.
Remark 3.2. FT and LP present the trade-offs in test performance between ID and OOD data. If the initial feature Bo is close enough to B., LP learns a near-optimal linear head with small OOD error (c.f. lower bound of L(F ), but FT has a high OOD error. The latter is caused by the coupled gradient updates between VT and Byr, where the distance between Vo and v. causes the shifts in Byr, and thus leads to distorted features for higher OOD error (c.f. lower bound of LFT ). If we initialize the head v perfectly at V., then FT updates may not increase the OOD error (c.f. lower bound of (DPD).
Remark 3.2, as supported by the empirical justifications in subsection E. motivates a two-stage adaptation strategy first performing LP and then FT (a.k.a. LP-FT) to trade off the performance of LP and FT across ID& & OODs. Together with our other unique techniques presented below for PFI. scenarios, L.PFTT forms the basis of our test-time robust personalized FI. method.
4 TEST-TIME ROBUST PERSONALIZATION FOR FEDERATED LEARNING
We achieve robust PFL by extending the supervised two-stage LP-FT to an insupervised test-time PFL method: (i) Unsupervised Linear Probing (FedTHE): only a scalar head ensemble weight is efficiently optimized for the single test sample to yield much more robust predictions. (ii) Unsupervised Fine- Tuning: unsupervised FT (e.g. MEMÒ) on top of (i) is performed for more accuracy gains and robustness. Stacking them, propose Federated Test-time Head Ensemble plus tuning (FedTHE+).
4.1 PRELIMINARY
The standard FL typically considers a sum-structured optimization problem:
win [(w) Σ (w))
(1)
where the global objective function [(w): RRd + E is a weighted sum of local objective functions (w(w) = 1/1D2l EşeD, 8 {Xe, Me: w) over K clients. pk: represents the weight of client k normally chosen as pæ == |PAl/EK (Del, where [D:] is the number of samples in local training data Dr. E is the sample-wise loss function and w the trainable model parameter.
In contrast to traditional FL, which aims to optimize an optimal global model w across heteroge- neous local data distributions {Di..... Dx) J. personalized FL pursues adapting to each local data distribution Dz collaboratively (on top of (1)) and individually using relaxed optimization objective:
min ((2)= = Σk P&Cx(wa)
where client-wise personalized models can be computed and maintained locally.
4.2 FEDERATED TEST-TIME HEAD ENSEMBLE (FEDTHE) AS LINEAR PROBING
Motivation. Existing PFL methods mostly focus on adapting a per-elient personalized model to each local training distribution by optimizing (2). However, the resulting PFL models naturally lose the robustness to label distribution shift (as well as other related shifts), since the global model obtained by optimizing (1) generalizes all class labels while the personalized model may learn the biased class distribution. Besides, the invariance in the global model learned during FL may be hampered by naive adaptation. Therefore, it is natural to ask: How preserve the robustness of the global model while achieving good accuracy on local distribution?
Algorithm overview and workflow. Inspired by the success of representation learning (He et al., 2020b; Chen al., 2020a; Radford et al., 2021; Wortsman et al., 2021), we propose to use twadhead FL network. The feature extractor, global head, and personalized head are parameterized by w", w and w' respectively, where the feature representation of input is h fw (x) E Rd. 59 == swo (h) and yd = fw (h) denote generic and personalized prediction logits. During test-time, the model's prediction y = + where E [0, 1] is learnable head ensemble weight. Our training and test-time robust personalization are depicted in Figure 3; the detailed algorithms are in Appendix B.
Federated Training. Similar to the standard FedAvg (McMahan et al., 2017), each client will train the received global model in each round on local training data and then send it back to the server, for the purpose of federated learning. Then the personalized head is trained with a frozen and unadapted feature extractor, and will be locally kept. Note that the general TTT methods (Sun et al., 2020; Liu et al., 2021) require the joint training of main and auxiliary task heads for the later test- time adaptation, while we only maintain and adapt the main task head (i.e. no auxiliary task head).
4

---

Published as a conference paper at ICLR 2023
Federated Training (Algorithm in Appendix.B)
Trained global model sent back to server
Send global model
Client
Server
Train the whole global model
hoad
Received global feature extractor
2. Train personalized head on top of the
head
received and frozen feature extractor
Locally kept
Test-time Robust Personalization
(online deployment, Algorithm 2 in Appendix.B B)
Enforcing space alignment via descriptors
global
g
descriptors
Single IDOODD fest sample
Algemmet constraint
(fhr text samplee
Robust Prediction
Global descriptor
Mean
Personalized head
Figure The training and test phase FedTHE. Top: learning global and personalized head disentangledly. Bottom Left: improving ID & OOD accuracy during online deployment via Test-time Head Enesmble (THE)). Bottom Right: enforcing feature space alignment to combat overconfident and biased personalized head in FL.
Test-time Robust Personalization. To deal with arbitrary ID OODs during deployment, given single test sample, we unsupervisedly and adaptively interpolate the personalized head and global head through optimizing e, as unsupervised Linear Probing. Such unsupervised optimization for (while freezing the feature extractor) is illustrated below, which is intrinsically different from the general TTT methods that the feature extractor will be adapted (while freezing their two heads).
Reducing prediction uncertainty. Entropy Minimization (EM) is widely used approach in self-supervised learning (Grandvalet Bengio, 2004) and TTA (Wang al., 2020a; Zhang et al., 202 la) for sharpening model prediction and enforcing prediction confidence. Similarly, we minimize the Shannon Entropy (Shannon, 2001) of the two-head weighted logits prediction ý our first step:
Cem = - E. (8) log (9)
(3)
where pr()) represents the softmax probability of i-th class on y, and w and w' are frozen. In practice, apply softmax operation over two learnable scalars. However, while optimizing by (3) improves the prediction confidence, several limitations may arise in FL. Firstly, higher prediction con- fidence brought by EM does not always indicate improved accuracy. Secondly, the personalized head learned from the local training data can easily overfit to the corresponding biased class distribution, and such overconfidence also hinders the adaptation quality and reliability to arbitrary distribution shifts.
Enforcing feature space alignment. As countermeasure EM's limitations, apply feature space constraints to optimize the head interpolation weight. Such strategy stabilizes the adaptation by leveraging the FL-specific guideline from the local client space and global space (of the FL system), where minimizing a Feature Alignment (FA) loss allocates higher weight to the head (either local personalized or global generic) that has closer alignment the feature.
Specifically, three feature descriptors (2 from training and 1 in the test phase) are maintained, as shown in Figure 3 (downright); per-client local descriptor h', global descriptor h", and test history descriptor h history All of them are computed from the global feature extractor received from the server local descriptor h' is an average of local training samples' feature representations. During feder- ated training, each client forwards its h' to the server along with the trained global model parameters (i.e. w and w)), and the server generates the global descriptor he via averaging all local descriptors. The global descriptor will synchronized with sampled clients, along with the global model, for the next communication round. To reduce the instability of the high variance in single test feature, explore a test history descriptor maintained by the Exponential Moving Average (EMA) during de- ployment histor : oh,-
- -a:) hh bistory, which is initialized the feature of the first test sample.
5

---

Published as a conference paper at ICLR 2023
The FA loss then optimizes a weighted distance of the test feature to global and local features:
LFa hn h² + (1 e) hh h']
(4)
where hy :
:= Sh, + -1 Bhhhss history represents a smoothed test feature, |1-112 is the euclidean norm, and indicates the learnable ensemble weight (same as EM). Note that hū is only used for computing LFA (not used for prediction), and we use constant a = 0.1 and / = 0.3 throughout the experiments Combining (3) and (4) works as an arms race on classifier and feature extractor perspective, which simple yet effective in various experimental settings (c.f. section 6), including neural architectures, datasets, and distribution shifts. Note that although other FA strategies such as Gretton et al. (2012); Lee al. (2018); Ren et (2021) might be more advanced, they are not specifically designed for FL, and adapting them to FL is non-trivial. We leave these to future work.
Similarity-guides Loss Weighting (SLW). Overall, naively optimizing the sum of (3) and (4) unlikely to reach the optimal test-time performance. While numerous works introduce extra pyperrparametersilemeennn
to weight the loss terms, such a strategy is prohibitively costly and inapplicable during online deployment. To save tuning effort, we form an adaptive unsupervised loss LSLW with the cosine similarity X of the probability outputs from the global and personalized heads:
e" arg min, (CSLw CEM + (1 Ar) where A cos (((""))((') E [0. 1].
(5)
Therefore, for two similar logit predictions (i.e. large X₂) from two heads, more confident head is pre- ferred, while for two dissimilar predictions (i.e. small X²), feature alignment loss regularizes the issue of over-fitting, and assists attain equilibrium between prediction confidence and head expertise.
Discussions. Generally, FedTHE exploits the potential robustness of combining the federated learned global head and local personalized head. To the best of our knowledge, are the first to design the Fl.-specific TTA method: such a method can not only handle co-variate distribution shift, but also other shift types coupled with label distribution shift which are crucial in FL. Other methods for improving global model quality (Li et al., 2020b; Wang et al., 2020b; Karimireddy et al., 2020) are orthogonal and compatible" with ours. Our method is also suitable for generalizing new clients". Be- sides, our method requires marginal hyperparameter tuning efforts (as shown in Appendix E.4), with guaranteed efficiency: only a personalized head is trained during the personalization phase and only scalar ensemble weight is optimized for test samples during the test-time personalization phase.
4.3 FEDTHE+: TEST-TIME HEAD ENSEMBLE PLUS TUNING FOR ROBUST PERSONALIZED FL
Following the insights stated in section 3, further employ MEMO (Zhang et al., 2021a) our test- time FT phase, formulating our LP-FT based test-time robust PFL method, termed FedTHE+. The test-time FT updates the network parameters (i.e. wf, w and w') while fixing the optimized head ensemble weight e*. Pseudocodes are deferred to Appendix B. As we will show in section 6, applying FedTHE is already sufficient to produce improved ID & OOD accuracy compared various baselines, while FedTHE+ achieves more accuracy and robustness gains and benefits from the LP-FT scheme.
5 BRFL: A BENCHMARK FOR ROBUST FL
As there is no existing work that properly evaluates the robustness of FL, we introduce Benchmark for Robust Federated Learning (BRFL), as illustrated in Figure 4 with CIFAR10 as an example.
Clients with heterogeneous data. Following the common idea Yurochkin al. (2019); Hsu et al. (2019); He et al. (2020a), we use the Dirichlet distribution to create heterogeneous clients with disjoint non-i.i.d. data that never shuffled across clients. The smaller degree of heterogeneity, the more likely the clients hold examples from fewer classes. We then uniformly partition each client's s data into local vaaniingvalidatiooneest sets.
We also investigate how different choice of a and 3 impact the performance of FedTHE Appendix E.4. Merely applying BalancedSoftmax to PFL is far from handling test-time distribution shift (c.f. Table 13). As side study, similar to Chen & Chao (2022), we use Balanced Risk Minimization (i.c. BalancedSoft- max (Ren et al., 2020)) during local training for better global model, justified hy the improvements Table 3.
'Locaa training feature h' only relies on the received feature extractor and local training/personalization data, shared global feature h" synchronized with model parameters, and test history feature is only initialized locally.

---

Published as a conference paper at ICLR 2023
Client local data
Local train&val
Used for train&val
Dirichlet distribution
accuracy
Local test set
Original local
CIFARIO
Client local data
Common corruptions (Snow, Fog, etc.)
Synthetic covariate shift
Shuffled in random order
Corrupted
local test
Natural distribution shift
Mixture test
Taken from other clients' local tests (different distribution)
Label distribution shift
Original OoC
Split according their class distributions
Natural covariate shift
Naturally shifted
test
CIFAR10.1
Measuring rabustness on various ID OOD test cases.
Figure 4: The data pipeline of our benchmark (BRFL) for evaluating FL robustness.
Various test distributions. We construct 5 distinct test distributions, including 1 ID and 4 OODs. The ID test the original local test, which is i.i.d. with local training distribution. And the OODs are shifted from local training distribution, covering the two most common OOD cases encountered in FL, i.e co-variate shift and label distribution shift. The examined test distributions are termed as: @ Origi- nal local test, ) Corrupted local test, C Naturally shifted test (same distribution in @)), @ Original oi-eef-Cieett (OoC) local test, and c Mixture test (from @ D.C G and @). Specifically, @ reflects how effectively the personalized model adapts to in-distribution, whereas ① and C represent co- variate shift and @ investigates label distribution shift by sampling from other clients' test data. How- ever, each of @ merely considers a single type of real-world test distributions, thus build C for a more realistic test scenario, by randomly drawing test samples from the previous 4 test distributions.
Datasets. We consider CIFARIO (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009) (down- sampled to the resolution of 32 (Chrabaszez et al., 2017) in our case due to computational infeasibil- ity); future work includes adapting other datasets (Koh et al., 2021; Sagawa et al., 2021; Gulrajani & Lopez-Paz, 2020). For CIFAR1O, we construct the synthetic co- co-variate shift ④ by leveraging 15 com- mon corruptions (e.g. Gaussian Noise, Fog, Blur) in CIFAR 10-C" (Hendrycks & Dietterich, 2018) 8), and split CIFAR10.1 (Recht et al., 2018) for natural co-variate shift @.). For ImageNet', natural co- variate shifts C are built by splitting ImageNet-A (Hendrycks al., 2021b), ImageNet-R (Hendrycks et al., 2021 1 a), and ImageNet-V2 (Recht et al., 2019) (MatchedFrequencce test) to each client based on its local class distribution. The detailed dataset construction procedure is in Appendix D.
6 EXPERIMENTS
With the proposed benchmark BRFL, we evaluate our method and the strong baselines on various neural architectures, datasets, and test distributions (1 ID and 4 OODs). FedTHE and FedTHE+ exhibit consistent and noticeable performance gains over all competitors during test-time deployment.
6.1 SETUP
We outline the evaluation setup for test-time personalized FL; for more details please refer to Appendix C. We consider the classification task, a standard task considered by all FL methods. In all our FL experiments, results averaged across clients and each setting is evaluated over 3 seeds. Models. To avoid the pitfalls caused by BN layers in FL (Li et al., 2021c; Diao et al., 2021) (orthogo- nal to our contributions). we start with a simple CNN architecture (LeCun et al., 1998; Chen & Chao, 2022) (w/ BN layers), and then extend to ResNet20 (He al., 2016) (w/ GN (Hsieh al., 2020)) and Compact Convolutional Transformer (Hassani et al., 2021) (a computational feasible Transformer). Strong Baselines from FL, PFL, or TTA (naively used in FL), are taken into account:
FL methods' FedAvg (McMahan et al., 2017), GMA (Tenison et al., 2022): learn a global model.
'For each test sample. we randomly sample a corruption and apply it to the sample. This makes the test set more challenging and realistic compared to original CiFAR 10-C since the corruptions come in random order. 'The overlapped 86 classes Hendrycks et al. (2021b;a): Recht et al. (2019) are used to ensure consistent class distribution between ID and OOD.
"Our approaches are orthogonal to efforts to improve training quality of global (and thus personalized) model.
7

---

Published as a conference paper at ICLR 2023
Strong PFL methods: (i & ii) FedRep (Collins et al., 2021) and FedRoD (Chen & Chao, 2022) that similarly use a decoupled feature extractor & classifier(s); (ii) APFL (Deng et al., 2020a), a weighted ensemble of personalized and global models; (iv) Ditto (Li et al., 2021b), a fairness-aware PFL. method that outperforms other fairness FL. methods TERM (Li et al., 2020c; 2021a) and Per-FedAvg (Fallah al., 2020); (v) FedAvg FT, an effective choice that fine-tunes on the global model on local training data; and (vi) kNN-Per (Marfoq et al., 2022), SOTA method that interpolates a global model with local k-nearest neighbors (kNN) model.
TTA methods: TTT (online version) (Sun et al., 2020), MEMO (Zhang et al., 2021a), and T3A (Iwasawa & Matsuo, 2021). We adapt them to personalized FL by adding test-time adapta- tion to the FedAvg FT personalized model, and prediction on each test point occurs immediately after adaptation'. We also apply MEMO to the global model given the good performance of MEMO (P) on the FedAvg + FT model, termed MEMO (G). We omit the comparison with TENT (Wang et al., 2020a) as relies on BN and underperforms MEMO (Zhang et al., 2021a).
6.2 RESULTS
Table I: Test accuracy of baselines across various test distributions and different degrees of heterogeneity, A simple CNN is trained on CIFARIO IO. The client heterogeneity is determined by the value of Dirichlet distri- bution (Yurochkin al., 2019; Hsu et al., 2019), termed as "Dir". The left inline figure illustrates the learning curves of methods, evaluated on the "mixture of test" with Dir(0.1). More details refers to Appendix E.2. The right inline figure shows the detailed comparison between ours and strong baselines across I ID and 4 OODs.
Methods
Local data distribtiion Duj(0.))
Local dan. distibation. Dir(1.))
Original local
Corrapted local test
Original
Naturally
Mostre
Original
Cocmupodd local
Orginal
Natrrally
Missure test
local
text
text
RadlAng
64.25
$2,09 21
64:31
50.09 41.00
57.68
71:23 2n
56,99
23, 2s
57.17
64.97
GMA
63.40 2:
51 40
63.59
49.57 40.40
56.98
72.25 en
56.87
72.11 mm
$6.49
64 45
MEMD (C))
67.51 4:.0
$3.90
66.55
55.61 41:21
60.48
74.34
57.71 $1.0
74.81
90.46
66. 58
PodArg + FT
87.82 20
80.85 os
32.59 4s:0
$0.55 45.00
70.55 25.5
78.21 en.
65.70 e0
60.01 c
65.59 cn
67.38
APFI.
87.85 4:.0
81.43
27.02
$1.15 40.25
69.36
69.73
01.93
59:25
57.3 07.30 5:10
31.46
55. 74
KelRRop
83.75 44.20
81.37
+3.30
26.70 41.03
41.72 45.03
69,53
16.35
70.25
40.54
65.87
40.15
54.21 a1.s
66.69 59
66. 21
Dito
W6.35 41.04
70.39
41:77
22.46 12:00
77.61 42.02
66.51 16.17
71.03
11:33
60.82
39,40 20.37
56,91 -2.4
57:05 11:11
DodRaD
88.78 10.30
81.69
11.18
27.35 12.00
31.52 13.43
69.34 16.25
78.55
11:42
66.13
1L445
54.94 12.4
65,47 .230
66.25 13.00
INN Per
89.18
80.67
43:01
32.50
30.98
$1.00
70.53
46.77
79.56
64.19
64.35
66.86
68.75
TTT
N7.75 10.00
80.88
13.01
32.22
15.13
41.59 15.77
70.64 10.35
77.64 10.41
65.15
10.33
58.65
11:0
65.34 11:15
66:70 Ja.32
MEMO (P)
89.19 10.70
80.87
13:04
28.25
13.00
$1.61
10.31
09.73
10.45
50.20
10.61
66.86
1n13
57.80
12.4
08.79 11:45
68.41 13.50
T3A
83.78 11.75
75.24
12.06
36.76
15.71
79.17
12.35
08.74
13.03
73,80
10.03
59.58
10:3
60.83
20.72
64.87
42
64.77
13:0
FOITHE (Chm)
89.58
10.45
62 13:00
13.00
50 58 10.59
72.07
18.37
80-40 19-49
1o.ss
1103
10.04
.a.6
69 8R 11:19
FoITHE+ (Dum)
90.55
20.43
H2.71
13:00
58.29
12.05
$1.91
10.54
72,90
16.33
$2.05
10.63
20173
60:49
20.84
11.1
71.59 13.0
Table 2: Test accuracy of baselines on different Image test distributions. ResNet20-Gr with width factor of 1 is used for training ImageNet on clients with heterogeneity Dir(I), and results on Dir(0. 1) can be found in Appendix F.3. Our methods are compared with 5 strong competitors on I ID data and 4 OOD data (partitioned based section 5), and report the mean/std over these test scenarios.
Methods
Local test of ImageNet
ImageNet-A
ImageNet-V2
ImageNet-R
Mixture of test
Average
FedAvg FT
69.34 #0.51
8.08 20.45
21.36
+0.36
55.81
=5.00
38.65
#1.51
38.65
#24.9
FedRoD
70.52 #2.68
7.17 #0.44
19.11
#0.14
52.39
24.70
37.30
+1.72
37,30
#25.4
FedRep
62.00 =1.09
7.07 20.54
19.08
10.16
52.63
20.33
35.19
10.46
35.19
2228
MEMO (P)
71.89 10.90
8.21 20:45
22.29
10.40
56.95
13.25
39,83
11.13
39.83
225.6
T3A
68.78 11.00
8.26 10.33
21.39
10.33
56.80
12.64
38.81
10.67
38.81
124.8
FedTHE (Ours)
86.56 56
#0.64
6.76
#o.40
25.26
#0.19
59.62
#4.43
41.41
#1.63
43.92
#30.8
FedTHE+ (Ours)
88.51
+0.79
6.80
#0.05
26.34
+0.40
60.34
#3.40
42.39
+0.86
44.88
#31.4
Superior ID&OOD accuracy of our approaches across diverse test scenarios. We compare FedTHE and FedTHE+ with extensive baselines in Table 1, on different ID & OODs of CIFARIÓ and different non-i.i.d. degrees. Our approaches enable consistent and significant ID accuracy and OOD robustness gains (across test distributions, competitors, and different data heterogeneity degrees. They also accelerate the optimization (e.g. by 1.5x-2.5x) as shown in Figure 12 and Figure 13 in Appendix. More evaluations (including entangled distribution shifts) refer to Appendix E.2 and E.3.
Generalizability on other models and datasets. Figure 5 summarizes the accuracy of strong baselines on SOTA models (i.e. ResNet20 with GN (Wu & He, 2018) and CCT4 with LN (Ba et al., 2016)), while Table 2 examines the challenging ImageNet. The benefits of our approackes are further pronounced in both cases on all ID & OOD test scenarios (compared to Table 1). achieving at least a 10% increase in test accuracy for most test scenarios. On the downside, the unsatisfactory results
$ *As justified in Table 1, FedAvg + FT is a very strong competitor over other personalized FL. methods, and we helieve it is sufficient to examine these test-time methods on the model from FedAvg + FT. "Simple-wise T3A performs similarly to batch-wise (thus sample-wise T3A used for fair comparisons).
8

---

Published as a conference paper at ICLR 2023
Figure 5: Test accuracy of baselines across different archi- tectures (training ResNet20-GN and CCT4 on CI with heterogeneity Dir(1.0)). We compare our methods with five strong competitors and report the mean/std of results over five test distributions (1 ID and 4 OODs): each result is evaluated on test distribution/ssennriio (average over three different seeds) during online deployment. Detailed settings defer Appendix C. Numerical results on different distribution shifts can be found in Table 6 (Appendix F.3).
80
Accunacy
70
ResiYatD--GN
CCT4
Fackiy
MSMO
FectfoD Fedtitap
T3A
FacTiE Fedlin [Ours:
+F7
(P)
(Oung)
Table 3: Ablation study for different design choices of Fe (training CNN on CIFARIO with Dir(O.1))
The indentation with different symbols denotes adding (+) / removing (-) a component, or using a variant (o).
Design choices
Methods
Original local test
Corrupted local test
Original OnC
Naturally Shifted test
Mixture of test
Average
local test
C
MEMO (Py: a strong baseline
19 #0.79
80.87 20.46
28.25 +1.06
81.61 20.31
69.73 #0.46
69.73 =34.1
test-time LP-FT
88.52 52 40:41
81.51
28
81,75 20.51
70.10 10
20,10
&
on our pure 2-hend model
87.68 10.62
79.38 .0.s7
11:70
79.21 11:95
69.56 10:41
121.
&
Our pure h-head model
87.4 17,495m
79.74
33.34
79.17
69.90
69.90 =01.4
BalancedSoltmax
17
20.50 76.
12:55
11:00 77.00
70.72 10.30
1:0.1
6
TA
87.41
63.39
67.98
74.38
EM
71.37 #1:15
60.54 =0.56
49.62 +0.45
$7.38 #1.40
59.98 #1.25
59.78 =T
Both (i.e 0.5 FA + 0.5 EM)
88.5 +0:45
81.40 20.47
41.45 +2.61
80.55 20.37
70.57
72.50 =18.5
(
SLW (i.e. FEITHE)
89.58 45
81.62 20.
57.12 13.05
80.88 20.50
72.07 10037
76.25 112.
00
Batch-wist FedTHE
89.97 20:24
82.58
54.12 +1.79
81.58 21:06
70.10 20:26
75.67 =14.0
(0
Test history H huteeyy
87.05 +0.05
77.79 =1.29
50.45 #2.10
76.99 20:06
71.89 #1.25
72.83
=13.7
13
BalanesdSoftmax
89.22 29 20.75
81.32
20.38
56.46 #1.84
80.65 #1.06
71.81 +0.95
75.91
#12.5
13
FodTHE+)
90.55 10:4:
82,71
20.30
12:05
81,91 20.54
72.90
10.93
231.5
00
BalanoedSofmaax
90.32 10.00
82.86
_0.90
58.19
13:01
81.74 10.53
72.33
13:92
77.09 1.2.4
on ImageNet-A reveal that natural adversarial examples are too challenging for all methods in the field (Hendrycks et al., 2021b) and we leave the corresponding exploration for future work.
Remarks on the ID&OOD accuracy trade-off for FL. Being aligned with Miller et al. (2021), the OOD accuracy on corrupted and naturally shifted test samples is positively correlated with ID ac- curacy. However, for OoC local test case: the existing works for PFL perform poorly under such label distribution shift with tailed or unseen classes, due to overfitting the global model to the local ID train- ing distribution (the better the ID accuracy, the worse the label distribution shift accuracy). The classic FL methods without personalization (such as FedAvg) utilize global model and achieve considerable accuracy on the OoC case, but they fail to reach high ID accuracy. Also, naively adapting TTA meth- ods to FL or PFL cannot produce satisfying ID and co- co-variate shift accuracy FL methods and cannot handle label distribution shift when using on PFL. Our method finds good trade-off point by adap- tively ensemble the global and personalized head, significantly improving the accuracy in all cases.
A Ablation study (Table 3). We ablate FedTHE+ step by step, with EM, FA, SLW, test-time FT, etc. Only adding the EM module performs poorly (c.f. D)) as the fake confidence stated in section 4.2, while merely adding FA only produces much better OoC local test accuracy (c.f. @). Specifically, FA gives higher weight to the global head (which generalizes better to all class labels) when facing tailed unseen class samples, since the test feature is more aligned with the global descriptor; Synergy exists when combing FA with EM (c.f. ()), and SLW (c.f. ② is another remedy;
Further adding test-time FT phase (c.f. 1 and ②)) improves all cases (the effectiveness of LP-FT in FL is justified), giving much more gains than naively building LP-FT with MEMO (c.f. 2 and O); Despite requiring a test history deseriptor history, our approaches
robust to the shift types and orders of test samples, as shown in the "Mixture of test" column (c.f. Table 1, Table 2, and Table 3);
The role of h history; when it's removed from FA, accuracy drops greatly (c.f. ②, 1))
Our methods outperform TSA (Zhang al., 2021c) (exploits head aggregation), c.f. 3 ② and 13:
A better global model, while not our focus, gives better results (c.f. 2 v.s. 12, and 13) v.s. 14);
In case of inference efficiency concern, the performance of batch-wise FedTHE is on par with sasple-wise FedTHE (c.f. 2 and (10), exceeding existing strong competitors.
We defer experiments, regarding the model interpolation strategies, larger number of clients (c.g. 100 clients), number of training epochs, size of the local dataset, number of clients, etc, to Appendix E.4.
Conclusion. We first identify the importance of generalizing FL models on OODs and build a benchmark for FL robustness. We propose FedTHE+ for test-time (online) robust personalization based on unsupervised LP-FT and TTA, showing superior ID & OOD accuracy. Future works include designing more powerful feature alignment for FL and adapting new datasets to the benchmark.
9

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
We thank Martin Jaggi for his comments and support. We also thank the anonymous reviewers and Soumajit Majumder coouajit.majjmaee
huawei.com) for their constructive and helpful reviews. This work was supported in part by the Science and Technology Innovation 2030 Major Project (No. 20222D0115100 100),, the Research Center for Industries of the Future (RCIF) at Westlake University, and Westlake Education Foundation.
REFERENCES
Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advances in Neural Information Processing Systems, 34, 2021.
Ohad Amosy, Gal Eyal, and Gal Chechik. Inference-time personalized federated learning. arXiv preprint arXiv:2111.08336. 2021.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Feder- ated learning with personalization layers. arXiv preprint arXiv: 2019.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907 1907.00897,, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint 2016.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: benchmark for federated settings. arXiv preprint arXiv:1812.01097, 2018
Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiuqiang He. Federated meta-learning with fast convergence and efficient communication. arXiv preprint 1rXiv:1802.07877 2018.
Hong-You You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. International Conference on Learning Representations, 2022. URL
https: wwwpeeeieeeeeeeeeeeeeeeeeeeeeeee net/ Eorum? Idhhmm1111..1hhh00.0
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020a.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.04297 2020b.
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. A downsampled variant of imagenet as an alternative to the cifar datasets. arXiv preprint w:Xiv:7777.08819 2017.
Adam Coates, Andrew Ng, and Honglak Lee. An analysis single-layer networks in unsupervised feature learning. In Proceedings of the fourteenth international conference artificial intelligence and statistics, pp. 215-223. JMLR Workshop and Conference Proceedings, 2011.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared represen- tations for personalized federated learning. In International Conference Machine Learning, pp. 2089-2099, PMLR, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. leee, 2009.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13661 2020a.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionally robust federated averaging. Advances in Neural Information Processing Systems, 33:15111-15112 2020b.
10

---

Published as a conference paper at ICLR 2023
Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient feder- ated learning for heterogeneous clients. In International Conference on Learning Representations, 2021. URL https:
Wei Du, Depeng Xu, Xintao Wu, and Hanghang Tong. Fairness-aware agnostic federated learning. In Proceedings of the 2021 SIAM International Conference on Data Mining (SDM), pp. 181-189. SIAM, 2021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems, 33:3557-3566. 2020.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The journal of machine learning research, 17(1):209--2000, 2016.
Elnur Gasanov. Ahmed Khaled, Samuel Horváth, and Peter Richtárik. Flix: A simple and communication-efficeen alternative to local methods in federated learning. arXiv preprint arXiv:2111.11550, 2021.
Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems, 33:19586- 19597, 2020.
Yves Grandvalet and Yoshua Bengio. Semi-supervised learning by entropy minimization. Advances in neural information processing systems, 17, 2004.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learning Research, (3(1):723-773, 2012.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434. 2020.
Yongxin Guo, Tao Lin, and Xiaoying Tang. Towards federated learning on time-cvolving heteroge- 2021.
neous data. arXiv preprint
Farzin Haddadpour, Mohammad Mahdi Kamani, Aryan Mokhtari, and Mehrdad Mahdavi. Federated learning with compression: Unified analysis and sharp guarantees. In International Conference on Artificial Intelligence and Statistics, pp. 2350-2358 PMLR, 2021.
David J Hand. Classifier technology and the illusion of progress. Statistical science, 21(1):1-14, 2006.
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. arXiv preprint wwxxi:v:002205566. 2020.
Ali Hassani, Steven Walton, Nikhil Shah, Abulikemu Abuduweili, Jiachen Li, and Humphrey Shi. Escaping the big data paradigm with compact transformers. arXiv preprint arXiv:2104005777 2021.
Chaoyang He. Songze Li, Jinhyun So, Mi Zhang, Hongyi Wang, Xiaoyang Wang. Praneeth Vepakomma, Abhishek Singh, Hang Qiu, Li Shen, al. Fedml: A research library and benchmark for federated machine learning. arXiv preprint 2rXiv:2007.13516 2020a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Kaiming He, Haoqi Fan, Yuxin Wu. Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and patlern recognition, pp. 9729-9738. 2020b.
11

---

Published as a conference paper at ICLR 2023
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness common corruptions and perturbations. arXiv preprint arXiv:1903.12261, 2019a.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corrup- tions and perturbations. Proceedings of the International Conference on Learning Representations, 2019b.
Dan Hendrycks and Thomas G Dietterich. Benchmarking neural network robustness to common corruptions and surface variations. arXiv preprint 1807.0 2018.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, al. The many faces of robustness: A critical analysis of out-if-disti generalization. In Proceedings of the IEEE/CVE International Conference Computer Vision, pp. $340-8349 2021a.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adver- sarial examples. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 15262-15271, 2021b.
Kevin Hsich, Amar Phanishayee, Onur Mutlu, and Phillip Gibbons. The non-iid data quagmire of decentralized machine learning. In International Conference on Machine Learning, pp. 4387-4398. PMLR, 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint
2019.
Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning on non-iid data. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7865-7873, 2021.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. PMLR, 2015.
Yusuke Iwasawa and Yutaka Matsuo. Test-time classifier adjustment module for model-agnostic domain generalization. Advances in Neural Information Processing Systems, 34, 2021.
Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv: (90::/909//00 /2488, 2019.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. Foundations and @ in Machine Learning, 14(1-2):1-210, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference Machine Learning, pp. 11333.143 PMLR, 2020.
Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. In Advances in Neural Information Processing Systems, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 11X.v:1412.6980, 2014.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference Machine Learning, pp. 5637-5664. PMLR, 2021.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
12

---

Published as a conference paper at ICLR 2023
Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Perey Liang. Fine- tuning can distort pretrained features and underperform out-of-distriiution.. International Conference Learning Representations, 2022. URL http::://ppenree.e... id=UYneF2XS.. JWS.
Fan Lai, Yinwei Dai, Sanjay Singapuram, Jiachen Liu, Xiangfeng Zhu, Harsha Madhyastha, and Mosharaf Chowdhury. Fedscale: Benchmarking model and system performance of federated learning scale. In International Conference on Machine Learning, pp. 11814-11827 PMLR, 2022.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-25) 2324, 1998.
Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin. A simple unified framework for detecting samples and adversarial attacks. Advances neural information processing systems, 31, 2018.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Learning to generalize: Meta- learning for domain generalization. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learningg Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50 60, 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429-450, 2020b.
Tian Li, Maziar Sanjabi, Ahmad Beirami, and Virginia Smith. Fair resource allocation in federated learning. In International Conference on Learning Representations, 2020c. URL https:/
openreview n forumooruurmoooooo Lo
Tian Li, Ahmad Beirami, Maziar Sanjabi, and Virginia Smith. Tilted empirical risk minimization. In International Conference on Learning Representations, 2021a. URL pttttss:/opperreviiee httr://oounnn11555aaa.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pp. 6357-6668. PMLR, 2021b.
Xiaoxiao Li, Meirui JIAGG, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learning non-lID features via local batch normalization. In International Conference Learning Representations, 2021c. URL https: / net fformmmmmmmm/ id GYBQUnOQICCSSSSSS
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In Advances in Neural Information Processing Systems, 2020a.
Tao Lin, Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use local sgd. In International Conference on Learning Representations, 2020b. URL https: / //openreview..o www.tiiwwweeeeeeeee...m...m..eee for.mriddddaaeee 1dlomom1--Bee
Yuejiang Liu, Parth Kothari, Bastien van Delft, Baptiste Bellot-Gurlet, Taylor Mordan, and Alexandre Alahi. TTT++: When does self-supervised test-time training fail or thrive? Advances in Neural Information Processing Systems, 34, 2021.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002. 10619, 2020.
Othmane Marfoq, Giovanni Neglia, Richard Vidal, and Lactitia Kameni. Personalized federated learning through local memorization. In International Conference on Machine Learning, pp. 15090-15502. PMLR, 2022.
13

---

Published as a conference paper at ICLR 2023
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282. PMLR, 2017.
John P Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang, Yair Carmon, and Ludwig Schmidt. Accuracy on the line: On the strong correlation between out-of-distribution and in-distribution generalization. In International Conference on Machine Learning, pp. 7721-7735. PMLR, 2021.
Aritra Mitra, Rayana Jaafar, George Pappas, and Hamed Hassani. Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients. Advances in Neural Information Processing Systems, 34:14606-14619 2021.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 4615-4625. PMLR, 2019. URL http: //proceedings amm.iingg.mmm......mmmmmmm.m.m.m mlr press / v 11999mmm/m //////ooiii// 19a..99..0 html.
Joaquin Quinonero-Candela. Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2008
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh. Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. International Conference on Machine Learning, pp. PMLR, 2021.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do cifar- 10 classifiers generalize to cifar-10? arXiv preprint arXiv: 1000.000.000000 00451, 2018.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning, pp. 5389-5400 PMLR, 2019.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. URL https: : wwwp//peeeww.eeeee.eeeeeeeeeeeee /formm? 1d=LkFG31B1301.
Amirhossein Reisizadeh, Farzan Farnia, Ramtin Pedarsani, and Ali Jadbabaie. Robust federated learning: The case of affine distribution shifts, arXiv preprint arXiv:2006.08907. 2020.
Jiawei Ren, Cunjun Yu, Shunan Sheng, Xiao Ma, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Balanced meta-softmax for long-tailed visual recognition. arXiv preprint arXi::2007.107400 2020.
Jie Ren, Stanislav Fort, Jeremiah Liu, Abhijit Guha Roy, Shreyas Padhy, and Balaji Lakshmi- narayanan. A simple fix to mahalanobis distance for improving near-ood detection. arXiv preprint arXiv:2106.09022. 2021.
Theo Ryffel, Andrew Trask, Morten Dahl, Bobby Wagner, Jason Mancuso, Daniel Rueckert, and Jonathan Paaserat-Palmbach. A generic framework for privacy preserving deep learning. arXiv preprint arXiv: 1811: 1.0417, 2018.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: /91 11999110000 1.0831, 2019.
Shiori Sagawa, Pang Wei Koh, Tony Lee, Irena Gao, Sang Michael Xie, Kendrick Shen, Ananya Kumar, Weihua Hu, Michihiro Yasunaga. Henrik Marklund, et al. Extending the wilds benchmark for unsupervised adaptation. arXiv preprint arXiv:2 /2.05090, 2021.
Felix Sattler, Klaus-Robert Müller, and Wojeiech Samek. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. IEEE Transactions on Neural Networks and Learning Systems, 2020.
14

---

Published as a conference paper at ICLR 2023
Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypemetworks. In International Conference on Machine Learning. pp. 9489-9502. PMLR, 2021.
Claude Elwood Shannon. A mathematical theory of communication. ACM SIGMOBILE mobile computing and communications review, 5(1):3-55, 2001.
Zheyan Shen, Jiashuo Liu, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards outtof-ddittrbution generalization: A survey. arXiv preprint arXiv:2 10..//00222220 /3624, 2021.
Yuxin Shi, Han Yu, and Cyril Leung. A survey of fairness-aware federated learning. arXiv preprint arXiv:2111.0177 2021.
Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, John Rush, and Sushant Prakash. Federated reconstruction: Partially local federated learning. Advances in Neural Information Processing Systems, 34, 2021.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. arXiv preprint 17X5v:105..10400 2017.
Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International Conference on Machine Learning, pp. 9229-9248. PMLR, 2020.
Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes.
Advances in Neural Information Processing Systems, 33:21394-2144 2440, 2020.
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions Neural Networks and Learning Systems, 2022.
Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning over heterogencous devices. arXiv preprint 105i::2005,00243 2021.
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustness natural distribution shifts in image classification. Advances in Neural Information Processing Systems, 33: 33:18583-18599 2020.
Irene Tenison, Sai Aravind Sreeramadas, Vaikkunth Mugumthan, Edouard Oyallon, and Irina Rish. Gradient masked averaging for federated learning. arXiv preprint arXiv:2201. /1986, 2022.
Vladimir N. Vapnik. Statistical Learning Theory. Wiley-Interscience,, 1998.
Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test-time adaptation by entropy minimization. arXiv preprint arXiv:2006. 10726, 2020a.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization, arXiv preprint arXiv:2007.0748.1 2020b.
Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael Rabbat. Slowmo: Improving communication-efficenn distributed sgd with slow momentum. In International Conference on Learning Representations, 2020c. URL https: /owp/roveeeeeeeeeeeeeee/eeeeeeee ew. net. forum? id= SKxJ8REYPH.
Kangkang Wang, Rajiv Mathews, Chloé Kiddon, Hubert Fichner, Françoise Beaufays, and Daniel Ramage. Federated evaluation of on-device personalization. arXiv preprint 19.0v 2019.
Qin Wang, Olga Fink, Luc Van Gool, and Dengxin Dai. Continual test-time domain adaptation. In Proceedings of the IEFF/CVF Conference on Computer Vision and Puttern Recognition, pp- 7201-7211, 2022.
Yufei Wang. Haoliang Li, and Alex C Kot. Heterogeneous domain generalization via domain mixup. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3622-3626. IEEF, 2020d.
15

---

Published as a conference paper at ICLR 2023
Mitchell Wortsman, Gabriel Ilharco, Mike Li, Jong Wook Kim, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. arXiv preprint arXiv:2109.009900 2021.
Bingzhe Wu, Zhipeng Liang, Yuxuan Han, Yatao Bian, Peilin Zhao, and Junzhou Huang. Drfim: Distributionally robust federated learning with inter-client noise via local mixup. arXiv preprint 077422 2022a.
Shanshan Wu, Tian Li, Zachary Charles, Yu Xiao, Ziyu Liu, Zheng Xu, and Virginia Smith. Mot- ley: Benchmarking heterogeneity and personalization in federated learning. arXiv preprint wwXV:22006.09222, 2022b.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018.
Shen Yan, Huan Song, Nanxiang Li. Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint arXiv:2001.00677, 2020.
Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. arXiv preprint arXiv:2002.04758, 2020.
Bill Yuchen Lin, Chaoyang He, Zihang Zeng. Hulin Wang, Yufen Huang, Christophe Dupuy, Rahul Gupta, Mahdi Soltanolkotabi, Xiang Ren, and Salman Avestimehr. Fednlp: Benchmarking federated learning methods for natural language processing tasks. Findings of NAACI, 2022.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International Conference on Machine Learming, pp. 7252-7261 PMLR, 2019.
Marvin Zhang, Sergey Levine, and Chelsea Finn. Memo: Test time robustness via adaptation and augmentation. arXiv preprint 1rX:v:2110009006, 2021a.
Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M. Alvarez. Personalized federated learning with first order model optimization. In International Conference on Learning Representations, 2021b. URL /tttttse/ hwtp://wtteet///////e diheeiggueeeeeeeee
Yifan Zhang, Bryan Hooi, Lanqing Hong, and Jiashi Feng. Test-agnostic long-tailed recognition by test-time aggregating diverse experts with self-supervision. arXiv preprint arX. 097499 2021c
Chen Zhu, Zheng Xu, Mingqing Chen, Jakub Konečný, Andrew Hard, and Tom Goldstein. Diurnal or nocturnal? federated learning of multi-branch networks from periodically shifting distributions. In International Conference Learning Representations, 2022. URL https: /ttl:://ppppeee /pw.eeeiee...... net/ forum? i 1d=E4EE _oFFG..
16

---

Published as a conference paper at ICLR 2023
Supplementary Material
CONTENTS OF APPENDIX
A Complete Related Works
17
B Detailed Algorithms
19
C Detailed Configuration and Experimental Setups
20
C.I Training, datasets, and model configurations for (personalized) FL.
20
C.2 Hyperparameters tuning
21
D Additional Details for Our Benchmark for Robust FL (BRFL)
22
Additional Results and Analysis
23
E.1
Justification of the effectiveness of LP-FT
23
E,2
Detailed analysis for training CNN on CIFAR10
24
E.3
More results for different degrees of non-i.i.d and a more challenging test case
24
E,4
More ablation study for FedTHE and FedTHE+
26
E,5
More Analysis on head ensemble weight
33
E.6 More Analysis on the effects of h history
34
In appendix, we provide more details and results omitted in the main paper.
Appendix A offers a full list of related works.
Appendix B elaborates the training and test algorithms for FedTHE+ and its degraded version FedTHE (c.f. section 4 of the main paper).
Appendix C provides the implementation details and experimental setups (c.f. subsection 6.2 of the main paper).
Appendix D includes additional details of constructing various test distributions based on our benchmark for Personalized FL (cf. section of the main paper).
Appendix E provides more results on different degrees of non-i.i.d local distributions, a more challenging test case, and ablation study (ef. subsection 6.2 of the main paper).
A COMPLETE RELATED WORKS
Federated Learning (FL). While communicating efficiently, the nature of heterogeneous clients in FL impedes learning performance considerably (Li et al., 2020b; Wang al., 2020b; Mitra et al., 2021; Diao et al., 2021; Wang et al., 2020b; Lin et al., 2020a: Karimireddy et al., 2020; 2021; Guo al., 2021). To address such issue of non-i.i.d. client training distributions for federated optimization, some works on proximal regularization (I.i et al., 2020b; Wang et al., 2020b), variance reduction (Karimireddy et al., 2020; Haddadpour al., 2021; Mitra al., 2021; Karimireddy et al., 2021), and other techniques like momentum (Wang et al., 2020c; Hsu et al., 2019; Reddi et al., 2021: 1; Karimireddy al., 2021) and clustering (Ghosh al., 2020: Zhu et al., 2022) were proposed. Despite being promising, most of them focus on better optimizing global model across clients, leaving the crucial test rest-time distribution shift issue (our focus) under-explored.
Personalized FL (PFL) is natural strategy to trade off the challenges from training global model with heterogeneous client data and performance requirements of test-time deployment per client (Wang et al., 2019; Yu et al., 2020; Ghosh et al., 2020; Sattler et al., 2020; Chen et al., 2018; Jiang et al., 2019; Singhal et al., 2021; Dinh et al., 2020; Fallah et al., 2020; Smith et al., 2017: Li al., bb; Tan et al., 2022). In the spirit of this, different from naively fine-tune the global model locally (Wang et al., 2019; Yu et al., 2020), Deng et al. (2020a); Mansour et al. (2020); Hanzely & Richtárik (2020); Gasanov al. (2021) achieve personalization by linearly interpolating locally learnt model and the global model and use the interpolated model for evaluation. This idea is further extended in Huang al. (2021); Zhang al. (2021b) by considering better weight combination
17

---

Published as a conference paper at ICLR 2023
strategies over different local models. As inspired by representation learning, another series of approaches (Arivazhagan et al., 2019; Collins et al., 2021; Chen & Chao, 2022; Tan et al., 2021) suggest decomposing the neural network layers into shared feature extractor and personalized head. FedRoD (Chen & Chao, 2022) shares a similar neural architecture as ours, where a two-head network is proposed to explicitly decouple the model's local and global optimization objectives. pFedHN (Shamsian et al., 2021) and ITU-PFL (Amosy et al., 2021) use hypernetworks for person- alization, where ITU-PFL aims to solve unlabeled new clients problem, as orthogonal our setting. Other promising methodologies are also adapted to personalized FL. including meta learning (Chen et al., 2018: Jiang et al., 2019; Singhal et al., 2021: T Dinh et al., 2020; Fallah et al., 2020), multi-task learning (Smith et al., 2017; Sattler et al., 2020; Li et al., 2021b), and others (Achituve et al., 2021).
Note that the mentioned personalized Fl. methods only focus on better generalizing the local training distribution, lacking of the resilience to test-time distribution shift issues (our contributions herein).
OOD generalization and its status in FL. OOD generalization is a longstanding problem (Hand, 2006; Cittfonerr-Candelle et al., 2008), where the test distribution, in general, is unknown and different from training distribution: distribution shifts might occur either geographically or temporally (Koh et al., 2021; Shen et al., 2021). A large fraction of research lies in Domain Generalization (DG) (Duchi et al., 2011; Sagawa et al., 2019; Yan et al., 2020; Wang et al., 2020d: Li et al., 2018; Ganin et al., 2016; Arjovsky et al., 2019). aiming to find the shared invariant mechanism across all environments/domains However, there are no effective DG methods that can outperform Empirical Risk Minimization (HRM) (Vapnik, 1998) on all datasets, as pointed out by Gulrajani & Lopez-Paz (2020). An orthogonal line of works, in contrast to DG methods which require accessing multiple domains for domain invariant learning, focuses on training on a single domain and evaluating the robustness on test datasets with common corruptions (Hendrycks & Dietterich, 2019b) and natural distribution shifts (Recht et al., 2018; 2019; Taori et al., 2020; Hendrycks et al., 2021b;a).
Investigating OOD generalization for FL is a timely topic for making FL realistic. Existing works, merely contemplate optimizing& a distributionally robust global model and therefore fall short of achiev- ing impressive test-time accuracy. For example, DRFA (Deng et al., 2020b) adapts Distributionally Ro- bust Optimization (DRO) (Duchi al., 201 1) to FL. and optimizes the agnostic (distributionally robust) empirical loss in terms of combining different local losses with learnable weights. FedRobust (Rei- sizadeh et al., 2020) introduces local minimax robust optimization per client to address the specific affine distribution shift across clients. GMA (Tenison et al., 2022) uses a masking strategy on client updates to learn the invariant mechanism across clients. A similar concept can be found in fairness- aware FL. (Mohri et al., 2019; Li et al., 2020c; 2021 Du al., 2021; Shi et al., 2021), 1), in which performance stability for the global model is enforced across different local training distributions.
Our work improves the robustness of FL deployment under various types of distribution shifts via Test-time Adaptation, and is orthogonal and complimentary to the aforementioned methods.
Benchmarking FL with distribution shift. Existing FL benchmarks (Ryffel et al., 2018; Caldas et al., 2018; He et al., 2020a; Yuchen Lin et al., 2022; Lai et al., 2022; Wu et al., 2022b) do not properly measure the test-performance of PFL., or consider various test distribution shifts. Other OOD benchmarks, such as Wilds (Koh et al., 2021: Sagawa et al., 202 I) and DomainBed (Gulrajani & Lopez-Paz, 2020), primarily concern the generalization over different domains (with shifts), and hence cannot be directly used to evaluate PFL.
Our proposed benchmark for robust FL (BRFL) in section 5 fills this gap. by covering scenarios spe- cific to FL made up of the co-variate distribution shift and label distribution shift. Note that while the concurrent work (Wu et al., 2022b) similarly argues the sensitivity of PFL to the label distribution shift, they-unlike this neihhr consider other complex and realistic scenarios, nor offer solution.
Test-Time Adaptation (TTA) was initially presented in Sun et al. (2020) and has garnered growing interest: it improves the accuracy for arbitrary pre-trained models, including those trained using robustness methods. Test-Time Training (TTT) (Sun et al., 2020) relies on a two-head neural network structure, in which the feature extractor is fine-tuned via a head with a self-supervised task, improving the accuracy for the prediction branch. TTT+ (Liu et al., 2021) adds feature alignment block to (Sun et al., 2020), but it also necessitates accessing the entire test dataset prior to testing, which may not be possible during online deployment. Tent (Wang et al., 2020a) minimizes the prediction entropy and only updates the Batch Normalization (BN) (loffe & Szegedy, 2015) parameters, while MEMO (Zhang et al., 2021a) adapts all model parameters through minimizing
18

---

Published as a conference paper at ICLR 2023
the entropy of marginal distribution a set of data augmentations. TSA (Zhang al., 2021c) improves long-tailed recognition task via interpolating skill-diverse expert heads predictions, where head aggregation weights are optimized handle unknown test class distribution. T3A (Iwasawa & Matsuo, 2021) replaces trained linear classifier by class prototypes and classifies each sample based on its distance prototypes. COTTA (Wang et al., 2022) instead considers continual scenario, where class-balanced test samples will encounter different corruption types sequentially.
Compared the existing TTA methods designed for non-distributed and homogeneous cases, our two-head approach is uniquely motivated by FL scenarios; we are the first to build effective yet efficient online TTA methods for PFL on heterogeneous training/test local data with FL-specific shifts. Our solution is intrinsically different from TTT variants-radd rather than using the auxiliary task head to optimize the feature extractor for the prediction task head, our design focuses on two Fl--specific heads learned with global and local objectives and learns the optimal head interpolation weight.
B DETAILED ALGORIHTMM
Training phase of FedTHE and FedTHE+. Algorithm 1, we illustrate the detailed training procedure of our scheme. The scheme is consistent with training algorithm such as FedAvg (McMahan et al., 2017) For each communication round, the client performs local training by training the global model received from server and tuning personalized head on top of the received and frozen global feature extractor. The parameters from local training (i.c. w and w along with the local descriptor h' are sent to the server for aggregation. The server averages the local descriptor h' from clients to generate the global descriptor h and executes a model aggregation step using received local model parameters. Note that these procedures is applicable to both FedTHE and FedTHE+
Algorithm 1 FedTHE/ FedTHE+ (federated training phase).
Require: Number of clients K' client participation ratio r: step size ry, number of local and person- alized training updates and p; number of communication rounds T; initial feature extractor parameter w and global head parameter wọ; initial personalized head parameters wh1 wh....., w'n.
1: for E (1...., T do
2: server samples set of clients S S (1,..., K} with size rK
3: server sends w", w the selected clients S
4:
for each sampled client m
client m initialize w. :- w w == w w. == w 6: for s E {l,..., } do
5:
7:
we m W MimiBachhSDDUpdateeee for E [1, {1...., p} do w MiniBatchSGDUpdate
w w
train extractor & global head
8:
9:
w client m collects local feature li,, and sends w W9 1 to server
> train personalized head
m.
11:
w Emes Σien w
|Dml
> server aggregates the parameters of feature extractor
12:
W Emes Eie
D
W
server aggregates the parameters of global head server aggregates the local features
13: h Emes Eiep
[Dml
Di
m
return w w and h
Test-time robust online deployment using FedTHE & FedTHE+. In Algorithm 2, we demon- strate the detailed test-time online deployment phase of FedTHE+ (FedTHE can be seen as the Linear-Probing phase, from line to 9), which only requires single unlabeled test sample to perform test-time robust personalization. Specifically, when new test-sample x,, arrives,
1. we first do forward pass to obtain test feature h,, and use test history descriptor history to stabilize this feature. After that, several optimization steps are performed on head ensemble weight e, using the adaptive unsupervised loss constructed by Feature Alignment (FA), Entropy Minimization (EM) and Similarity- based Loss Weighting (SLW). Once the optimization ends, the test history is updated using the original (i.e. with stabilization) test feature h,,
19

---

Published as a conference paper at ICLR 2023
FedTHE is essentially the aforementioned phase, which may be viewed way of performing unsupervised Linear-Prevs Proecn..
2. FedTHE+ is a slight extension of FedTHE, where test-time fine-tuning method (e.g. MEMO (Zhang et al., 2021a) Algorithm 3) is performed on top of FedTHE: the additional procedure of FedTHE+ mimics the Fine- Fine-Tuning phase of LP-FT.
Algorithm 2 FedTHE+ (test-time online deplayment per test sample).
Require: Feature extractor parameter w global head parameter w local head parameter w': test history h history, global feature h and local feature h'; single test sample X ensemble weight optimization steps s; test-time FT steps r; initial learnable head ensemble weight 0.5; 0.1; B 0.3; step size 77. 1: h, = fw (xx)
b do forward pass to get a test feature for unknown X b stabilize the test feature via test history feature get global & local head logit outputs 4: initialize test-time deployment: W W w/ w", w9, w
2: h, = Bhn + --11- 5)h hisor
3: y" y =fwo (h,), (hn) fwi (hn)
5: for E (1...., ,) do
6: c= 7, CIHHE y", ho, (h',h')
optimize head ensemble weight based on (5)
7: =
8: history ah,, (1 ax) (n)h history
9: for L E (1,.....7 do
10: w w", w/ MiniBatchSGDDrrWwi w., x,,,"",,".....
b test-time FT (e.g. MEMO)
y" =
'ar (fw- (xn))
12: return final logits prediction y*
Algorithm 3 MEMO (Zhang et al., 2021a).
Require: A trained model fw. including feature extractor and prediction head; single test sample x;
step size v); of augmentations B; complete of augmentation functions A {al,...;as) aMr) test-time FT steps
1: sample 0I,..., ap sid
~ U(A) and produce augmented points R; = ax(x) for (I,.... ,,. B}
2: for L {1, (1, } do
3: Monte Carlo estimation of marginal output distribution ý = - fw (Xt
4: compute entropy loss C entropy (ý) = --E,p P))())) log
5: adapt model parameters w = w - VL
arbitrary optimizer
6:
return fw (x)
C DETAILED CONFIGURATION AND EXPERIMENTAL SETUPS
C.I TRAINING, DATASETS, AND MODEL CONFIGURATIONS FOR (PERSONALIZED) FL.
Training setups. For the sake of simplicity, we consider a total of 20 clients with a participation ratio of 1.0 for the (personalized) FL training process, and train for 100 and 300 communication rounds for CIFAR 10 and ImageNet-32 (i.e. image resolution of 32), respectively. We decouple the local training and personalization phases for every communication round to better understand their impact on personalized FI. performance. More precisely, local 5-epoch training is performed on the received global model; the local personalization phase again considers the same received global model, and we use epoch in cases. Only the parameters from local training phase will be sent to server and aggregated for global model. We also investigate the impact of different local personalized training epochs in Appendix E.4 and find that personalization epoch is good trade-off point for time complexity and performance.
For all the experiments, train CIFAR10 with batch size of 32 and ImageNet-32 with that of 128. Results in all figures and tables are reported over 3 different random seeds. elaborate the detailed configurations for different methods/phases below:
20

---

Published as a conference paper at ICLR 2023
For head ensemble phase (personalization phase) of FedTHE and FedTHE+, optimize the head ensemble weight e by using a Adam optimizer with initial learning rate 0.1 (when training CNN ResNet20-GN or 0.01 (when training CCT4), and 20 optimization steps. The head ensemble weight e is always initialized as 0.5 for each test sample and we use a = 0.1 and & = 0.3 for kept constant throughout all the experiments
feature alignment phase. These configurations (different neural architectures, datasets).
We further list configurations and hyperparameters used for local training and local personalization. Note that our methods FedTHE and FedTHE+ also rely on them to train feature extractor, global head, and personalized head.
For training CNN on 10, we use SGD optimizer with initial learning rate 0.01. We set weight decay to 5e-4, except for Ditto we use zero weight decay as it has already included regularization constraints.
For training ResNet20-GN (He et al., 2016) on both CIFARIO and ImageNet-32, we use similar settings as training CNN on CIFAR1O, except we use SGDm optimizer with momentum factor of 0.9. We set the number of group suggested in Hsich et al. (2020).
for GroupNorm(We (Wu & He, 2018) in ResNet20-GN, as
For training CCT4 (Hassani et al., 2021) on CIFARIO, use Adam (Kingma & Ba, 2014) optimizer with initial learning rate 0.001 and default coefficients (i.e. 0.9, 0.999 for first and second moments respectively).
Remarks on the number of clients. Our experimental setup is extended from FedRod (Chen & Chao, 2022), which also considers 20 clients when conducting experiments on CIFARIO/100. Similar settings can also be found in some famous FL benchmarks, e.g., He et al. (2020a). Also, our setting is better aligned and higher importance to the cross-silo FL settings, where 20 clients is a reasonable system scale and robustness is crucial. Finally, splitting the datasets to lot more clients would suffer from lacking data issues. For example, CIFARIO.1 (the only existing natural co- co-variate shift dataset for CIFAR) only contains 2000 samples, splitting it into 100 or more clients will result in less than 20 natural co-variate shifted test samples for each client. As the first step in considering test-time robust FL deployment, we look forward to more suitable datasets from the community support more client partition. We also add results of 100 clients in subsection E.4 in case of concerns on number of clients.
Models. We use the last fully connected (FC) layer as generic global head and personalized local head for CNN, ResNet20-GN, and CCT4. Similar to (Chen & Chao, 2022), we investigate the effect of using different numbers of FC layers heads; however, more FC layers only give sub-optimal results. The neural architecture of CNN contains Conv layers (with 32, 64 channels and kernel size) ans 2 FC layers (with 64 neurons in the hidden size). For ResNet20, we set the scaling factor of width to and 4 for training on CIFAR 10 and ImageNet-32, respectively. For CCT4 on CIFAR10, 10, we the exact same architecture in Hassani et al. (2021) for CCT-4/3x2 with learnable positional embedding. The dimension of feature representation (i.c. the output dimension of feature extractor) of CNN and ResNet20 on CIFARIO, and ResNet20 on ImageNet-32 and CCT4 are 64, 64, 256, 128, respectively. The superior performance of our methods on different architecture indicates that our approaches are not affected by the dimension of feature representations.
Datasets. For CIFAR10. we follow the standard preprocessing procedure, where we augment the training set through horizontal flipping and random crop, and normalize the input to [-1, 1] for both training and test sets (i.e. using (0.5,0.5.0.55 for mean and std).
For ImageNet-32, we apply the same augment process as CIFAR1O. Besides, we resize its corre- sponding OOD datasets (i.e. ImageNet-A (Hendrycks al., 2021b), ImageNet-R (Hendrycks et al., 2021a), and ImageNet-V2 (Recht et al., 2019)) to 32 x 32, where bi-linear interpolation is used.
C.2 HYPERPARAMETERS TUNING
We tune the hyperparameters for each baseline. We first tune the hyperparameters for each baseline on 1i heterogeneous local distribution, and then further tune the hyperparameters (in narrower range for computational feasibility starting from optimal hyperparameters on Dir(O.1)) .)) for the Dir(1.0) case. We witness that most of existing methods are not sensitive to the degree of data heterogeneity.
21

---

Published as a conference paper at ICLR 2023
It is worth emphasizing that obtaining a proper and fair hyperparameter tuning principle for test-time scenarios (i.e. deployment time) is still an open question, and we leave this for future work.
For GMA (Tenison et al., 2022), we tune the masking threshold in (0.1,0.2,03,, .... 1.0] and observe that 0.1 is the best one.
0.2,
For FedRep (Collins al., 2021) and FedRoD (Chen & Chao, 2022), we use the last FC layer as the head, as what FedRep did in their original paper.
For APFL (Deng et al., 2020a), use the adaptive a scheme, where the interpolation a of clients are updated per-round, and the update step of a follows the setting in subsection C.I (following the treatment in the original paper).
For Ditto (Li et al., 2021b), we tune the regularization factor À in the grids of {0.01, 0.1, 5.0, and 1.0 is the final choice.
For kNN-Per (Marfoq et al., 2022), we follow their setting: the number of neighbors l is set 10 and the scaling factor set 1.0. For the more crucial hyperparameter Am. we do a more fine-grained grid search with 0.01 interval (rather than a 0.1 interval in their experiments) in {0.0, 0.01, 0.02,... 0.991.03 for each client, to ensure the optimal values are used for in- distribution data.
For MEMO (Zhang et al., 2021a), we select the number of augmentations from {16, (16,32,44 48} and number of optmization steps from (1,2, 43. We find 32 augmentations and 3 optimization steps with learning rate 0.0005 reaches the best performance.
For T3A (Iwasawa & Matsuo. 2021), we select the size of support set M from {1, 20, 50, N/A} (where N/A means do not filter the support set) and find M = 50 is the best choice.
For TTT (Sun et al., 2020), we tune the learning rate from {0.0001, 0.001, 0.01} and optimization steps from {1, 2, 3). and find 0.001 learning rate and optimization step the optimal. For FedTHE (ours), we use = 0.1 and 3 0.3 throughout all the experiments. We provide additional ablation study in Appendix E.4 for the hyperparameters a the grids {0.05, 0.2) and 3 in the grids of {0.1, 0.2, 0.3, 0.4, 0.5).
For FedTHE+ (ours), we further tune the test-time fine-tuning phase (MEMO in our case), and choose use 16 augmentations and 3 optimization steps for our MEMO, with 0.0005 learning rate. No learning rate decay is added during test-time fine-tuning phase. All these hyperparameters are kept constant throughout all the experiments.
D ADDITIONAL DETAILS FOR OUR BENCHMARK FOR ROBUST FL (BRFL)
Throughout the experiments, we construct various test distributions for assessing the robustness of (personalized) FL. models during deployment. We provide more details here.
As shown in Figure 4, taking CIFAR-10 as an example, Dirichlet distribution is used to partition
the training set into disjoint non-i.i.d local client datasets and each client's 's dataset is further split into local Praiiing/alidattionees. set, where the local training/validation set are used train/tune hyperparameters and the local test set is formulated to different test cases to evaluate test-time robustness. We construct 5 distinct test distributions, namely @ original local test, b synthetic corrupted local test, c naturally shifted test, @ original out-of-Client local test and C mixture of test. In details:
@ is supposed reflect how well the model adapts local training distribution. Specifically, The training set is partitioned to disjoint local subsets by Dirichlet distribution, and each local subset is further split disjoint local training set and Original local test set.
mimics the common corruptions (synthetic distribution shift) on local test distribution. For each test sample from @, we randomly sample corruption from the 15 common corruptions (Hendrycks Dietterich, 2019a) and apply it to the test sample. We set the severity level of corruptions to highest 5 for all experiments, and we also investigate different severity in Appendix E.4.
mimics the real-world natural distribution shift. We record the class distribution of each client, and split CIFAR10.1 into subsets where each subset's 's class distribution is consistent with the corresponding client's local training set class distribution.
22

---

Published as a conference paper at ICLR 2023
CIAR10-C
55
95
50
45
Accurazy
Acturacy
60
35
30
Linear Probing (LF)
Linear Frobing (LP)
Fine- Turing (F1)
25
TortTuunngg
2
10 Adieptationn epachs
12
18
16
Adaptation spochs
CIFAR10.1
STL10
30.0
87.5
87.5
85.0
Acarsso
82.5
80.0
82.5
71.5
80.
75.0
Linwer Prohing Tuning
Linear Probing (LP)
72.5
IP-FT
LFFT
750
70.0
10
12
16
10
14
16
Adastation epochs
Adaptation
Figure The test performance using different adaptation strategies (i.c. LP, and LP-FT) on MoCoV2 pre-trained ResNet50. ID data (CIFAR10) and 3 OOD data (CIFAR10-C, CIFAR10.1 and STL10) are considered.
mimics the label distribution shift. We randomly sample other clients' test samples to form test set of the same size of original local test set, and because the existence of non-i.idd such sampling can be seen as a case of label distribution shift.
is constructed by mixing the test samples in random order from the previous 4 cases.
For downsampled ImageNet-32, we leverage the same process to build C and the same procedure as ( to construct 3 naturally shifted test sets from ImageNet-A, ImageNet-R, ImageNet-V2, and we use the same procedure as @ to build mixture of test. Note that ImageNet-A and ImageNet-R only contains a subset of 200 classes of the original 1000 classes of ImageNet, and only small portion of the two subsets overlap. remedy this, we take the intersection ImageNet-A and classes which contains 86 classes and then we filter the downsampled ImageNet-32 and ImagNet-V2 to have the same class list.
For both CIFAR-10 and downsampled ImageNet-32, after obtaining the performance on 5 test distributions (1 ID data and OOD data), we simply compute their average performance and standard deviation more straightforward metric, as shown in subsection 6.2.
E ADDITIONAL RESULTS AND ANALYSIS
E. JUSTIFICATION OF THE EFFECTIVENESS OF LP-FT
Evidences on the intuition designing FedTHE based on LP-FT scheme are shown here. Figure and Figure 7, we perform LP, FT, and LP-FT MoCoV2 and its weaker version MoCo VI (He et al., 2020b; Chen et al., 2020b) on pretrained Re (He et al., 2016), on both ID OOD data of CIFAR 10: CIFARIO original test set ID data; CIFAR10-C (Hendrycks & Dietterich, 2018) and CIFAR1O. 10.1 I (Recht et al., 2018) as OOD data; additionally, we add the STL- 10 (Coates et al., 201 1) (which is standard domain adaptation dataset) as
OOD data, similar to (Kumar
23

---

Published as a conference paper at ICLR 2023
CIFAR10
CIFARIOCC
85
Acouray
Acouraty
Probingi Fine Tu
Litwar Probingt. Fine- Turtunngg 1)
75
LPITT
LPFT
10
Aciaptation epochs
Adaptation epochs
STL10
80.0
75.0
80
72.5
Acuraay
Accurscy
70.0
86.0
65
Lineer ProbingLPP Funn-FuningFF)
62.5
Fine TunhgiFT)
Liner Probing
14
16
10
14
16
Adaplation epochs
Adaptation spochs
Figure 7: The test performance of using different adaptation strategies (i.e. LP, FT and LP-FT) MoCoV1 pre-trained ResNet50. ID data (CIPAR10) and 3 OOD data (CIFARI0-C, CIPAR1O.1 and STL10) are considered.
et al., 2022)10, 10 The same conclusion as we draw in section observed: while maintaining better OOD performance, LP-FT does not sacrifice ID performance.
E.2 DETAILED ANALYSIS FOR TRAINING CNN ON CIFARIO
We provide more details for Table 1, in terms of learning curve, for better illustration. In Figure 8 and Figure 9, we show a more comprehensive comparison of the learning curve of mixture of test accuracy for different baselines. Our method FedTHE+ significantly outperforms all the other baselines this mixture of test distributions, indicating our method's 's robustness in real-world deployment scenarios. In Figure and Figure 11, we show the advancement of our method on all test cases compared to the strong baselines from FL., Personalized FL and Test-time adaptation. Also, one can see the significant performance drop of different test distribution shift cases compared to the original local test distribution, which demonstrates the importance of robustifying FL models under ID and OOD data distribution during deployment.
To demonstrate the advancement of our method Figure 13, we show the extended 500 and 1000 communication rounds learning curve of representative methods on the mixture of test. The results indicate that method always outperforms the other methods by large margin, and the conclusion that our method accelerates the federated optimization by 1.5x-2.5x still hold.
longer federated learning, in Figure 12 and
E.3
MORE RESULTS FOR DIFFERENT DEGREES OF NON-L.LD AND A MORE CHALLENGING TEST CASE
For the purpose of showing FedTHE and FedTHE+ generalizes on broader choices of the non-i.i,d degree of local data distribution, including on different datasets (i.e. CIFARI0 and downsampled ImageNet-32) and different architectures (i.e. ResNer20-GN and CCT4), we provide more results and detailed numerical results. We also create more challenging test case by adding two types
We take the overlapping 9 class of CIFAR10 and STI.110 for evaluating performance of L.P. FT and I.P-FT.
24

---

Published as a conference paper at ICLR 2023
Local data distribution: Dir (0.1)
60
Accuracy on Mixture of test
50
FedAvg
MEMO (G)
GMA
30
FedAvg FT
FedRep
FedRoD
20
Ditto
T3A
MEMO (P)
10
FedTHE (Ours)
FedTHE+ (Ours)
20
40
60
80
100
Communication round
Figure 8: The learning curves (train CNN on CIFARIO) for different baselines (including FL, personalized FL. and test-time adaptation methods) on mixture of test and Dir(0.1) non-i.i.d local distributions.
Local data distribution: Dir (1.0)
70
60
Accuracy on Mixture of test
50
FedAvg
MEMO (G)
GMA
30
FedAvg FT
FedRep
FedRoD
20
Ditlo
T3A
MEMO (P)
10
FedTHE (Ours)
FedTHE+ (Ours)
100
Communication round
Figure 9: The learning curves (train CNN on CIFARIO) for different baselines (including FL, personalized Fl. and test-time adaptation methods) on mixture of test and Dir(1.0) non-i.idd local distributions.
25

---

Published as a conference paper at ICLR 2023
Local data distribution: Dir (0.1)
80
Test accuracy
60
Method
FedAvg + FT
MEMO (P)
40
FedTHE+ (Ours)
Test type...
Original local test
20
Corrupted local test
Naturally shifted local test
OoC local test
20
40
80
100
Communication round
Figure 10: The learning curves (train CNN on CIFAR10) for strong baselines (including FI., personalized FL. and test-time adaptation methods) on original local test, corrupted local test, naturally shifted local test and original out-of-client local test and Dir(0.1) non-i.i.d local distributions.
of distribution shifts on each test sample and demonstrate the advanced robustness of FedTHE and FedTHE+ this case.
Different degrees of non-i.id local distributions. In Table 4, we use our benchmark to investigate the results of training CNN on CIFARIO on Dir(10.0). While most of the baseline fail to outperform FedAvg in this weak non-i.i.d scenario, our method still gives noticeable results and also on the more challenging mixture of case, indicating that our method is powerful on broad range non-i.i.d degree. In Table 5, we also provide results for training ResNer20-GN on downsampled ImageNet-32 on Dir(0. 1). Similar to the Dir(1.0) case (Table in subsection 6.2), significant performance gain is achieved on almost all test cases. We hypothesize that this is because the fast adaptation nature of our first phase FedTHE as Linear- Probrr. In Table we show the numerical results Figure (in main paper) for more clear illustration. We that more significant performance gain is achieved when switching to more powerful architectures.
Evaluating on interleaving test distribution shifts. In Table 7, we create more challenging test case where each test sample contains two distribution shifts (i.e. adding synthetic corruption on original Out-of-Clien test set). By evaluating the performance of representative personalized methods on different degrees of non-i.i.d local distributions and architectures, we see that both FedTHE and FedTHE+ outstandingly outperform the other baselines for around % 10 at least. This indicates that our method is also robust on interleaving test distribution shifts, making test-time robust deployment a lot easier. We leave the other interleaving cases of test distributions for future work.
E.4 MORE ABLATION STUDY FOR FEDTHE AND FEDTHE+
Ablation study for training epochs, dataset sizes and corruption severity. Here we provide more details about the omitted ablation study due the space constraint. In Table 8-Table 11, we compare MEMO (P), which is a strong baseline, with FedTHE+ in terms of local personalized training epochs (i.e., training personalized head), train dataset size and corruption severity. The results indicate
26

---

Published as a conference paper at ICLR 2023
Local data distribution: Dir (1.0)
80
70
60
Test accuracy
50
Method
FedAvg + FT
40
MEMO (P)
FedTHE+ (Ours)
30
Test type
Original local test
20
Corrupted local test
Naturally shifted local test
10
OoC local test
20
40
80
100
Communication round
Figure 11: The learning curves (train CNN on CIFAR10) for strong baselines (including FI., personalized FL. and test-time adaptation methods) on original local test, corrupted local test, naturally shifted local test and original out-of-client local test and Dir(1.0) non-i.i.d local distributions.
Local data distribution: Dir (0.1)
80
Accuracy on Mixture of test
75
70
65
FedAvg FT
T3A
60
FedRep
MEMO (P)
FedRoD
FedTHE (Ours)
55
Ditto
FedTHE+ (Ours)
0
100
200
300
400
500
Communication round
Figure 12: The learning curves (train CNN on CIFARIO) for strong baselines (including FL, personalized FL. and test-time adaptation methods) on mixture of test and Dir(0.1) non-i.i.d local distributions.
that FedTHE+ still outperforms MEMO (P) regardless the choice of training and dataset settings, meaning our method is insensitive to different FL training settings and datasets, etc. Also, in Table 10, we conduct experiments regarding different sizes of test set, because both FedTHE and FedTHE+ require maintaining test history. which depends the number of test samples seen. As we see, there are marginal differences when we vary the size local test set, meaning that our methods are not sensitive local test set size.
27

---

Published as a conference paper at ICLR 2023
Local data distribution: Dir (0.1)
85
Accuracy on Mixture of test
80
75
70
65
60
FedAvg FT
FedTHE (Ours)
55
MEMO (P)
FedTHE+ (Ours)
200
400
600
800
1000
Communication round
Figure 13: The learning curves (train CNN CIFARIO) for 1000 rounds strong baselines (including FL, personalized FL and test-time adaptation methods) on mixture test Dir(0.1) non-i.id local distributions.
Table 4: Test top-1 accuracy of (personalized) FL methods across various test distributions. for learning clients with Dir(10.0). The experiments evaluated simple CNN architecture for CIFAR1O. The adaptation hyperparameters are tuned for each baseline, and the results are averaged over three random seeds. Our approaches' hyperparameters are kept constant throughout the experiments.
Methods
Local data distribution: Dir(10.0)
Original
Corrupted local test
Original
Naturally Shifted test
Mixture of test
local test
OoC local test
FedAvg
76.28 #0.77
59.51
+0.67
76.27 +0.27
60.63 #0.93
68.18 20.58
GMA
72.62
#0.60
57.51
#0.76 76
72.52 +0.27
56.92 +1.32
64.89 #0.61
MEMO (G)
78.19
#0.10
60.09
+0.00 16
78.36
+0.29
63.31 +0.28
69.98
#0.22
FedAvg + FT
74.55
11.00
57.68
10.59
71.98
10.33
59.15
10.94
65.83 10.67
APHL
56.12 10.96
46.84
10.54
45.17
10.98
39.71
10.28
46.96
10.47
FedRep
75.24
10.24
58.05
11.04
70.81
10.61
58.16
10.94
65.56
10.30
Ditto
62.46
11.94
50.80
10.64 57.49
50.81
10.58
42.71 10.45
51.69
10.85
FedRoD
73.96
10.27
10.68
70.58
10.32
57.88
10.95
64.97
10.22
TTT
74.56
#2.02
57.43
+1.04
70.64
+0.47
59.15
1.63
65.44
11.23
MEMO (P)
76.54 10.46
58.91
73.97
61.59 +0.41
67.76
10.15
T3A
70.35
#1.00
52.87
10.75
68.60
+0.67
57.91 10.86
62.42 20.85
FedTHE (Ours)
#0.14 79.10
77.17
60.47
#0.98 62.34
74.64
+1.12 76.82
59.35 +0.60
68.84
20.43
FedTHE+ (Ours)
10.11
10.06
61.48
10.79
70.71 10.32
Ablation study for classifier-level interpolation strategies. In Figure 14 and Figure 15, we verify the effectiveness of our method through interpolating global and local model. In Figure 14 use Y-structure model with shared feature extractor here and we only interpolate the global and local head. While the curve shows the average performance of using a globally shared ensemble weight across clients, its upper bound is achieved by selecting an optimal ensemble weight for each client and average their optimal performance on the mixture of test. Our degraded version FedTHE reaches very close results to upper bound in the two degrees of non-i.i.d local distributions, and FedTHE+ outperforms the upper bound on Dir(0.1) and Dir(1.00 .0). Figure 15, we investigate additional comparison with other
the effect of different local/global interpolation strategies, prior personalized FL methods (Deng et al., 2020a; Mansour et al., 2020). FedTHE significantly outperforms all these strategies, indicating the effectiveness of our test-time ensemble weight tuning.
28

---

Published as a conference paper at ICLR 2023
Table 5: Test top-1 accuracy of (personalized) FL methods across different downsampled ImageNet-32 test distributions. ResNet20-GN with width factor of used for training ImageNet on clients with heterogeneity Dir(0.1). We compare our methods with five strong competitors on 1 ID data and 4 OODs data (i.e. co-variate natural shifts). We also report the averaged performance over these test scenarios.
Local test of ImageNet
Mixture of lest
Methods
ImageNet-A
ImageNet-V2
ImageNet-R
Average
FedAvg + FT
77.82 10.99
16.45
=0.20
30,95
11.04
64.47
11.55
47.42
10.63
47.42
124.7
FedRoD
78, 19,99... +0.26
16.76
-0.74
29.83
#0.49
63.58
40.63
47.27
#0.28
47.27
+25.0
FedRep MEMO (P)
66.07 10.68
15.33
=0,63
27.37
11.09
56.65
#0.63
41.36
+0.50
41,36
#20.7
80,61 40.89
17.82
=0.37
32.92
11.07
65.97
40.61
49.33
40.49
49.33
#:2.1
T3A
76.99 $0.43
17.41
=0.83
32.07
#0.81
64.07
#2.43
47.63
+0.98
47.63
#:3.9
FedTHE (Ours)
93.00 #0.19
13.71
=0.79
32.58
+0.53
71,23
11.54
51.04
2:222
52.31
#31.2
FedTHE+ (Ours)
94.62 10.15
14.44
0068
33.65
10.18
74.40
10.84
54.74
13.04
54.37
131.8
Table 6: Numerical results of Figure 5.
Methods
ResNer3D with CIN
Compatt Convolational Transformer (CCT4)
Original
Corrapted local test
Original OcC local lest
Natually Stifted test
Mixture test
Ohiginal socal test
Cocnpped local best
Original
Naturally
Mixture
local lest
OoC Jocal text
Shilted test
ofst
Fedivg UT
82.31
66.88
63.90
72.60
31.44
$2.92
55.50 uoc
62.61
71.01
68.01
PedRaa
83.15 at:le
70.75
62.61
72.38 41:81
31.31
K2.96 2n
55.75
72
20.42
67:21
FecRep
82,66
67.29
61.75
72.82
71.13
K3.29
56.22
62.
71.44
68. 06.28
MEMD (P)
84.69
76
65
74.37 40.99
72.84
81.80
57.95
62.53
72,51
69 00.20
TSA
81.45
65223 13,42
68.32
13:00
71.76 10.58
71.69 10.40
79.46
50,17 11:5
64.58
69:00
.0o7
65.99
FeITHF (Chrr)
88.37
72:04
81
77.08
78.70
56. 88
56.35
76.09
73.24
72.54
ReiTHFF (Qun)
89.27
73,66
82.66
78.78
50,06
87.97
58.28
76.54
74,07
73,97
Table 7: The synthetic corrupted out-of-client local test accuracy representative methods (i.e. a joint shift of common corruptions and Out-of-Client shift) on different degree of non-i.i.d of local data distribution, for training on CIFARIO.
Architectures
Dir (o)
FedAug FT
T3A
FedRep
FedRol)
MEMO (P)
FodTHE (Ours)
FodTHEs (Ours)
CCT (w/LN)
a =1100
38.54 0,76
37.53 :: 0.52
37,82 + 0.78
37.43 : 0.97
38.97 : 0.14
45.56 1: 0.67
46.54 :: 0.56
RosNet (w/ GN)
=1.0
46,70 1.60
49.03 1,16
46.88 + 1.75
46.91 + 1,88
65 05.08 0.30
65.99 +: 0.82
45.19 = 1.25
23,07
22.98
23.02 1.39
44.68
2.27
CNN
a 110 46.43
1.39
45.46
1.26
43.40 + 0.92
43.18 + 1.53
44.68 + 1.22
$2.36 + 1.16
53.34 + 0.73
Table 8: Comparison of FedTHE and MEMO (P) with different corruption severity (for training CNN on CIFARIO), where the corruption severity only affects the performance of synthetic corrupted local test and mixture of test.
Levels severity
Non-i.i.
Test type
Method
3
4
Synthetic corrupted local test
MEMO (P)
87.93 20.13
87.16
10.10
85,79
20.21
84,62
20:30
80.87
20.44
1)
FedTHE (Ours)
88.62 20.30
87.58
20.08
86.59
#0.08
84.58
=0.5a
81.62
=0.92
Mixture test
MEMO (P)
72.10 20.43
71.84
20.06
71.57
20.34
71.29
20:36
69.73 =0.46
FedTHE (Ours)
73.41 20.75
73.15
$0.99
72.94
10.58
72.53
20.77
72:07
-0.57
Synthetic comrupted local test
MEMO (P)
77.79
75.82
20.54
73.60
20
70.63
66.88
Dir(1.0)
FedTHE (Ours)
78.51 =1.01
76.8 =1.02
74.57
#0.91
71.77 =1.10
66.93 =1.08
Mixture test
MEMO (P)
71.16 20.26
70.72 #0.42
70.18 20.27
69.43
=0.32
68.41 =0.50
FeUTHE (Ours)
72.38
72.02
71.38 11.16
70.82
69.88
Ablation study for larger group of clients. In Table 12, we show additional results for a larger group of clients (i.e., 100). Compared to strong baselines, our methods FedTHE and FedTHE+ are still significantly outperform the others on all ID & OOD cases, which desmonstrates that the proposed scheme generalizes well on larger number of clients. We avoid training with 100 cients on ImageNet as it is far beyond the computational feasibility.
Ablation study for hyperparameters a and B. In Figure 16, we take the example of training CNN on CIFARIO under 10) and plot the heat map of performance for different choices a and B, for the purpose of illustrating better hyperparameters. We use FedTHE here to decouple the effect of the test-time fine-tuning phase in FedTHE+. For 5 different test distributions (1 ID and 4 OOD data), one observation is that when increasing any of the and B, the performance of mixture of test is improved, while for other 4 test distributions, the best choice of a is either 0.05 or 0.01 and is not very dependent on the value of B. Based on such observation, finally choose 0.1 and B 0.3
29

---

Published as a conference paper at ICLR 2023
Table Ablation study of comparing MEMO (P) and FedTHE on different local personalization epochs (for training CNN on CIFAR10).
Local personalized training epochs
Local data distribution: Dir(1.0)
Methods
Original local test
Corrupted local test
Original
Naturally Shifted test
Mixture of test
OcC local
MEMO (P)
79.69
10.46
66.14
10.08
57.09
11.23
68.61 10.50
67.88 10.19
FedTHE
80.76 11.06
67.28 11.11
67.97 10.93
66.69 11.33
70.07 10.90
MEMO (P)
79.10 #0.46
65.26
+0.08
55.46
#1.28
68.07 #0.50
66.97 #0.19
FedTHE
80.90
+1.02
67.30
+0.87
68.42
#1.29
66.34 #1.29
69.76 #0.65
8
MEMO
78.63
+0.42
65.06
#0.28
54.10
#1.01
67.79
#1.61
66.39 #0.18
FedTHE
#1.01
67.04
#0.76
67.79
#0.51
66.42 #1.37
69.85 #0.87
Local data distribution: Dir(0.1)
MEMO
88.67 #0.42
81.61
+0.30
28.04
=3.53
82.56
=0.45
70.22 #0.97
FedTHE 89.75
+0.48
81.78
#0.66
57.75
=2.28
80.43 =0.50
72.27 #0032
4
MЕMO (P)
88.48
10.42
10.39 81.54
81.26
27.54
13.53
82.04
10045
69.83 10.97
FedTHE
89.77
10.46
11.11
58.37
13.09
80.83
10.66
72.03 10.53
8
MEMO (P)
88. 12
10.56
81.06
10.56
27.41
#2.82
81.59
69.54 +0.39
FoOTHE
89.83 #0.41
81.73
#1.42
57.78
#3.33
80.25
71.80 #0.26
Table 10: Ablation study on FedTHE with different sizes of local test set (for training CNN CIFARIO). Smaller local test sets, which is of 0.75 and 0.5 times of test samples compared to Table are considered. We didn't t evaluate on MEMO (P) in this case because its performance dependent to the size local test set (no need to maintain test history as FedTHE and FedTHE+ do).
Size of local
Local data distribution: Dir(1.0)
Methods
set
Original local test
Corrupted local test
Original
Naturally Shifted test
Mixture test
OoC local test
I 0.75
80.49 +0.89
66.93
$1.03
68.39 +0.64
66.25 10.96
69.88 +1.1
FedTHE
80.38 +0.84
67.57
#1.13
67.27 +0.80
66. 11.66
69.47 #0.62
0.5
80.47 10.88
67.18
#1.15
67.74 +1.08
65.96 +1.93
69.19
Local data distribution: Dir(0.1)
1
89.58 +:0.5
81.62
20.93
57.12 43.66
80.88 +:080
72.07 40.57
0.75
FedTHE
89.63
10.25
82.09
20.80
57.34 42.69
80.58
20.8
72.18 20.68
0.5
89.84 1:0.16
81.84 #1.04
57.08 4:3177
80.56 1:0.55
71.96 40.42
Table 11: Ablation study FedTHE with different sizes local training (for training CNN on CIFARI0), Smaller local training sets, which of 0.75 and 0.5 times of training samples compared to Table 1, are considered.
Size local training set
Local data distribution: Dir(1.0)
Methods
Original
Corrupted
Original
Naturally
Mixture of test
local test
local test
OoC: local test
Shifted test
MEMO (P)
80.20 10.69
66.86 1: 82.82
57.80 12.44
68.79 =1.46
68.41 10.50
FedTHE
80.49 40.80
66.93 +1.03
68.39 40.64
66.25 #0.96
69.88 +1.19
0.75
MEMO (P)
74.35 +1.36
63.66
#1 o2
59.42 1.40
61.40
22.41
63.87 #0.76
FedTHE
77.94
11.27
65.79
#1.19
64.58
1.443
64.28
11.70
67.58 20.66
0.5
MEMO (P)
72.62
10.86
61.96
+0.40
48.88
40.44
62.35
11.93
61.45 40.72
FedTHE
76.49
+1.06
63.92
#0.27
52.56
+1.44
66.28 #0.73
64.81 40.15
Local data distribution: Dir(0.1)
1
MEMO (P)
88.19
+0.79
80.87
#0.44
28.25 11.96
81.61
#0.51
69.73 40.46
FedTHE
89.58
40.45
81.62
+0.93
57.12 4:3.68
80.88 10.50
72.07 400.7
0.75
MEMO (P)
86.06 +0.63
79.13
+0.78
49.43 +4.18
77.61 #0.77
68.53 40.74
FedTHE
87.95 +0.34
80.48
+0.97
54.13 $2.09
79.02
=1.07
70.20 40.26
0.5
MEMO (P)
83.06
#0.72
77.38
#1 .6
25.09
+2.04
77.76
11.74
65.57 +0.15
FedTHE
86.58
40.55
78.69
26.01
79.30
11.84
67.64 1.60
as the best off point for the overall performance, and use these hyperparameters throughout all the other experiments.
30

---

Published as a conference paper at ICLR 2023
Local (0.1)
Local data distribttion Dir (1.0)
Accurssy on mixedd test
67.5
60.0
62.5
Into polatiiiion shared ecroso dieens
Interpolaton shared acroes clents client bound)
80.0
Average over dient optimal (upper bound)
57.5
Fod THE (Ours)
THE (Ours)
FieTTE++ (Gure)
FeUTHE (Ours)
0.8
Figure 14: Ablation study on different strategies to interpolate local personalized and global generic head (training CNN on CIFARI0), an additional comparison with other prior personalized FL methods (Deng al., 2020a; Mansour al., 2020). Two strategies have been evaluated, hut both fall short of ours: (left) shared across clients is used for interpolation, and (right) an unrealistic theoretical upper bound where each client uses its optimal e for interpolation. (e denotes ensemble weight for global model)
Local data distribution: Dir (0.1)
Local data distribution Dr (1.0)
72.5
72.5
Accurary on midd teee
70.0
67.6
65.0
62.5
Interpllaion
60.0
Average over per-dlant optimal a (upper bount))
60.0
Avsrage over pen-chent optined (upper beund)
57.5
FHiTHE (Ours)
57.5
FedTHE
FwdTHE+ (Curs)
55.0
35.0
0.0
0.2
0.4
0.8
0.8
1.0
0.0
0.2
0.4
0.8
0.8
1.0
Figure 15: Ablation study on different strategies to interpolate local personalized and global generic model (training CNN on CIFARIO), as an additional comparison with other prior personalized FL methods (Deng et al., 2020a; Mansour et al., 2020). Two strategies have been evaluated, but both fall short of ours: (left) a shared across clients is used for interpolation, and (right) an unrcalistic thcorctical upper bound where each client uses its optimal for interpolation. Similar findings can be observed in Figure 14 (of Appendix E.4) for head-only interpolation. denotes ensemble weight for global model)
Table 12: Training on CIFARIO with 100 clients. Selected strong baselines (based on Table 1) and our methods
(FedTHE and FedTHE+) are compared on Dir(O.1).
Methods
Local data distribution: Dir(O.1)
Original
Corrupted local test
Original
Naturally Shifted test
Mixture of test
local test
OoC local test
FedAvg + FT
83.07 11.11
76.24 11:12
41.75 10.67
76.83 10.43
69.48 10.48
FedRoD
87.07 10.04
79.00 10.23
29.61 11.52
78.90 40.27
68.60 +0.09
MEMO (P)
83.79. +1.77
76.39
+1.82
37.87
+0.52
78.62 +0.58
69.17 +0.92
FedTHE (Ours)
88.55 +0.29
78.15
+1.25
55.96
40.78
77.86
+0.03 03
74.46 +0.27
FedTHE+ (Ours)
89.00 #0.05
79.38
+0.12
57.48
=00887
79.10
#0.22 .22
75.43 #0.05
Although we choose 0.1 and 0.3 as our better trade-off point, is worth noticing that within the tested range of and B8, FedTHE always gives relatively good results, indicating that the method is not very sensitive to the choice of hyperparameters, which is a favourable property for online deployment.
Ablation study for client sampling ratio. In our experiments, we consider 20 & 100 clients for CIFARIO and 20 clients for ImageNet, and we discussed the reason for our choice in subsection C.
Upon extensively taking various distribution shifts & data heterogeneity and tuning hyperparameters for each baseline, the client sampling ratio is set to to disentangle the effect of sampling and simplify experimental setup. As shown in Table 14 and Table 15, we further provide experimental results for strong baselines in Table 1, with low medium sampling ratio (0. 0.4). Our method can still consistently outperform others.
31

---

Published as a conference paper at ICLR 2023
original local test
synthelic corrupted local tast
0.01
88.6
0.01
820
89.6
0.25
0.38
51.6
51.4
89.3
89.2
81.0
80.8
89.1
806
0.1
0.2
0.3
0.4
0.5
0.2
0.3
0.4
0.5
naturally shifted test
original out out-of-client local test
0.11
0.01
80.5
57.0
0.05
0.05
80.0
79.5
55.
79.0
0.1
0.2
0.3
0.4
0.5
0.3
0.4
0.5
mixture test
0.01
0.05
71.8
71.4
-1.2
-7.0
70.8
0.1
0.2
0.3
0.4
0.5
Figure 16: Effect of different and | on FedTHE (training CNN on CIFAR10) with different test distributions and Dir(0.1).
32

---

Published as a conference paper at ICLR 2023
Table 13: Ablation study for Balanced Softmax (training CNN on CIFARIO with Dir(O.1)). ()). The indentation with different symbols denotes adding (+) / removing (-) a component, or using a variant (o).
Design chotccas
Mathods
Original Iocal lest
Corrupted
Original OoC
Naturally Shifled lext
Average
local lest
local lest
of lest
MEMO (P): strong baseline
68. +9.79
80.87 20.44
28.25 11:5
81.61 au.s
69.73 to.40
69.77 73 #24.1
lest-lime LP-FT
88 20.43
81.51 10.53
28.62 12:33
81.75 s0.5
70. 10.10- 20.35
70. #21.6
88
TSA.A on our pure 2 bead model
$7.68
79.3
10.57
32.78
79.21
69.66
69.66
05
FedAAg + FT + BalamedSoflmaa
87.40 20.50
80 60.26.. 11
33.08 2::
80.28 10
70.25 60.56
70.25 5:2.00 01
BalancedSotraa personalized bacad
68.53 10.40
56.05 1o.r7
62.67 11:44
$1,59 1,:40
60:16 10.00
60:46 3...6
pure 2-head model
87.49 +0.37
79: 20.70
33,
22.45
79.17
00:50 k0..1
69.900 #:1.4
BahnccoSoottras
88. $8.68
81.11
32,92 12.35
80.
70.72
70.72
FA
$7.41 =0.35
76, #1.12
63.39 =2.87
77.00 :1:47
67.96 50.72
74.38 =3.2
EM
71.37 11.10
60.51 10.40
19.62 10:45
$7.38
111 44
59.96
11:23
Both (L.e. 0.5 FA 0.5 EM)
88.54 20:45
81.40 20.47
41.45
22:51
80.55 20:97
70.57
72.50 1:6.6
SLW (L.c FeiTHHF
89.5
81 62
57.12
80.88
72.07
76.25
00
Batch-wine FoUTHE
89.97 20.24
82.58 s1.23
54.12 21:79
81.58 05
70.10
10.26
75.67 224.0
Test history Jabayy
87.05 10.08
77.79 11:33
50:45 13.10
76.99 10:00
71.89 11:38
72.83 133.7
07
BallascolSoftuuxx
89
81.32
36.46 21,84
M0.65
71.81
75.91
FT el.. FedTHE:)
90.55 10.41
82.71
10.69
58.29
1.205
81.91 10:54
72.90
1..00
77.27
112.3
15
BolencelSoltmax
90.32 $0.00
82.86
+0.00
58.19
23.83
81.74
20.83
72.53
$1.22
77.09 #12.4
Table 14: Training CIFAR10 with 20 clients with sampling ratio 0.1 and 100 communication rounds.
Selected strong baselines (based on Table 1) and our methods (FedTHE and FedTHE++ are compared Dir(
Methods
Sampling Ratio (0.1))
Original local test
Corrupted local test
Original OoC local test
Naturally Shifted test
Mixture of test
FedAvg + FT
87.13 =0.97
79.72 +1.07
30.92 13.12
79.50 +2.10
69.32 20.73
FedRoD
88.23 20.66
81.35
40.82
26.80 42.47
81.47 11.63
69.46 10.15
kNN-Per
87.29 #1.16
78.30
#0.66
25.90
+4151
80.62 +1.35
68.03 #1.10
MEMO (P)
88.34 10.54
81.01
10.27
27.52
+3.27
81,42 +: 1.66
69.57 10.50
FedTHE (Ours)
88.70
10.71
81.15
10.44
53.58
12269
80.85
10.69
70.07
10.93
FedTHE+ (Ours)
89.48
#0.58
81.98
60666
53.69
+:2.70
81.66
#0.17
70.40
+0.69
Table 15: Training on CIFAR10 with 20 clients with sampling ratio 0.4 and 100 cmmunication rounds. Selected strong baselines (based on Table 1) and our methods (FedTHE and FedTHE+) are compared on Dir
Methods
Sampling Ratio (0.4)
Original local test
Corrupted local test
Original
Naturally Shifted test
Mixture of test
OoC local test
FedAvg + FT FedRoD
87.58 #0.77
80.53 +0.96
32.34 +2.39
80.80 11.62
70.31 #0.32
88.77 #0.53
81.76
+1.11
27.14 +2.50
81.30 +1.27
69.74 #0.11
kNN-Per
88.92 #0.67
80.40
#0.37
30.49
+4.52
81.23 +1.43
70.26 #0.72
MEMO (P)
88.77 10.20
81.86
10.35
28.59
#2.84
81.49 #1.13
70.18 20.35
FedTHE (Ours)
89.21 20.55
81.87
+1.11
57.80
+3.52
80.71
+1.33
71.54 10.69
FedTHE+ (Ours)
89.91
10.46
82.16
11.16
58.03
13.42
81.64
10.72
71.79 10.74
E.5 MORE ANALYSIS ON HEAD ENSEMBLE WEIGHT
Here we provide visualization and analysis on head ensemble weight e, which is the key of our proposed FedTHE, where in Figure 17 we visualize the distribution (probability density) of 1 (which is the ensemble weight for the local head) across various test distributions including 1 ID and 4 OOD distributions. Note that each e corresponds to a single test sample from the corresponding test distribution. We can see that:
For in-distribution case: for most samples (major classes), large weights on the local head are learned since the local head fits local distribution; for a very small portion of samples that belong to tailed local classes, a smaller ensemble weight on local head is learned, in order to better utilize the knowledge from the global head.
For in in-distribution/ common corruptions/ / natural distribution shift cases, the distribution of reflects the insights in 'Accuracy-on-the-inne (Miller et al., 2021), thus similar patterns
33

---

Published as a conference paper at ICLR 2023
Original local lest (in-diatribution)
Corrupted local test (common comuptions)
20
Probabiitty denaty
Probataiity densty
15
15
10
10
8.
0.0
0.2
1.0
0.4
0.8
0.8
1.0
(denctes ensemble weight for local head)
1 e (denotes ensemble weight for local head)
Original OoC test (label distribution shift)
Naturally shifted test (natural distribution shift)
20
Probabiitt dansit
Probability density
15
10
0.0
0.2
0.6
1.0
- (denates ensemble weight for local head)
1
1-e (denates ensemble weight for local head)
Mixture of test distributions
20
Probability
10
1 1-e (denotes ensemble weight for local head)
Figure 17: The distribution (probability density) of e across clients.
with some differences (e.g. the usage of global model increases in common corruptions and natural distribution shift) can be observed.
For label distribution shift case, our method learns to utilize global head as the main power of prediction when encountering tailed unseen classes, while for smaller portion of test samples whose labels fall in the present set local training labels, the local head is favored.
E.6 MORE ANALYSIS ON THE EFFECTS OF history
As a variance reduction is quite important in per-sample adaptation, history and the moving average supposed to stabilize the choice of in FedTHE and FedTHE+. We demonstrate the effectiveness and effects of history here. As shown in Figure 18, we visualize the feature representation with and without h history with t-SNE in 2-dimension, and when history considered, the clusters more concentrated, and the effects of this concentrated pattern are discussed below. For the effects of history, show that, in Figure 19, the history makes difference on choosing a proper e. Accompanying with the 1 in Table we see that such difference is crucial on the preformance of the method.
34

---

Published as a conference paper at ICLR 2023
Original local test (n-distribution
Corrupted (common corruptions)
A Aitcy N/stooy
0.75
0.50
0.5
X
-0.5
by Nittory
wr h Nottoy
-1.00
-1.0
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.78
1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
x1
x1
Original GoC test (labal distribution shift)
Naturally shitted test (natural distrbution shift)
0.75
0.75
0.50
0.50
X
0.00
0.00
-0.25
0.50
wio Aittary
wo Arittoy
-075
w
Anittor
-1.00
-0.75
-0.50
-0.25
1.00
0.25
0.50
0.75
-1.00
-0.75
-0.50
0.25
0.50
x1
x1
Mixture test distributions
1.0
0.5
0.0
-0.5
W/o i Puttey
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
1.00
x1
Figure 18: The distribution (probability density) across clients,
35

---

Published as a conference paper at ICLR 2023
Original local lest (in-iitribuutin)
Corrupted local test (common coruptions)
Wt hittry
w/ h
20
wio
wio A hotory
Probabiltty domalty
Probability donnity
15
1.5
10
e (dendtes ensemble weight for local head)
1 e (denates ensemble weight for local head)
Original OoC (labal distribution shift)
Naturally shifted (natural distribution shift)
W/ Ahtttryy
wi h
20
w/o
/)
Probabiity density
Probabiity density
1.5
10
- (denctes ensemble weight for local head)
1
e (denates ensemble weight for local head)
test distributions
20
Probabiitty density
0.6
0.8
1 (demotes ensemble weight for local head)
Figure 19: The distribution (probability density) across clients.
36


Reviews:
[{"review_id": "wi-mAAIpSt", "reviewer": "Reviewer_wWXU", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper considers the federated learning problem under the test-time distribution shift. Federated learning is a distributed learning paradigm with a server and many clients such that each client has its private data, and the server aggregates the trained model from clients without revealing any private information. Many previous works focused on the performance of the server model, but recently personalization on FL has been studied. However, most works assume distribution matching between training and testing, whereas, in practice, the class distribution can change, unseen classes can appear, and common corruptions or natural distribution shifts can happen. Therefore, it is crucial to design robust FL algorithms for dealing with such distribution shifts. In this work, the authors propose Federated Test-time Head Ensemble plus tuning (FedTHE+), which personalizes FL models robustly to the distribution shifts. They evaluate FedTHE+ with various neural architectures on CIFAR10 and ImageNet with diverse test distributions. ", "summary_of_the_review": "This paper suggests an interesting direction, FL robustness. However, this paper does not provide a good intuition for the topic.", "strength_and_weaknesses": "Strengths\n\n-This paper combines test time training and federated learning, which is an interesting direction.\n\n-The authors provide a benchmark for FL robustness\n\n-The proposed algorithm, FedTHE+, outperforms baselines in most of the scenarios.\n\n\nWeakness\n\n- The authors do not clearly explain why each step of the algorithm is required and for what reason the algorithm can be good for FL. \n\n- FedTHE, the proposed algorithm, consists of many parts. Many parts have synergy between them, and the synergy produces good performances. However, such a complicated structure makes it difficult to understand which part is the key to dealing with distribution shifts. \n\n- Since this paper has too much content with a page limit, many details are omitted in the main body and deferred to the appendix. So it took a lot of work to read.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "l5ZMyUIu2E", "reviewer": "Reviewer_129X", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work introduces FedTHE(+), a way to allow for test-time personalization of federated learning models. FedTHE works by introducing and training a scalar that interpolates global and local classifier predictions. This scalar is tuned at test time by minimising the entropy of the interpolated logits while being regularised by an alignment loss between the features. The latter giudes the scalar according to the difference of the representation of the test point to the local and global feature means; if the representation is closer to the global one it increases the scalar otherwise it decreases it. To balance the two losses, the authors propose to use the cosine similarity between the local and global logits as a “weight” that linearly interpolates the two. After this specific scalar is optimized, the other parameters of the network are also fine-tuned using an entropy-minimisation method, which constitutes the FedTHE+. \n", "summary_of_the_review": "This is an interesting work and discusses a relatively simple method FedTHE and its extension FedTHE+. FedTHE is novel and the experimental evaluation is quite extensive, including multiple distribution shift scenarios, with good results and the authors do extensive ablations to show the effects of each of their choices in FedTHE, which is appreciated. For this reason, I am recommending acceptance of this work. ", "strength_and_weaknesses": "Strengths\n- Extensive evaluation and ablation studies\n- Hyperparameter tuning for the baselines as well\n- Simple method and computationally efficient to apply, which is an important aspect for FL\n- Good results\n\nWeaknesses\n- A lot of moving parts\n- The motivation behind some choices is a bit unclear", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "SDMldC4Ffwo", "reviewer": "Reviewer_JEHY", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed a new personalization method in federated learning: Federated Test-time\nHead Ensemble plus tuning (FedTHE+). FedTHE trained a global feature extractor, a global head and local heads for classification, and during the personalization time, a scalar e is learnt to interpolate between predictions by global and local heads. \n", "summary_of_the_review": "I would be happy to discuss the clarification questions with the authors. ", "strength_and_weaknesses": "Distribution shift/robust personalization is an important problem that only draws attention recently. \n\nThe proposed method of interpolation between local and global heads sounds intuitive. \n\nI appreciate the extensive experiments, but also have some concerns on clarity and reproducibility listed below. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Ig20LuwOYfl", "reviewer": "Reviewer_bhMe", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the OOD generalization problem in the FL setting. Specifically, the author presents to use a test-time-training based strategy to ensemble the global and local prediction heads. The object function use in the TTT phase consists of two parts (1) the entropy of the predictive distribution. (2) feature space constrains. The empirical results have shown significant improvements over various baseline models.", "summary_of_the_review": "See above.", "strength_and_weaknesses": "# Pros\n1. The idea of using test-time-training to improve the model robustnees in FL is novel and effective.\n2. The empirical results are comprehensive and can validate the superiority of the proposed method.\n# Cons\n1. Acctually, this paper is not directly for OOD generalization, i.e., the local and global training is still based on standard FedAVG. There are some previous work [1,2] focus on optimizing the worst-case local loss needed to be discussed. In addition, is it possible to train different prediction heads by distributionnaly robust optimization?\n2. There are many critical hyperparameters need to be tuned.\n# Comments\n1. It seems that this paper is a re-submission of NIPS, can the author provide the critical difference between the current and previous version?\n# Refs\n[1] Distributionally Robust Federated Averaging. \n[2] DRFLM: Distributionally Robust Federated Learning with Inter-client Noise via Local Mixup", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster