Federated Learning (FL) is an emerging paradigm that allows a model to be trained across a number of participants without sharing data. Recent works have begun to consider the effects of using pre-trained models as an initialization point for existing FL algorithms; however, these approaches ignore the vast body of efficient transfer learning literature from the centralized learning setting. Here we revisit the problem of FL from a pre-trained model considered in prior work and expand it to a set of computer vision transfer learning problems. We first observe that simply fitting a linear classification head can be efficient and effective in many cases. We then show that in the FL setting, fitting a classifier using the Nearest Class Means (NCM) can be done exactly and orders of magnitude more efficiently than existing proposals, while obtaining strong performance. Finally, we demonstrate that using a two-phase approach of obtaining the classifier and then fine-tuning the model can yield rapid convergence and improved generalization in the federated setting. We demonstrate the potential our method has to reduce communication and compute costs while achieving better model performance.
Paper
Similar papers
Peer review
Summary
The work proposes to address the training of the classifier (last layer of the network) separately in the setup of federated learning with a pretrained weights initialization. In particular, the suggestion is to use a nearest class means (NCM) approach, which saves a lot of computation time. The proposed approach is also supposed to be helpful in heterogeneous setup and thus evaluated on datasets separated into local ones by Dirichlet distribution. The paper provides an empirical evaluation of the linear probing and NCM approaches for training classifier, compared to standard federated learning. The conclusion made from the empirical results is that proposed method saves communication and leads to better performance.
Strengths
The proposal to use only the last layer adaptation when the weights are pretrained is a viable and interesting idea for federated setup. It seems to mitigate heterogeneity and save computational time (with NCM) and communication (when no further finetuning is performed).
Weaknesses
1) The explanation of the approach can be improved---especially the mixed one (is the final tuning happening for all the weights or without head?). 2) Heterogeneity is presented as one of the main reasons to introduce a new approach, but the method is not compared to any specific federated algorithms devised for it. 3) The evaluation is trying to highlight both LP and NCM, while overall they do have rather different approaches; this makes it unclear what exactly is the main message of the paper: to train only the last layer or to train it in the particular way? Moreover, only the combination with further tuning results in best performance. Minor - table3 is never referenced - figure4 is referenced before figure3 - figure5 is called table5
Questions
1) It seems that only head tuning + full tuning leads to the best results? But then it is not communication cheap? 2) It is claimed that NCM is more secure than standard FedAvg, but it is sharing the mean of the features for classes for each local learner. I think this can be a revealing information about local weights. Can you comment on this? 3) How the approach relates to layerwise personalization federated algorithms?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
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
Evaluation for heterogeneity is done only with very particular non-iid with respect to labels, but it is not highlighted.
Summary
This paper studies strategies to fine-tune models in Federated Learning, starting from a pre-trained model. It has been shown in the literature that pre-training is beneficial in FL, that it improves convergence speed and robustness to heterogeneity. This paper shows that existent fine-tuning pipeline are not optimal. Further, this paper proposes a two stage head-tuning + fine-tuning technique, where the first head-tuning stage is compute and communication efficient. The authors show that their proposed two state approach are better than Linear Probing and full Fine-Tuning in terms of convergence to accuracy and robustness to hyperparameters.
Strengths
1. Experiments are done with rigor and described in detail the paper. 2. The authors ran into a discrepancy in terms of experiment results on one dataset reported in (Nguyen, 2023). They provided a concise reason and experiments to explain the source of discrepancy, which is due to the image size in the training data.
Weaknesses
1. I would like to see how the proposed method compared to other simple two-stage fine-tuning methods, for example, running LP for a few epochs then perform full FT. This simple modification is not compared in the experiments. 2. high-dimensional classification: the output dimension on all problems studied in the experiments are not too big. In some high-dimensional classification problem, e.g. language model, some classes only have a few examples. I would like to understand how they affect the centroid-based initialization and the applicability of FedNCM.
Questions
In Cars and Eurostat dataset, FedNCM + FT starts with the highest initialization accuracy, but the accuracy during training decreases first. Non-monotonicity are expected due to the stochastic nature of optimizer, but looking at Figure 1, the non-monotonicity here are more than what is produced by stochasticity. In particularly Eurostat, the accuracy decreased for the first 1000 rounds to half of the starting accuracy before it improves. I would like to understand more about this behavior. In practice when you are running FL in production where communications are costly, if one observes the accuracy continues to drop for the first 1000 rounds, the FL job will be killed and some bug will be suspected somewhere.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
4 excellent
Presentation
4 excellent
Contribution
3 good
Limitations
N/A
Summary
This paper proposes to use nearest class means (NCM) with pre-trained models in federated learning, coined as FedNCM. Experiments show that FedNCM is effective in terms of convergence and communication cost, due to the application of pre-trained models and light-weighted last layers.
Strengths
1. Using pre-trained models in federated learning is a new and interesting direction, especially with the popularity of foundation models. This paper explores pre-trained models from a federated learning perspective, which shows the effectiveness of saving communication. 2. It provides code and experimental results on several vision datasets, including CIFAR-10, Flowers102, Stanford Cars, CUB and EuroSAT-Sub. 3. The experimental results show the improvement over Random baseline.
Weaknesses
1. The algorithm seems to be a mere combination of NCM and FedAvg, which is not novel enough. This is not vital though as long as the performance is great. 2. However, the application of pre-trained models in FL has already been proposed in Nguyen et al. ICLR 2023. 3. The experimental comparison is not considering Federated Learning baselines, like FedProx, SCAFFOLD, FedYogi, LG-FedAvg [1], FedPer [2], etc. 4. The experiments only include small scale datasets, in total less than 50,000. 5. There is not much theory on the communication cost, or convergence rate analysis. 6. Minor: Line 201 FedAVG -> FedAvg; Line 234: imagenet -> ImageNet; [1] Liang, Paul Pu, et al. "Think locally, act globally: Federated learning with local and global representations." arXiv preprint arXiv:2001.01523 (2020). [2] Arivazhagan, Manoj Ghuhan, et al. "Federated learning with personalization layers." arXiv preprint arXiv:1912.00818 (2019).
Questions
1. Line 165: Could you explain NTK or add a reference? 2. Could you explain more clearly the notations: FedNCM, FedNCM + FT, FedNCM + FT Pretrain, LP Pretrain, Random? Otherwise, the readers can only guess from these words. 3. Could you add more FL baselines as well as recent works on FL + pretrained model finetuning? For example, why is NCM preferable compared to softmax classifier? 4. Could you add more explanation about Fig 1, such as the different behavior on different datasets?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
2 fair
Contribution
2 fair
Limitations
The limitation discussion is a bit limited and the authors just spend one sentence on this. Suggestions include the negative societal impact of using pretrained models and the lack of theoretical analysis and large-scale datasets, etc.
Summary
This paper studies the effect of tuning the linear head on federated feature learning. They show that well-tuning the linear head will make the base feature extractor converge closer to the optimal solution while reducing the total communication cost in FL. They evaluate their method on multiple FL datasets and show promising results.
Strengths
Investigating the effect of a linear head on the final FL model tuning is interesting.
Weaknesses
One thing is missing. When learning the linear head in the first stage, it is possible to get the closed-form solution of V. Why not do that? This will give a smaller upper bound for the right-hand side of Eq (2). There are some works to improve the compute/communication efficient FL, such as https://arxiv.org/pdf/2206.08671.pdf (not limited to this). Some comparisons should be conducted.
Questions
In Eq (2), the triangular inequality may be valid for a single client. Will this still be valid considering the FedAvg/FedAdam aggregation for w? Can you put a theoretical convergence analysis of the two-stage FedNCM to show its advantage compared with the one without first-stage head tuning?
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
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
2 fair
Contribution
2 fair
Limitations
Please see above.
Thank you for the replies! Taking into account your W2 answer, can you please once again clarify the motivation for the new method for fighting heterogeneity: if in the setup with a pretrained model all the federated approaches are already good, why introduce another method? Is the main benefit then in saving communication in the initial phase and speeding up convergence?
Clarifications
Thank you for taking the time to read and reply to our rebuttal. First of all we want to clarify a potential misunderstanding: Nguyen et al observed that in the pre-trained setting methods not designed for heterogeneity (e.g. FedAdam) can closely match and even outperform those that are designed for it, and also that the gap of iid vs non-iid is reduced (but importantly not closed). Despite performing well compared to other methods and outperforming random initialization, this does not mean that performance, communication efficiency, convergence speed, nor even robustness to heterogeneity is perfect, for example the performance in CIFAR-10 still lags behind centralized training. In our work our main message is about the importance of HeadTuning methods for achieving drastic improvements in multiple factors: communication, computation, convergence improvements, and often overall better performance. We note this can be a more critical factor than the choice of FL base algorithm. Robustness to heterogeneity however is another important observed side effect of our proposal. Firstly, FedNCM is completely robust to any of the typical non-iid data shifts (as opposed to LP or FT), and FedNCM+FT due to starting with a fixed head is more robust to heterogeneity than naive applications of FL algorithms. Note that **in our Figure 2 we consider more extreme non-iid case than in Nguyen et al** with $\alpha$=0.01, here performance of base FL algorithm starts to degrade substantially, while with our two phase method the degradation is much more graceful
Thank you for the reply. I am wondering if you somehow connect your research to model soups (Wortsman, Mitchell, et al. "Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time." International Conference on Machine Learning. PMLR, 2022.)? Also, your comment on CIFAR10 reminded me of the known problem with averaging networks with batch normalization (e.g., Ainsworth, Samuel K., Jonathan Hayase, and Siddhartha Srinivasa. "Git re-basin: Merging models modulo permutation symmetries." arXiv preprint arXiv:2209.04836 (2022).) - maybe it is worth taking it into account.
Follow-up
Thank you for the reference to Model soup, we will add it. Notably Wortsman et. al. focus on fine-tuning on a *single* dataset, observing that combining models finetuned with different hyperparameters can be beneficial. This is distinct from our setting where we simultaneously fine-tune models with federated constraints on as many as 100 heterogeneous datasets. Their method also requires a centralized validation set. It can however be a potential future work to combine multiple models fine-tuned by our rapid fine-tuning procedure with different FL hyperparameters. However this would come at a higher communication and computational cost, although may be practical in some settings. We leave this exploration for future work. Regarding the connection between the issue we observed on L233 with respect to cifar-10 and normalization in git-rebasin. We want to emphasize the issue we observed is not unique to the federated learning or model averaging situation and can be observed in a standard centralized setting. Models trained on Imagenet are learning to extract features with input at 224x224 resolution, thus feeding heavily downsampled images for feature extraction can lead to distortion in the natural image hierarchy of CNNs, leading to large gaps in the linear model (or in NCM) performance when using the features. We note our image experiments do not use normalization layers as squeezenet the pre-trained model we focus on (inspired by Nguyen et al) does not have any. Our new NLP experiments with Distillbert use layernorm as this is part of the pre-trained model. We hope that we have addressed the concerns of the reviewer.
Thank you for the discussion. I am therefore raising my score.
Thank you for the response
Thank you for your comprehensive response. I understand now the difference compared to Nguyen et al ICLR 2023. Overall I think fine-tuning is more important than NCM, from the experimental results. For example, on CIFAR-10 FedNCM + FT only improves FT by less than 2%. On other datasets like Flowers, however, the improvement of FedNCM on FT is more significant. I think it is interesting to understand when and how FedNCM helps, although I appreciate the extensive experiments done in this work. Additional question: why is FedNCM + FT not implemented for CIFAR-10 $\times$ 32?
Communication, Compute cost and Robustness to Heterogeneity are greatly improved by FedNCM+FT
Thank you for the follow-up. We want to emphasize that our FedNCM+FT does not just improve final accuracy (as noted quite substantially in some datasets) but also critically the convergence speed over naive FT which leads to large gains in communication and compute cost particularly under a budget. In addition we obtain larger performance gains in highly heterogeneous settings (see discussion with reviewer 1Br4). The NCM is critical for obtaining these large advantages since it is a negligible cost (and robust to heterogeneity) initial step. Communication, compute cost, heterogeneity are all essential performance factors in FL. Note that in Table 1 for fairness of comparison we run the FT and FedNCM+FT both for the same amount of rounds even if FedNCM+FT has converged, thus the total compute cost is not representative of the drastic advantages in communication that are seen when observing Figure 1. We refer the reviewer to consider the large gains observed in FT vs FedNCM+FT in Figure 1 and their practical consequences in FL. We will clarify this further in the text. Regarding CIFAR-10 x 32, the results for that shown in Tab 1 are taken from Nguyen et al (as clarified in L238-242) to directly show that when doing the native model resolution the FT acc is much higher and the gap of FT and LP is smaller (see L238-241). Thus in all our experiments we use CIFAR-10 upsampled to 224x224 (as done in all other datasets). We will more clearly indicate the CIFAR-10 x 32 results are added from Nguyen et al. as reference and add the FedNCM+FT result for this case which we have run now giving 67.9% (vs 63% for naive FT)
Sorry I don't understand the argument. As far as I understand the clients have to compute all the feature embeddings on the local training datasets, which can involve heavy computation if the training set is large. Also, I think the advantage of NCM in FL is due to the class imbalance in the datasets (in which case some other methods can be compared like Asymmetric loss e.g., https://arxiv.org/pdf/2106.03110.pdf). The authors should make this point more clearly. Given the new experiments I will update my score to 5. I still think this paper, though interesting, may need another round of revision.
FedNCM requires only one forward through all data vs many rounds and multiple local epochs of forward backward (FT)
Thanks for considering our rebuttal and increasing the score. We are happy to clarify the efficiency of FedNCM. We want to emphasize that FedNCM (without FT) is negligible in communication and compute cost as compared to FT (the prior work), while already obtaining good results (that can be further improved by a 2 phase procedure). The large gap is because FedNCM requires just one forward through all data and sends only one time the parameters, while FT does multiple local epochs and rounds of computation and communication. We illustrate this with an example below: Consider a standard FL FT based training (as done e.g. in Nguyen et al), with 5 local epochs of training and 500 rounds. Let’s also denote F as the number of FLOPs for a forward pass on all data and for simplicity consider 2*F the number of FLOPs for a backward pass on all data. **Compute (NCM vs FT)**: Each round consists of all clients forwarding and backwarding 5 times on all data. The total compute cost of FT is 3*(local epochs)*(rounds)* F = 7500*F **Thus NCM will cost only F Flops while FT will cost 7500 times more**. Note that from Table 2 NCM alone can already sometimes exceed FT (e.g. flowers). To improve this further the excellent NCM starting point then leads to rapid convergence if followed by an FT phase (leading to savings in comm cost over FT due to the rapid convergence) **Communication (NCM vs FT)**: If we compare simply NCM to FT; NCM requires 1 round of sending weights to the clients, clients return only their centroid averages which is typically negligible cost. On the other hand FT involves 500 rounds of sending the entire model to the clients and the clients sending the entire model back to the server. The reviewer is right that with respect to heterogeneity advantage of NCM (only one of multiple advantages!) this is best seen when label distributions vary amongst clients, the most common heterogeneity considered in the FL literature (e.g. Nguyen et al, Legate et al, McMahan et al). We will further emphasize this in Sec 3.2.
Thank you for future clarification. I understand FT takes more computation, but in the paper NCM works well only on top of FT. Therefore we could only say that the overhead is small compared to FT.
FedNCM and FedNCM+FT have large gains over FT given a compute and communication budget
- FedNCM does already exceed or match FT for Flowers and in the new experiments on AG News (see attached pdf in top level message). In other cases it would be the best performer if the application has a constrained compute/comm budget - Our two phase (FedNCM +FT) leads to large gains in compute and communication without ever losing acc over FT. For example observing Figure 1 for CIFAR-10 FedNCM+FT achieves the final accuracy of FT in 60% less rounds, for Flowers in 99.5% less rounds, for CUB in 90% less rounds, for Cars in 75% less rounds, for Eurosat in 50% less rounds
I appreciate the follow-up on the issues by the authors. Here are my comments regarding the author response: 1. response to weakness 1"Comparisons to first stage LP": this concern is sufficiently addressed. 2. response to few shot cases in NLP: the question is partially addressed. The authors provided experiments on AG-news classification tasks, and additionally considered setups with few examples per class. I want to clarify that, what my original suggestions was to consider cases where heterogeneity is high, and some classes has very few examples (extremely unbalance classification with some classes with only a few examples). Overall, my rating for the paper is unchanged.
Decision
Accept (poster)