Paper: All for One and One for All: A Collaborative FL Framework for Generic Federated Learning with Personalized Plug-ins (ICLR-2024)

Reviewer comment (Reproducibility):
The test-time strategy of this method is somewhat impractical and unclear, as whenever a test sample appears, the scheme requires the client to download all the features from the server (which is the same as the number of clients in the system) and compute (8). This can incur communication, computation, and latency issues.

Author response:
**Authors' Response:** Thanks for pointing out potential problems when HPFL meets with an FL system with a great number of clients. As the number of clients grows to a large number, clients can choose to add noise to their local test features and send the noised test features to the server to select the plug-ins. In this way, each client can get the required plug-in without the need to download all the plug-ins and the noised training features, which cause a great communication cost with a great number of clients in the FL. We conduct experiments to testify the feasibility of this method against the communication burden of HPFL in FL systems with plenty of clients. Please find them in the [Solution toward shared scalability and privacy concern](https://openreview.net/forum?id=8rhHI6C8iC&noteId=ULBhGHenqN).

Reviewer comment (Reproducibility):
During inference, each client may need to download multiple plug-in modules, which will again incur additional communication and latency. For example, suppose there are three clients, each having classes (1, 2), (3, 4), (5, 6) in its local dataset, and produce plug-in modules 1, 2, 3, respectively. Now for an arbitrary client X, when the first test sample belongs to class 1, client X will download the plug-in module 1. Then, if the next test sample belongs to class 2, client X will download the plug-in module 2. Finally, it will also download plug-in module 3 whenever classes 5 or 6 appear in the test set. It is unclear whether this is really happening in the proposed method.

Author response:
**Authors' Response:** HPFL is designed to infer in batches, online inference where test samples arrive one by one is not our main application scenario [1, 2, 3]. Selecting plug-ins for every sample may incur expensive computation costs. However, when met with a similar situation and the computation costs are unavoidable, clients can give up downloading plug-ins. Similar to method stated in [Solution toward shared scalability and privacy concern](https://openreview.net/forum?id=8rhHI6C8iC&noteId=ULBhGHenqN), following instructions free clients from the storage burden brought by downloading plug-ins and corresponding training features for every sample: simply send the noised test feature to the server; the server select the appropriate plug-in; **then infer at the server side; and finally return the inference result back to the client**. In this way, communication and latency issues with online inference can be solved with slight modifications in HPFL. We have added ***Appendix F.2.2 in our revision*** to discuss how HPFL can adjust to special circumstances like online inference.

Reviewer comment (Experiments):
Related to the previous comment, the authors' scheme might have an advantage compared to (Chen & Chao, 2021) when considering a mixed dataset of personalized and generic datasets, similar to what is done in [Ref]: each client's test dataset has more local classes when the distribution shift is small, while having more classes outside of the client's local dataset when the distribution shift is large. This potential advantage should be discussed and validated.

Author response:
**Authors' Response:** **(a)** While FedRoD [4] can handle both GFL-GM and PFL-PM by two plug-in heads, our proposed setting GFL-PM is different from GFL-GM [5] and PFL-PM [6], because in our setting distribution of local test data may not the same as that of training data, in this situation, a single personalized classifier is not enough for the diversified test data. **(b)** Our experiment also shows that FedRoD's local classifiers trained on local training data can't handle test data in different distributions, leading to poor performance in our proposed GFL-PM setting. It fails to provide a more fine-grained solution compared with our methods. As shown in ***Table 2 of origin main text***, even using the global classifier to infer on every client in the GFL-PM setting (it is the same as the performance in the GFL-GM setting as the model in every client is identical), our method shows a significant performance gain compared with FedRoD in most cases. **\(c\)** Our proposed set-up GFL-PM is totally different from the setting in [7], which is an interpolated dataset of local test data and global test data. In our set-up, clients may meet with local test data from both other clients and themselves. Since the problem we are trying to solve is quite divergent, we are not sure if our proposed method HPFL is able to perform well in setting in [7], as we assign a single plug-in from other clients or the client itself to the client, instead of interpolating the local head and the global head.

Reviewer comment (Novelty):
The combination of using a separate backbone and classifier in PFL for improvement, along with the idea of customize the classifier on local client, is not novel. The former has been widely validated for its effectiveness in FedPer and FedRep, while the latter's idea is highly relevant to that of KNN-Per [1] and FedBABU [2]. However, the authors do not provide relevant discussion or experimental comparisons of the latter ones.

Author response:
Thanks for providing us with two excellent methods in the PFL setting for comparing our methods with. **KNN-Per (a)** [1] and **FedBABU (b)** [2] share the similar problem of classic PFL algorithms [3-6] in our GFL-PM setting. 
- **KNN-Per (a)** build a local datastore $S_m$ with local training data and query $S_m$ to assist inference using a knn method. Imagine a circumstance where some classes in local test data are not presented in local training data, in this situation, the test samples of those classes are likely to retrieve a wrong label from $S_m$ and mislead the final decision. 
- In **FedBABU (b)**, classifiers of local clients are randomly initialized and never updated during federated training, and for personalization during the evaluation process, FedBABU fine-tuned its classifiers on local training data. Here, FedBABU also requires local test data to be in the same distribution as local training data to fine-tune the classifiers on local training data to perform well on local test data. 

**These two methods only customize the classifier to fit static local test data with the same distribution as its training data.** Instead, our proposed setting GFL-PM is faced with the problem of changing local test data distributions, whose distribution is different from local training data, and HPFL is proposed to solve this problem. These two methods don't propose **the adaption mechanism to test data, which is implemented with plug-in selection in HPFL**, thus probably face with similar difficulty of classic PFL algorithms like FedPer and FedRep when being exploited in GFL-PM setting.

Reviewer comment (Experiments):
Since the authors compared the performance of both GFL and PFL, the experiments should include more commonly used GFL-related baselines.

Author response:
We'd like to thank you for pointing out the potential deficiency in GFL baselines. For comparison in the GFL setting in our experiments, FedAvg is a classic and most commonly used GFL algorithm, and FedRod can simultaneously achieve leading generic and personalized performance. For more comprehensive comparison in the GFL setting, we prepare to compare HPFL with a more recent SOTA GFL method FedSAM [7]. However, due to the time needed to carry out the experiments, we may need several days to report the result of FedSAM and add the result to our revision.

Reviewer comment (Experiments):
The performance of PerFedMask (Table 1) in the personalized setting is much lower than other baselines. I wonder if there was an error in implementing this method, or if its hyper-parameters were not appropriately fine-tuned?

Author response:
Our results of PerFedMask is presented in ***Table 2 of origin main text*** instead of ***Table 1 of origin main text***. We followed the open-source implementation of PerFedMask [6]. To fairly compare it with other baselines and HPFL, we only change the data partition method, the learning rate scheduler, and the communication round to align with other methods. To testify the scalability of methods, we increase the number of clients, and raise batch size from 50 to 128 as the value we use in other methods for better convergence. All other hyperparameters, including the initial learning rate, stay the same as the origin implementation of PerFedMask. We have added these information about hyperparameters to ***Appendix C.4 in our revision***.

Reviewer comment (Evaluation):
About the experiments with added noise (Table 4). Why do experiments with different levels of noise produce nearly the same results? This appears somewhat counterintuitive. Can the authors provide a more detailed analysis?

Author response:
For this counterintuitive phenomenon, we provide our explaination in ***Section 5.3 of origin main text***. Since our selection method only chooses plug-ins with selection scores, if the changes in selection scores are not significant enough to affect the final choices of plug-ins, which can be observed in the MMD scores heatmaps in ***Figure 3 of origin main text***, the performance of HPFL will remain the same, which adds robustness towards noise to HPFL without significant loss on accuracy performance. We have added a detailed discussion about this phenomenon in ***Section 5.3 of our revision***.

Reviewer comment (Experiments):
Some parts of the design choices of the proposed selective FL could be elaborated more to give the paper a broader scope. For example, the part of the general feature extractors can be trained with any GFL algorithm. Considering there are many GFL methods proposed, what are their performances in comparison? Are they compatible with selective FL? For the selection methods in section 4.3, could the author provide more insights and discussion besides the used distance-based methods? Any other potential replacements?

Author response:
**Authors' Response:** Thanks for the advice on improving the ablation study of our method. **(a)** As long as the used GFL methods are able to train a strong general feature extractor, HPFL is able to utilize the feature extractor to train the personalized plug-ins and extract features. Due to the great availability of GFL methods used for training backbone [1-5], we didn't carry out experiments to see how different GFL backbone training methods affect the overall performance of HPFL. We add experiments using FedRoD [6] to testify HPFL's compatibility with other GFL methods. The results are given in ***Table 2*** below. Number of clients equal $10$, local fine-tuning epoch $E_p=10$, local datasets are partitioned in $Dir(0.1)$. Other settings remain the same as the main experiments in ***Table 2 of origin main text***:

**Table 2**

|                   |                  | HPFL (FedAvg) | HPFL (FedRoD) |
| :-----------------: | :----------------: | :-------------: | :-------------: |
| **CIFAR-10**      | GFL-GM           | 81.5          | 85.3(+3.8)     |
|                   | GFL-PM           | 95.7          | 96.0(+0.3)     |
|                   | PFL-PM           | 95.7          | 96.0(+0.3)     |
| **FMNIST**        | GFL-GM           | 86.0          | 87.9(+1.9)     |
|                   | GFL-PM           | 98.4          | 98.4(+0)       |
|                   | PFL-PM           | 98.4          | 98.4(+0)       |
| **CIFAR-100**     | GFL-GM           | 68.6          | 69.9(+1.3)     |
|                   | GFL-PM           | 72.2          | 68.5(-3.7)     |
|                   | PFL-PM           | 85.7          | 85.5(-0.2)     |
| **Tiny-ImageNet-200** | GFL-GM        | 56.5          | 57.4(+0.9)     |
|                   | GFL-PM           | 50.9          | 56.0(+5.1)     |
|                   | PFL-PM           | 73.7          | 74.7(+1.0)     |


From the overall performance of HPFL(FedRoD), we can see that HPFL using FedRoD as its backbone training method is comparable to that using FedAvg, which confirms HPFL is compatible with the GFL methods other than FedAvg. Following your valuable comments, in ***Section 5.2 of main text of our revision***, we have directed readers to results of HPFL(FedRoD) shown in ***Appendix D.6 in our revision*** to discuss the ablation study of GFL backbone training methods. **(b)** As for selection methods other than distance-based methods, due to the page limit, we choose to introduce other selection methods in the main body. In the ***Appendix D.2***, we elaborately discuss an OOD-Detection (Out of Distribution Detection)-based method [7] to select suitable plug-ins by judging whether the test data is OOD for plug-ins and choose the plug-in with the lowest OOD score. Aside from OOD-Detection-based methods, if a method can be used to judge how appropriate plug-ins are for certain test data, it is available for selection. Even a supervised classifier trained with the noised training features as data and corresponding client indexes as labels may be able to select plug-ins.

Reviewer comment (Experiments):
Considering FedRoD is the sota in this topic evidenced by Table 2, can the authors provide a more fine-grained comparison? Specifically, FedRod (also followed by FedTHE) proposes to train the general model with a class-balanced loss; will this improve HPFL as well? FedRoD also fine-tunes the whole general model for each client with local training; how many personalized epochs are trained like the reported HPFL (e.g., E_p=1 or 10)?

Author response:
**Authors' Response:** Yes, and thank the reviewer for the suggestions about a more elaborated discussion on FedRoD. **(a)** From the results of HPFL(FedRoD) given in ***Response to Q1***, we observe an interesting fact that even if FedRoD shows excellent performance in GFL-GM (surpasses FedAvg in many settings), fine-tuning the backbone trained with it is not advantageous as shown in PFL-PM (only comparable with fine-tuning on the backbone trained with FedAvg). From this phenomenon, we presume the advantage of FedRoD in GFL-GM should mainly be attributed to its global head trained with a class-balanced loss instead of its backbone. **(b)** As is stated in ***Section 5.1 of origin main text***, the local training epoch of all methods is $1$. To avoid the possible misunderstanding, we would like to clarify that the local training epoch is different from the $E_p$ in HPFL, because the backbone will not be updated when fine-tuning the plug-ins in HPFL. Fine-tuning of HPFL happens after the GFL backbone training and is completely done in local clients. We apologize for not clearly indicating personalized epochs of other baseline methods with personalized epochs like PerFedMask. We have specified the value of personalized epochs for those methods with personalized epochs in ***Table 2 of main text in our revision***.

Reviewer comment (Writing):
There are many typos and inconsistent wording.

Author response:
**Authors' Response:** We apologize for the typos and inconsistent wording. We will check the whole paper thoroughly and carefully, and revise mistakes in it.

Reviewer comment (Theory):
Can the authors provide a discussion on the data/client assumptions on when should HPFL work or not work? Suppose the new client has very different distributions unseen in the training, how possible will HPFL work well?

Author response:
**Authors' Response:** Sure. **(a)** HPFL works well in situations where the distribution of local training data and test data is quite different, but its test data distribution is similar to the distribution of other clients' training data, HPFL can help the client by borrowing knowledge from other clients, i.e., utilize plug-ins trained by other clients. **(b)** For the clients that meet with totally different test data distributions that haven't been seen in the training data of all clients (include those clients themselves; if those clients have similar distribution on training data and test data, they can choose the plug-in trained on their own local data). However, if there is really no similar data pattern even in the whole FL system, HPFL can only attempt to choose the most appropriate plug-in, though in this situation, even if the most appropriate plug-in may not perform well on the totally strange test data due to not identical training and test distribution [8]. But it is worth noting that with the increasing number of clients in the FL system, this situation happens significantly less with the abundance of clients' training data distribution.

Reviewer comment (Presentation):
The main weakness of the paper is its presentation. In particular, the manuscript is not explicit about the form of the test distribution that the clients will encounter. I was confused about this, and could only infer it until I reached Section 3.3 and I finally understood the problem's objective.

Author response:
**Authors' Response:** We are sorry that we failed to highlight the test scenarios of our problem, which is briefly illustrated in the ***3rd paragraph in the introduction of origin main text (Section 1)***: "i.e., clients equally meet with the test data from all clients." We chose to elaborate more in our theory part for this problem setting. We have added a detailed explanation, "In GFL-PM setting, the test set every client encounters comes from the local test data of other clients, whose distribution is the same as that of clients' local training data." in the ***3rd paragraph in the introduction of our revision*** to help readers to more easily understand our proposed problem.

Reviewer comment (Theory):
Storing all the clients' plug-ins may also be a privacy risk, as there is not any aggregation that protects the clients' privacy from a malicious server. The threat model is not discussed by the paper.

Author response:
**Authors' Response:** Due to page limit, we discussed the potential privacy risk of sharing plug-ins in ***Appendix E.1***. 
- As stated in ***Section 4.3 of origin main text***, **sharing part of models (i.e. plug-ins) is exposed to actually no more privacy risk than sharing the parameter gradients or the complete global model as adopted in classic FL algorithms like FedAvg** [7], which in practice is often protected with regular privacy protect techniques like differential privacy (DP) [8]. 
- **As clients in HPFL don't upload the gradient during plug-ins training**, no gradient of models is accessible for the server, so there is no need to worry about gradient inversion attacks [9-12] towards plug-ins. 
- **We also carry out experiments to testify the safety of our method toward the model inversion attack** [13,14,15] when with both the noise features and the model trained on raw features in ***Appendix E.2***. 
- **Even extreme concern on potential privacy risk of keeping plug-ins in the server can be solved by only requesting plug-ins after selection** as described in ***Response to Q2***, clients providing the plug-ins can ask the server to delete the plug-ins after sending the plug-ins to the clients in need. To increase the visibility of our discussions about the privacy problem of sharing plug-ins in ***Appendix E.1***, we have directed the reader to them in ***Section 4.2 of main text in our revision***.

Reviewer comment (Novelty):
Although the given set-up is novel, I am still not convinced of its importance, at least not justify the drawbacks of the proposed solution. I recommend the manuscript update its introduction to include more motivating examples that are explicitly related to their problem formulation.

Author response:
**Authors' Response:** We appreciate the suggestion given by the reviewer to justify the motivation of our proposed set-up. We give some examples exhibiting the significance of our proposed set-up below: 

- **Case 1: Some clients may have insufficient computing resources or local training data to fine-tune a deep learning model in a cross-device setting.** In these situations, training distribution can be regarded as an empty set $\emptyset$. In this way, the client cannot get a personalized model by locally fine-tuning the global model. In traditional GFL and PFL settings, the client has no choice but to adopt the global model and endure the lack of personalization. This problem is caused by the mismatch of training data distribution and test data distribution, as assumed in our proposed set-up, and is solvable with our proposed method HPFL by exploiting personalized plug-ins from other clients.
- **Case 2: A car with a personalized automated driving system (ADS) has driven out of the previous city it used to be.** It requires to personalize on geometric data from the present city it is now in for improving the performance of the ADS in this new city. Classic GFL and PFL in this situation leave the ADS no option but to collect the geometric data and personalize on it after the collection completes, and accept the temporary performance loss using the previous personalized model before finishing the new personalization, since the distribution of test data has greatly changed. It's another example where the discrepancy between training data (geometric data from the previous city) and test data (geometric data from the present city) threatens the availability of FL systems. While with our proposed method designed to solve the problem, the ADS can attempt to access the plug-ins from car owners living in the present city.
- **Case 3: Imagine a person is traveling from a high latitude area to an equatorial region, and the recommender system on their phone is supported by federated learning.** If the recommender system uses the personalized model trained when in the high latitude area, it will continue to prompt thick down jackets for the person, which is clearly an unexpected and unreasonable recommendation. With our method, one can get the same recommendation as the local people with plug-ins on their phones without time to fine-tune the model again. 

Following your suggestion, we have added the examples listed above to the ***Appendix F.1 of our revision*** due to the page limit.

Reviewer comment (Evaluation):
Although the presented evaluation framework is good for the presented problem, a more exhaustive framework could be formulated to rigorously test the approach, e.g., one in which train client $m$ is given test data $n$, and we check if it indeed downloads plug-in $n$.

Author response:
**Authors' Response:** For our proposed methods HPFL, because all clients have the same selection method $C$. So the selection method is only determined by the test data $n$ and not affected by which client carries out the inference as stated in ***Equation 15 in origin main text*** $C_i(\xi_m,i)=argmax_n \space g\left (\xi_n, \xi_m\right)=C(\xi_m)$ in ***Appendix C.3***. In this situation, the evaluation method the reviewer have suggested is formularized as following equation: $\text{Accuracy} \left(\Omega, \theta _1, \ldots, \theta _M\right)=\frac{1}{M} \sum _{i=1}^M \sum _{m=1}^M p _m \mathbb{E} _{\xi _m \sim \mathcal{D} _m} \mathcal{T}\left(f\left(\theta _{C _i(\xi _m,i)}, \xi _m\right), \xi _m\right)$ is the plug-in n downloaded by client m; and the simplified evaluation method used in our implementation is accuracy  of  PM  selected  by  client  i  on  its  own  data: $\text{Accuracy} _{simplified}\left(\Omega, \theta _1, \ldots, \theta _M\right)=\sum _{i=1}^M p _i \mathbb{E} _{\xi _i \sim \mathcal{D} _i}\mathcal{T}(f(\theta _{C(\xi _i)} ^{pfl}, \xi _j), \xi _j)$. We have strictly proven the equivalence of our simplified evaluation framework and the origin evaluation metrics in ***Appendix C.3***.

Reviewer comment (Novelty):
The paper has low novelty. Federated learning and fine-tuning the sub-modules are really common and baseline personalized FL methods.

Author response:
We sincerely thank the reviewer for taking the time to review. We appreciate that you find our proposed problem novel, our paper well-written, and the experiments comprehensive. Based on your valuable comments, we provide detailed feedbacks as follows. 

**Question: novelty**
>The paper has low novelty. Federated learning and fine-tuning the sub-modules are really common and baseline personalized FL methods.

We not only propose a novel problem setting but also introduce a novel mechanism to solve this problem. We propose a novel problem named GFL-PM and reduce the problem into Selective FL (SFL), and finally solve the SFL problem with our proposed method HPFL, which introduces a novel plug-in mechanism. Our difference with other previous personalized FL methods, which carry out GFL and fine-tuning the sub-modules, is listed below: 
|            |                       Personalize part                       |                     Personalized method                      | Ability to adapt to test data | Ability to utilize personalized knowledge from other clients |
| :--------: | :----------------------------------------------------------: | :----------------------------------------------------------: | :---------------------------: | :----------------------------------------------------------: |
|   FedPer [1]   |                         last layers                          |                         fine-tuning                          |             **×**             |                            **×**                             |
| LG-FedAVG [2]  |                       previous layers                        |                  only aggregate last layers                  |             **×**             |                            **×**                             |
|   FedRep [3]   |                          last layer                          |                only aggregate previous layers                |             **×**             |                            **×**                             |
|   FedRoD [4]  |                          last layer                          |                only aggregate previous layers                |             **×**             |                            **×**                             |
| PerFedMask [6] |                 determine by calculated mask                 |                         fine-tuning                          |             **×**             |                            **×**                             |
|   FedBN [7]   |                          BN layers                           |               only aggregate layers except BNs               |             **×**             |                            **×**                             |
|   FedTHE [5]  |                          last layer                          |                         fine-tuning                          |             **√**             |                            **×**                             |
|    HPFL    | not limit (experiments are carried out with last layers, however, HPFL is compatible with all kinds of PFL) | not limit (experiments are carried out with fine-tuning, however, HPFL is compatible with all kinds of PFL) |             **√**             |                            **√**                             |

From the table above, we can clearly see most PFL algorithms lack the ability to adapt to test data. FedTHE utilizes an interpolation of global prediction and personalized prediction, whose coefficient is trained on test data. While our proposed method HPFL obtains similar ability with selection method and plug-ins from other clients, which enables the exploitation of personalized knowledge from other clients simultaneously. Therefore, the novelty of HPFL does not lie in how the backbone is trained or how the plug-ins are personalized, which is the key point of most previous methods, but in its overall framework and plug-ins selection mechanism.