Adaptive Test-Time Personalization for Federated Learning

Summary

This paper considers the federated learning setting of adaptive test-time personalization. Traditional test-time adaptation (TTA) can only handle specific target domain distributions, while federated learning requires flexible handling of multiple target domains. Existing TTA methods pre-define which modules to adapt, which limits the application of TTA in federated learning. Therefore, this paper proposes the Adaptive Test-time Personalization algorithm called ATP to automatically decide which modules to adapt and how much to adapt.

Strengths

1. The experiment in Section 3.2 is crucial as it effectively illustrate the inherent challenges encountered by current TTA methods when applied to federated learning, thus offering valuable insights and guiding directions to potential enhancements. 2. The proposed method is simple and easy to implement, and it has achieved impressive performance in the current experiments.

Weaknesses

1. The paper claims that it is the first to propose test-time personalized federated learning. However, this claim is questionable because previous works, such as [1], had already explored test-time personalized tasks in the context of federated learning. [1] Jiang, Liangze, and Tao Lin. "Test-Time Robust Personalization for Federated Learning." ICLR2023 (preprint arXiv:2205.10920 (2022)) 2. I am concerned about the significance of the “supervised refinement” step. If labeled data is available in this step, using the labeled data itself already leaks the distribution about the test samples in TTA tasks (while the test distribution is not known in the TTA setting), which obviously reduces the difficulty of TTA. If labeled data is not available, the proposed method for adaptively learning alpha seems unworkable. 3. The datasets used in the experiments are small-scale. It would be more convincing if experiments were conducted on real-world images, such as DomainNet, with at least a resolution of 224x224.

Questions

1. In Figure 2, why does setting "bn.running_mean (m=-0.1)" improve the accuracy of TTA in the presence of label shift, while "bn.running_mean (m=0.1)" significantly harms the accuracy? This seems counterintuitive. 2. The task addressed in this paper is more akin to source-free unsupervised domain adaptation rather than Test-Time adaptation. In my opinion, the focus of most TTA works is on an online setting where the entire test set cannot be obtained at once, and operations are performed on individual test samples to be predicted (as the details in Section 4.3). 3. See Weaknesses 2, how important is the step of "supervised refinement"? Can the authors provide any ablation experiments?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

There is no limitations or potential negative societal impact discussed in this paper.

Summary

The paper studies test-time personalization in a federated learning setting --- after training on participating clients, the goal is to locally adapt the global model given unlabeled test data. The paper's main idea is by pointing out that label non-IID and domain non-IID require adaptation on different layers of DNNs, and propose a novel way to learn the adaptation learning rates of each layer automatically in a data-driven fashion. A simple learning method that alternatively do SGD on the DNN parameters and the layer-wise learning rates shows effective improvements in test-time personalization.

Strengths

- The method proposes a novel aspect that to tailor different types of non-IIDness, the degrees of adaptation are different across layers. it is an interesting point. The solution is simple, sound, and effective. - Empirical results are overall satisfying. The experiments provide enough comparisons to centralized TTA algorithms; see some suggestions below.

Weaknesses

Although I am positive on this paper, I observe several important concerns. If the authors could address them, my score can be higher. [Major 1] Overclaims: considering TTA in PFL setting was first introduced in [15]. This is a natural extension of centralized TTA. It is unnecessary and imprecise for this paper to make "test-time personalized federated learning (TTPFL)" a new setting. The definition in L38 about the combination of distribution shifts of labels and styles itself is nothing to do with federated setting; centralized TTA can have both label and style shifts. [Major 2] Misleading section 4.2. I cannot understand why the proposed refinement has a connection with meta-learning. ClientTrain in Alg 1 is simply a coordinate gradient style optimization. Alternatively, a discussion about hyperparameter optimization should be more relevant. [Major 3] Theorem 5.1 is not informative. It's simply extending the classic FedAvg generalization bound [18] for the adaptive parameters. It will be more informative and fit to the context to discuss the generalization to the new clients of different distributions ''after test-time personalization'' in TTA sense. Otherwise, I would recommend avoid this laundry theorem. [Minor 1] The datasets are rather small scales and synthetic. it will be better to include natural federated datasets like FEMNIST or iNaturalist-GEO or large-scale common datasets in TTA (like ImageNet). [Minor 2] Why FedTHE or FedTHE+ in [15] are not compared? [15] seems to be the closest work about TTA in FL. Is the discussion in L76 faithful? It seems to me FedTHE does not need labeled data.

Questions

Section 3.2 is the limitation to TTA but not about PFL. Can the author discuss why the problem or the solution is dedicated to FL? I suggest adding a discussion around L134. In my opinion, the key is that in PFL we can figure out the ideal adaptive rates by leveraging the training clients who have diverse distributions.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

See weakness.

Summary

This paper proposes a new setting called test-time personalized federated learning (TTPFL) and proposes an Adaptive Test-time Personalization algorithm. The authors show effectiveness of proposed method over other test-time adaptation methods.

Strengths

The paper proposes an Adaptive Test-time Personalization algorithm and shows its effectiveness over other test time adaptation methods.

Weaknesses

1. The proposed setting is strange and not self-consistent. The authors claim that in this setting 'clients adapt a trained global model in an unsupervised manner without requiring any labeled data.' However, the proposed method involves labeled clients to learn the learning rates for different modules (Figure 3, left). At least, real-world scenarios as examples should be provided. 2. Missing representative federated learning baselines. This paper does not compare with any federated learning method. If the proposed method uses the labeled datasets for the training stage, then many federated learning methods can be seen as baselines that are trained on the labeled datasets. Baseslines include FedAvg, FedAvg with Fine-tuning, FedProx, FedProx with finetuning, Ditto, pFedMe. 3. Since the proposed method focuses on different operations on different modules of a model, model architecture is important to this paper, thus experiments on different model architectures are required.

Questions

See weakness.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Generally, I am really confused about the setting in this paper after reading. Please provide comprehensive explanations to correct me if I am wrong, and I would consider re-rating.

Summary

This paper introduces a novel setting where personalized FL during the test procedure is considered and multiple distribution shifts are involved. A method termed ATP is proposed to solve the challenges posed in this setting. Adaptive learning rates are learned for the model. Both theoretical and empirical studies are carried out to demonstrate the effectiveness of ATP.

Strengths

1. The paper is well-organized and easy to follow. 2. The main claims, e.g., the capability of dealing with multiple distribution shifts and the effectiveness of adaptive learning schemes, are well supported by empirical studies. 3. Theoretical analyses on both the convergence and generalization ability of ATP are provided.

Weaknesses

1. One of the key factors of TTPFL is confusing. Usually, we assume test data are unlabeled and the target of the classification task is to predict the label of the test data. However, in the summarization of TTPFL, it is emphasized that each testing client only has unlabeled data for personalization. It is better to further clarify this factor. 2. The relation between FL and test-time shift is weak. It seems that the proposed adapting trainable parameters, adapting running statistics, and adapting rates can also benefit centralized test-time shift problems. For FL systems, the unique challenges brought by test-time shift issues and how they motivate these adapting solutions are not explicitly demonstrated. 3. Lack of detailed discussion on the difference between this work and the previous study[15]. [15] also considers feature shift, label shift, and a mixture of these shifts in their recent paper. There are various feature shifts in [15], which can also be considered as part of the experimental setting in this work.

Questions

In Section 4.2, the authors have mentioned the smaller communication costs achieved by ATP, are there any corresponding experimental results that demonstrate this in a quantitive manner?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

This paper does not adequately address the limitations in terms of privacy issues, efficiency, etc.

Follow-up response regarding different model architectures

Dear Reviewer h3BM, Thanks again for your valuable feedback on our paper. Regarding the question of different model architecture, we would like to provide additional experiment results. Consider that smaller models are used in FL for their smaller communicational and computational cost, we test our algorithm with a 5-layer CNN on CIFAR-10 dataset. Other experimental settings are identical to our CIFAR-10 experiment in Section 6. As shown in Table D below, similar to our results in the paper, our proposed ATP achieves better performance than various baselines across three types of distribution shifts, which further validate the effectiveness of ATP. Table D: Accuracy (%) under three kinds of distribution shifts on CIFAR-10. We report the average accuracy on testing clients. | Method | Feature shift | Label shift | Hybrid shift | Avg. Rank | | ------------- | ------------- | ----------- | ------------ | --------- | | No adaptation | 64.33 | 69.15 | 61.87 | 7.7 | | BN | 66.51 | 54.60 | 50.20 | 7 | | SHOT | 65.70 | 49.93 | 45.89 | 8.3 | | Tent | 65.57 | 50.02 | 45.72 | 9 | | T3A | 64.33 | 66.34 | 59.59 | 8 | | MEMO | 65.64 | 71.77 | 64.16 | 5.7 | | EM | 61.65 | *76.16* | 67.06 | 5 | | BBSE | 56.92 | 76.13 | 66.21 | 6.3 | | Surgical | 64.45 | 73.58 | 65.35 | 5.7 | | ATP-Episodic | *66.88* | 76.14 | *68.48* | *2.3* | | ATP-Online | **67.11** | **78.38** | **70.86** | **1** |

Gentle Reminder

Dear Reviewers, Thank you very much for your time and efforts on our paper. We are writing to follow up on our rebuttal. We have carefully considered the raised questions and made our responses as below. As the discussion period is halfway through, we gently ask for a retrospect regarding our rebuttal and whether there are any remaining questions or concerns that you would like to address. Thanks again for your time and consideration!