PeFLL: Personalized Federated Learning by Learning to Learn

Thank you for your review, we answer your specific concerns below.

All experiments are conducted using LeNet-style models, which are relatively outdated. I would like to see the performance on ResNet-style models.

Please see our reply to the first question of reviewer Couf about the architecture and the size of the network.

In Table 2, the performance of "Local" is missing.

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.

Moreover, 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?

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.

That is why I think experiments covering these settings are needed to conclude the practical benefits of PeFLL.

We disagree that PeFLL should run with bigger target networks to be practical. As discussed above, small models are the default for personalized FL. For example, if in each round we had to send a ResNet50 with 24 million parameters to the clients, clients would suffer from large data transfer volumes already after a few rounds of participation.

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.

Thank you for your review, we address your comments below.

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

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.

[1] A. Shamsian, A. Navon, E. Fetaya, and G. Chechik. Personalized federated learning using hypernetworks. ICML, 2021.

[2] C. T. Dinh, N. H. Tran, and T. D. Nguyen. Personalized federated learning with moreau envelopes. NeurIPS, 2020.

[3] L. Collins, H. Hassani, A. Mokhtari, and S. Shakkottai. Exploiting shared representations for personalized federated learning. ICML, 2021.

[4] A. Fallah, A. Mokhtari, and A. E. Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In NeurIPS, 2020.

[5] A. Hard, K. Rao, R. Mathews, S. Ramaswamy, F. Beaufays, S. Augenstein, H. Eichner, C. Kiddon, and D. Ramage. Federated learning for mobile keyboard prediction. arXiv:1811.03604, 2018.

[6] X. Zheng, Y. Zhang, G. Andrew, C. A. Choquette-Choo, P. Kairouz, H. B. McMahan, J. Rosenstock, and Y. Zhang. "Federated Learning of Gboard Language Models with Differential Privacy." arXiv:2305.18465, 2023.

[7] B. Knyazev, M. Drozdzal, G. W. Taylor, and R. Soriano, A.. Parameter prediction for unseen deep architectures. NeurIPS, 2021.

[8] D. Ha, A. M. Dai, and Q. V. Le. Hypernetworks. ICLR, 2017.

Thank you for addressing my comments.

Note that I said “more modern” and not “larger”. Small models are indeed very relevant in FL. I totally agree with that, even though improvement in communication efficiency (either in terms of connection speed or algorithmic improvements) can allow the use for reasonably bigger models. However, I was actually more interested in seeing performance comparison on “modern” architectures regardless of size, particularly ones employing residual connections, attention, etc. This might be more relevant for cross-silo FL as models can in fact be a little big.

Thanks for sharing some details about training.I understand from you that bigger models can be generated efficiently by hyper-networks, which is good. It would be interesting then to see how that holds up in your experiments. The server could indeed have great computational capabilities, but at least you should demonstrate that scalability would not be an issue (e.g. by showing that compute and/or memory scales linearly with model size, which is doable based on what you said).

How does PeFLL fare in non-label skew settings? E.g. I would recommend the Shakespeare dataset

Note that FEMNIST is not a purely a label skew dataset, as different clients correspond to different writers. Thank you for suggesting the Shakespeare dataset, we will try to get experiments running and finished before the end of the discussion period.

During training, PeFFL is constrained to "FedSGD", i.e. a client cannot amortize communication with additional computation. Instead, forward-backward pass requires communication of various elements: Embedding, client-specific model and the gradients corresponding to both.

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.

Introduction of another baseline, in which a larger, more flexible model is trained through FedAdam (https://arxiv.org/abs/2003.00295) with properly tuned hyper-parameters. The hypothesis I would like to see tested is if a more flexible model is able to generalize across all train & test clients without requiring personalization. For a given communication budget, should I choose PeFFL or can I fall back to a larger FedAvg-trained model?

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?

Another interesting experiment (though not required), would be to see a combination of classical FedAvg and PeFFL. Consider applying PeFFL only to the last layer(s) of a model, while training the rest through standard FedAvg. For object classification, it is usually the last layers that require personalization, while the initial layers can share jointly relevant features. Such a setup would allow the embedding network to be a separate head of the FedAvg-trained network.

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!

Thank you for your detailed feedback, we address your concerns below.

I believe the discussion of related works should include the topic of Split-Learning

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

The authors only consider the label-skew setting of object-classification computer vision problems C10/C100 and FEMNIST. Specifically for C10/C100, the authors assign two labels per client. The provided code (mask_absent is True by default) suggests that their method consequently reduces to a binary classification problem at test time. Can you confirm this choice of evaluation is consistent with all baselines?

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.

It seems that for all (train & test) clients, you assume the presence of a labeled training set. [...] I would expect a "test-client" to not have any labelled data available.

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.

Can you clarify your evaluation procedure for me? Do "test-clients" use a different training-set for embedding computation? I would expect those to be identical. For the training-clients, I would be curious to understand if you use the labelled training data for embedding computation, or throw those away for the validation-set in Table 1a.

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.

You claim that the client descriptor network is not susceptible to overfitting. I would ask you to show evidence for that claim.

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.

We have now finished running the results for the experiments where the test clients do not possess labeled data. Specifically what we do is: the PeFLL embedding network is trained on the training clients using only the image data and no labels (embed_y = False in the code). The test clients then obtain their personalized model using only their unlabeled data and evaluate this model. Of course since they do not have labeled data there is no class masking (mask_absent = False in the code). We report below the results on the test clients. Note that the only baseline that is capable of addressing this setting is FedAvg, since all personalized FL methods require labeled data to be available at the test clients in order to obtain a personalized model.