Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper proposed a new personalization method in federated learning: Federated Test-time
Head 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. 

strength_and_weaknesses: Distribution shift/robust personalization is an important problem that only draws attention recently. 

The proposed method of interpolation between local and global heads sounds intuitive. 

I appreciate the extensive experiments, but also have some concerns on clarity and reproducibility listed below. 

clarity,_quality,_novelty_and_reproducibility: Is this method stateful, or stateless (A Field Guide to Federated Optimization Section 7.5 https://arxiv.org/abs/2107.06917)? Is the local head w^l maintained on each client? How many clients in experiments, and is there client sampling for each round?

How is h^{history} computed? Is it intra-client, or inter-client? Why is moving average used, which makes it depending on the order of “history”?

Though I appreciate the extensive results, table 2 is a bit hard to interpret. Maybe I missed it, I would personally see more analysis on e, which is the key of the proposed method. For example, what is the distribution of e for various clients, and in various distribution settings? Why did the proposed method work for various shifts while the previous methods fail. 

Both the proposed FedTHE+ and the dataset in section 5 sound somewhat complicated. Would the authors release both dataset and code for reproducibility? 

summary_of_the_review: I would be happy to discuss the clarification questions with the authors.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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. 
strength_and_weaknesses: Strengths

-This paper combines test time training and federated learning, which is an interesting direction.

-The authors provide a benchmark for FL robustness

-The proposed algorithm, FedTHE+, outperforms baselines in most of the scenarios.


Weakness

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

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

- 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.
clarity,_quality,_novelty_and_reproducibility: Due to the page limit, the algorithm is not precisely described in the main text. Many detailed steps are deferred to the appendix. 

They provide the source code, although I have not tried to run the code. 
summary_of_the_review: This paper suggests an interesting direction, FL robustness. However, this paper does not provide a good intuition for the topic.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
strength_and_weaknesses: # Pros
1. The idea of using test-time-training to improve the model robustnees in FL is novel and effective.
2. The empirical results are comprehensive and can validate the superiority of the proposed method.
# Cons
1. 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?
2. There are many critical hyperparameters need to be tuned.
# Comments
1. It seems that this paper is a re-submission of NIPS, can the author provide the critical difference between the current and previous version?
# Refs
[1] Distributionally Robust Federated Averaging. 
[2] DRFLM: Distributionally Robust Federated Learning with Inter-client Noise via Local Mixup
clarity,_quality,_novelty_and_reproducibility: See above.
summary_of_the_review: See above.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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+. 

strength_and_weaknesses: Strengths
- Extensive evaluation and ablation studies
- Hyperparameter tuning for the baselines as well
- Simple method and computationally efficient to apply, which is an important aspect for FL
- Good results

Weaknesses
- A lot of moving parts
- The motivation behind some choices is a bit unclear
clarity,_quality,_novelty_and_reproducibility: This paper is mostly clear and well written. One minor comment is that there seems to be a bit excessive use of \vspace that makes the text look cramped, so I would suggest for some things to be moved to the appendix so that the paper looks better. Apart from that, another minor comment is that MEMO (G) needs to be bold at Table 1 for Dir(0.1), OoC local test and Dir(1.0) OoC local test. 

While this work does have a lot of moving parts, the authors do extensive experimental evaluation and ablation studies to show the effects of each one. Besides that, the authors also perform an extensive hyperparameter optimization procedure for the baselines as well. I believe that there is enough information in this work for reproducibility. 

As far as I am aware, FedTHE is novel. Furthermore, its simplicity coupled with the extensive distribution shifts considered, makes it quite appealing for FL type of scenarios. One nitpick on the loss is that the “feature alignment” is a bit misleading, as the features are frozen. Something like adapting the scalar according to the feature distance would be more appropriate. The similarity-guided loss weighting is also an interesting and intuitive contribution. FedTHE+ is not novel per se, as the authors just use existing test-time adaptation methods (i.e., MEMO) to fine-tune the model. One catch for FedTHE though is that it needs to send summary statistics of the client data to the server ($h_l$), which I believe should be discussed in the context of the baselines (e.g., do baselines employ such summary statistics?) and in the context of the privacy / security aspects of FL.


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.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This paper proposes a new way to personalize the model in federated learning. All the reviewers think that personalization in federated learning is an important problem, and the proposed algorithm is well-motivated. The empirical results also support the claims. Despite having sufficient novelty, the clarity of the paper can be improved. Therefore, we recommend acceptance for this paper. 

As suggested by Reviewer wWXU, the paper can be improved by adding more explanations and ablation studies for each step of the proposed method. We hope the authors can address the concerns in the final version. 


justification_for_why_not_higher_score: The algorithm contains several parts, and there are insufficient explanations/ablation studies to demonstrate the usefulness of each part. 
justification_for_why_not_lower_score: The paper has sufficient novelty and is tackling an important problem.
Decision: Accept: poster