Paper: PeFLL: Personalized Federated Learning by Learning to Learn (ICLR-2024)

Reviewer comment (Presentation):
In Table 2, the performance of "Local" is missing.

Author response:
We can include this. It is the same procedure as in Table 1 operating on the same kind of data, so the values for local in Table 2 will not differ substantially from those in Table 1.

Reviewer comment (Experiments):
During experiments, what is the number of local updates for the baselines like Per-FedAvg and pFedMe? How do their performance improve as the number of local updates increases, and can they achieve similar performance compared to the proposed PeFLL?

Author response:
We assume by ‘local updates’ here you mean the number update steps used to personalize/finetune to a test client to obtain the results quoted in Table 2. The personalization on each test client was simply run locally until convergence, which does not usually require many steps due to the small amount of data on each client. The reported result is based on the best validation step. Therefore, we do not see that the results for Per-FedAvg/pFedMe could improve by even more steps.

Reviewer comment (Evaluation):
The comparison to baselines is somewhat unfair, due to a lack of measuring communication overhead.

Author response:
Thank you for clarifying your concern. Indeed, the communication cost of PeFLL per round is approximately twice that of FedAvg ($\nu_v$ and $\theta$ are of comparable or same size, the size of $v$ is just a few dozen bytes). 

As requested we have now included training curves that compare accuracy obtained for a given communication budget (measured in GB of transferred information). As we can see from the curves for any given communication budget PeFLL performs substantially better than FedAvg.

We do believe the comparison with respect to rounds also has merits, though, as this better reflects other aspects of federated learning training, such as recruiting clients to participate, client computation and client stragglers, which can be bottlenecks in practice.

Reviewer comment (Novelty):
The discussion of related works should include the topic of Split-Learning, a version of FL where parts of the forward-pass are performed across different devices (e.g. client and server). PeFFL is an example of split-learning under that definition.

Author response:
Thank you for bringing this to our attention. We agree that it is both relevant and highly interesting. We will add this topic to our related work section and also include citations to [1,2]. Please let us know if you have other suggested related work. However, we would like to emphasize that we are not doing split learning exactly as described because we are not communicating gradients, but rather model differences after multiple updates (exactly as FedAvg does). We discuss this in more detail below. 

[1] Otkrist Gupta, Ramesh Raskar. Distributed learning of deep neural network over multiple agents, Journal of Network and Computer Applications, 2018

[2] Chandra Thapa, M.A.P. Chamikara, Seyit Camtepe, Lichao Sun. SplitFed: When Federated Learning Meets Split Learning, AAAI, 2022

Reviewer comment (Experiments):
The authors only consider the label-skew setting of object-classification computer vision problems C10/C100 and Femnist, assigning two labels per client for C10/C100. The provided code ("mask_absent" is True by default) suggests that their method consequently reduces to a binary classification problem at test time. Clarification is needed on whether this choice of evaluation is consistent with all baselines. The evaluation in lines 374 & 375 of train.py concerns the test-clients, leading to evaluation in line 20, with default args embed_split=True suggesting that the training-set is used to compute embeddings and to test for presence of classes on that client (lines 58-76). It seems that for all (train & test) clients, a labelled training set is assumed, which allows for masking logits. Lines 24-38 suggest training & test clients are treated identically. The expectation is that a "test-client" should not have any labelled data available, meaning that logit-masking should not be allowed. Clarification is requested on whether "test-clients" use a different training-set for embedding computation, and whether these are identical. For training-clients, clarification is also sought on whether labelled training data is used for embedding computation, or excluded in favour of the validation-set in Table 1a. Additionally, evidence is requested for the claim that the client descriptor network is not susceptible to overfitting.

Author response:
First, what you describe is the setting for CIFAR10. For CIFAR100, each client actually has 10 classes, and for FEMNIST they can have all classes. 
Second, for all baselines the clients implicitly or explicitly use information about their subset of relevant classes. For example, kNN-Per performs kNN, so it can never predict a class not present in the client’s labeled training set. All other personalized methods train or finetune a classifier locally on a labeled training set, even for the test clients, so classes not present in their training data are quickly suppressed. 

Indeed, in our setting “test clients” differ from “train clients” only by the fact that they appear later, after the server is done with training. Maybe a better phrasing would be “clients seen during the server’s training phase” and “clients seen after the server’s training phase is finished”. 

In the supervised setting we study, labels are certainly needed during the training phase, and therefore it makes sense to assume that clients occurring later also have them. Indeed, none of the baselines to which we compare would be able to personalize to a test client that has no labeled data. 

However, we find the suggestion of learning personalized models for clients without labels for their data also quite exciting. Specifically, all baselines we report, except plain FedAvg, would fail to produce any classifiers for new clients in that setting. PeFLL, however, could work simply by using an embedding network that operates on the input data without labels. We are now running experiments in this setting (and without masking). We will report it as soon as it finishes, and we thank the reviewer for the suggestion.

Train and test clients behave identically in that they receive the embedding network, compute their client descriptor from a batch of their data, send the descriptor back to the server and receive a personalized model. The only difference is that “train clients” subsequently participate in the server’s training step by performing gradient computations and sending parameter updates to the server.

Please see the paragraph in the experiment section **Analysis of Client Descriptors** as we believe this answers your question. In this experiment we explicitly test if the client descriptor (embedding) network is able to embed the test clients into descriptor space in a way that is consistent with the embeddings of the train clients. As the experiment shows (Figure 3), the embedding network does indeed learn to do this. If the embedding network simply overfit to the train clients then the test client descriptors would be meaningless.

Reviewer comment (Experiments):
For a fair comparison of PeFFL against baselines, it is necessary to introduce the communication overhead of the different methods. Although the authors argue that "It has low latency and communication costs", this should be clarified to be true only at test-time (and even then, it is higher than alternatives that function offline, such as FedAvg). During training, PeFFL is constrained to "FedSGD", meaning a client cannot amortize communication with additional computation. A forward-backward pass requires communication of various elements: embeddings, client-specific model, and corresponding gradients. Two analyses are requested:  
1) GB of communication on the x-axis and test-client accuracy on the y-axis, to allow drawing iso-lines and understanding the impact of PeFFL's high-frequency communication.  
2) Introduction of another baseline where a larger, more flexible model is trained through FedAdam (https://arxiv.org/abs/2003.00295) with properly tuned hyper-parameters. The hypothesis to test is whether such a model can generalize across all train & test clients without requiring personalization, to compare with PeFFL for a given communication budget.

Author response:
This is not correct and we believe this could be leading to a **fundamental** misunderstanding of the communication costs of our method. We are **not** constrained to FedSGD. During training each client runs k local updates to their personalized model (line 10 of Algorithm 2). In our experiments k=50, while FedSGD would be k=1. The term that gets sent to the server to be ‘backpropped’ through the network is the difference between the initial personalized model and the final one **after** these 50 steps, (line 11 of Algorithm 2). We also refer you to our convergence results in Theorem 3.1, which shows that the number of local steps has the same effect as it has in FedAvg, and **proves that our procedure has equivalent convergence to FedAvg with k local steps per round** and not FedSGD. Since our convergence bound is in the same order of FedAvg, the number of global steps required for the convergence is the same as FedAvg, and we do not require more communication rounds than FedAvg. 

As we mentioned for the previous question, we have the same number of communication rounds as FedAvg for convergence, just in each round we send two (small) networks instead of one. Also, our optimization is in fact already done using Adam. 

Regarding a larger model: we do agree that a large enough global model might also achieve high accuracy in the setting of our experiments. Note, however, that such a baseline would not be a fair comparison in terms of the computational load on client devices (which often have very limited compute and memory)
In light of this fact and given our reply to your prior concern about the PeFLL’s communication cost, do you still think such a baseline is necessary?

Reviewer comment (Experiments):
An additional interesting (though not required) experiment would be to combine classical FedAvg and PeFFL, applying PeFFL only to the last layer(s) of a model while training the rest through standard FedAvg. For object classification, the last layers usually require personalization, while the initial layers can capture jointly relevant features. This setup would allow the embedding network to be a separate head of the FedAvg-trained network.

Author response:
We agree that this could be an interesting experiment, in particular as some early layers of the network could double as the embedding network. PeFFL would allow for such a setting without major modifications. Thank you for the suggestion!

Reviewer comment (Experiments):
I think that the experiments are not sufficient to conclude that PeFLL is superior. You have to demonstrate its improvement on architectures more modern than LeNet as well. It would also be great to report performance on federated datasets other than FEMNIST (e.g. other LEAF datasets).

Author response:
> I think that the experiments are not sufficient to conclude that PeFLL is superior. You have to demonstrate its improvement on architectures more modern than LeNet as well.

Note that the model we are using is not the classic “LeNet-5” with tanh-activation etc. that you might have in mind. It is a modern ConvNet inspired by the LeNet architecture but bigger and with ReLU activations. We use this architecture to allow a fair comparison of PeFLL to prior work on personalized FL, such as the baselines in [1,2,3,4].

Moreover, we do not agree that showing superior performance on small models is irrelevant. In Federated Learning it is very common in practice to use small models due to the fact that we need to keep communication costs low and the client devices themselves are often too limited in terms of compute to be able to train huge models. For instance Google’s Gboard model, which they use in production, is a single layer LSTM, see [5,6].

Reviewer comment (Experiments):
The computational efficiency of the model is questionable given that the hypernetwork is at least a 100 times bigger than the client's model, and this is only the case for these particular experiments. It might require a wider net to generate larger personalized models. In other words, LeNet on MNIST and CIFAR-10/100 might only need 100 "ranks" to perform well, whereas on more complicated models and datasets, a wider net would be needed for good performance, which makes scalability an issue. That is why I think experiments covering these settings are needed to conclude the practical benefits of PeFLL.

Author response:
> The computational efficiency of the model is questionable given that the hypernetwork is at least a 100 times bigger than the client's model, and this is only the case for these particular experiments. It might require a wider net to generate larger personalized models. In other words, LeNet on MNIST and CIFAR-10/100 might only need 100 "ranks" to perform well, whereas on more complicated models and datasets, a wider net would be needed for good performance, which makes scalability an issue.

If we understand you correctly, you are asking generally about the feasibility of using hypernetworks to generate large models. This is absolutely the case: there are a number of papers that use hypernetworks to generate very large models. For instance, [7] uses a hypernetwork to generate the parameters of a ResNet-50: “we can predict all 24 million parameters of ResNet-50 in less than a second either on a GPU or CPU” and [8] also generates ResNets using a hypernetwork. Note that -by design- in our method the hypernetwork is located on the server (as opposed to on client devices), so we are not so limited in terms of computational resources to run the hypernetwork.

Reviewer comment (Evaluation):
Can you share some details about training time, memory consumption, and computational efficiency of PeFLL vs. regular models?

Author response:
> Can you share some details about training time, memory consumption, and computational efficiency of PeFLL vs. regular models?

Can you clarify this question? Are you referring to the overall runtime, memory and compute of the method in our simulations? Our method runs in around 7 hours, while baselines took between 5 and 10 hours to run. The overall memory consumption is approximately 1 GB. However, we do not think that these statistics are very reflective of what will happen in practice in FL given that runtime is most greatly influenced by how much compute a client device has, scheduling overhead, stragglers etc. During training the amount of compute that is run on each client is roughly the same for all methods, 50 local updates to the personalized model. PeFLL additionally has a single forward and backward pass on the (small) embedding network during training. When running on test clients PeFLL has much lower compute than all other personalized FL methods since it only requires a forward pass on the (small) embedding network while other methods need to run finetuning locally. PeFLL also has additional compute that happens on the server, but this is minimal and consists of a single forward and backward pass through the hypernetwork for each global training round. This is negligible given the large compute power of the server, and the fact that it is only a single pass.