Guiding The Last Layer in Federated Learning with Pre-Trained Models
Authors: Gwen Legate
Nicolas Bernier
Lucas Caccia
Edouard Oyallon
Eugene Belilovsky
Venue: NeurIPS 2023 poster (2023)
Keywords: federated learning; transfer learning; nearest mean classifier; continual learning;
TL;DR: We show in the federated learning from pre-trained model setting that a simple approach of computing a nearest mean classifier, then finetuning  the model can have large advantages

Abstract: 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 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-stage 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.

Guiding The Last Layer in Federated Learning with Pre-Trained Models
Gwen Legate
Nicolas Bernier Concordia University, Mila Montreal, Canada
Concordia University, Mila
Montreal, Canada
gwendolyne. legate@mila. quebec
Lucas Caccia
Edouard Oyallon
Eugene Belilovsky Concordia University, Mila Montreal, Canada
McGill University, Mila
Sorbonne University, ISIR, CNRS Paris, France
Montreal, Canada
Abstract
Federated Learning (FL) is emerging paradigm that allows model be trained across a number of participants without sharing data. Recent works have begun consider the effects of using pre-trained models 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 set of computer vision transfer learning problems. We first observe that simply fitting linear classification head can be efficient in many cases. We then show that in the FL setting, fitting 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-stage approach of obtaining the classifier and then fine-tuning the model can yield rapid convergence and improved generalization the federated setting. demonstrate the potential our method has to reduce communication and compute costs while achieving better model performance. Code for our experiments is available
Introduction
In recent years increased focus data privacy has attracted significant research interest in Federated Learning (FL). is an approach in which a common global model is trained by aggregating model updates computed across set of distributed edge devices whose data is kept private. We desire to train federated global model capable of equivalent performance as one trained on the same set of centralized data. However, it is worth noting that FedAvg [MeMahan al. 2017], the most commonly used FL. baseline, has been shown to suffer from performance degradation when the distribution of data between clients is not iid [i al. 2020 Acar et al. 2021 Karimireddy et al.) 2020
Transfer learning from pre-trained models that have been trained sufficiently abundant and diverse data is well known to produce state-of-the-art results in tasks related to vision [He et al.] 2019 Girshick al. 2014). Natural Language Processing (NLP) [Radford et al. [019], and other domains. Indeed, pre-training combined with fine tuning to specialize the model for a specific downstream task often leads to better generalization and faster model convergence in the centralized setting (Weiss jet al. 6 Patel et al. [015). FL literature on the other hand, has been largely focused on models
trained from scratch McMaaan et al. 2017 Karimireddy al. 2020
Liet al. 2020 and the impact
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

of heterogeneity on algorithmic convergence. Recently, several studies have been conducted on the effect of pre- pre-training on the performance of standard FL algorithms, see e.g., Chen et al. (2023). Nguyen et al. [023]. Here, it was found that besides improving performance, pre-training can help to close the accuracy gap between a model trained in the federated setting and its non federated counterpart, particularly in the case of non-iid client data.
e.g.,Chen
Prior work on transfer learning in the federated setting has focused on treating the pre-trained model as a stable initialization for classical FL algorithms that adapt all the parameters of the model. Approaches from the transfer learning literature demonstrate that it is often more efficient to adapt only parts of the model such as just the last layers [Kornblihh et al. 2211), affine parameters (Lian al. [022] Yazdanpanah al. 2022), or adapters Houlsby et al. [2019] These approaches frequently yield combinations of better performance, computation time and better prevention against over- weer-fttting
The selection is often a function of architecture, task similarity, and dataset size (Evei et al. 2022 Shysheya et al., [0233 Yazdanpanah et al.) [022]. In many supervised learning tasks studied in the literature such as transfer learning from ImageNet, the representations are powerful and training only
the linear classifier is sufficient for strong performance [Kornblith et al.) 2019].
A key driver of algorithmic development in iL is the reduction of communication cost, having largely motivated the design of the FedAvg algorithm (see e.g., McMahan et al. (0077). Although not studied in the prior works, updating only the linear classifier can be highly efficient in the federated setting when starting from pre-trained model. It can allow for both high performance, limited communication cost (since only the linear layer needs to be transmitted), and potentially rapid convergence due to the stability of training only the final layer. Training the linear classifier in federated setting can lead to classical FL. problems such as client drift if not treated appropriately. An example of this is illustrated in Nguyen et al. [023) Appendix CJ.
We propose a two-stage approach based on first deriving a powerful classification head (HeadTuning stage) and subsequently performing a full fine-tuning of the model (Fine-Ture stage). Such two-stage approaches have been applied in practice and studied theoretically in the transfer learning literature [Kumar et al. [002 Ren et al. 2023b]. They have been shown to give improved performance in both in-distribution and out- -iistribbiititooonn settings [Kumar et al.) [022]. We highlight that the two-stage procedure can lead to many advantages in FL setting: (a) the fine-tuning stage is more stable when data is non-iid heterogeneous models, leading to substantial performance improvement (b) convergence of the fine-tuning stage is rapid (minimizing compute and communication cost)
For the HeadTunnig stage, our work highlights the Nearest Class Mean (NCM), classical alternative to initialize the classification layer which we denote FedNCM in the federated case. FedNCM can be computed exactly and efficiently in the federated setting without violating privacy constraints and we will demonstrate that in many cases of interest, using FedNCM to tune the classification head (without a subsequent fine-tumi mun-ttnng can even outperform approaches considered in prior work with significant communication and computation costs savings.
Our contributions in this work are:
We provide empirical evidence that, for numerous downstream datasets, training only the classifier head proves to be an effective approach in FL settings.
We propose employing a two-stage process consisting of HeadTuning (e.g., via FedNCM LP) followed by fine-tuning, results in faster convergence and higher accuracy without violating FL constraints. We further illustrate that it can address many key desiderata of FL: high accuracy, low communication, low computation, and robustness to high heterogeneity while being easier to tune in terms of hyperparameter selection.
We present FedNCM, a straightforward FL HeadTuning method that significantly reduces communication costs when used, either as stand alone technique, or as stage one (Head- Tuning) in our proposed two stage process which leads to improved accuracy.
2 Related work
Federated Learning The most well known approach in FL the FedAvg algorithm proposed by
McMahan et al. [2017]. In the random initialization setting, convergence of FedAvg and related
algorithms has been widely studied for both iid (Stich] 2019 Wang and Joshi 2018 and non-iid
settings (Karimiredyy et al. 2020 Li et al. 2020 Fallah et al. 2020 Yu al. 2019 A commonly
2

---

cited problem in the literature is the challenge of heterogeneous or non-iid data and a variety of algorithms have been developed tackle this [Li et al., 2020 Hsu et al. 2019 Legate et al., 2023 Karimireddy et al. 2020
Transfer Learning Transfer learning widely used in many domains where data is scarce Girshicc et al. 2014 Alyafeai et al. 2020 Zhuang et al. 20200 Yazdanpanah et al., 2022]. A number of
approaches for transfer learning have been proposed including the most commonly used full model fine-tuning and last layer tuning Korblith et al.] [2019] and some more efficient methods such as selecting features Evei et al. [022]. J. adding aihnn parameters Lian al.. [2022), Yazdanpanah et al.)
[2022], and adapters for transformers Houllsy et al. 2019). Transfer learning and the effects of pre-training in FL. have so far only been explored in limited capacity. In their recent publication,
Nguyen et al. [2023] show that initializing a model with pre-trained weights consistently improves training accuracy and reduces the performance gap between homogeneous and heterogeneous client data distributions. Additionally, in the case where pre- traited data is not readily available, producing synthetic data and training the global model centrally on this has been shown to be beneficial to FL
model performance [Chen et al. 2023
The use of the NCM algorithm artificial intelligence has a long history. Each class is represented as a point in feature space defined by the mean feature vector of its training samples. New samples are classified by computing the distances between them and the class means and selecting the class whose mean is the nearest. In 1990, Ratcliff proposed to use NCM to mitigate catastrophic forgetting in continual learning and since then the use of NCM has been widely adopted and extended by continual learning researchers. This is due to its simplicity and minimal compute requirements to obtain a final classifier when a strong representation has already been learnt. Some of these methods include Rebuffi al. [2017), [i and Hoiem [2017). Davari jet al. 2022 who maintain a memory of exemplars used to compute an NCM classifier. Related to work, recent literature in continual learning that have considered pre-trained models were shown to ignore a simple NCM baseline (Janson et al. 2022) which can outperform many of the more complicated methods proposed. In our work this NCM baseline, denoted FedNCM for the federated setting, demonstrates similar strong performance for FL while serving as a very practical first stage of training in our proposed two-step process.
Nearest Class Means Classifier
Methods
3.1 Background and Notation
In FL, distributed optimization occurs over K clients with each client k E {l, K} having data Xx. Ys that contains n& samples drawn from distribution De. We define the total number of samples across all clients as n = Σ Rik- The data X₁ at each node may be drawn from different distributions and/or may be unbalanced with some clients possessing more training samples than others. The typical objective function for federated optimization is given in Eq. T Koneččý et al. 2016 and aims find the minimizer of the loss over the sum of the client data:
K
Th C(a(s(w, XR). v))
w", v² E arg min W.V
Here we have split the model prediction into f, base parameterized by that produces representa- tions, and g, a task head parameterized by v. In this work we will focus on the case where the task head is a linear model, and the loss function, L represents standard classification or regression loss. The w are derived from pre-trained model and they can be optimized or held fixed.
One approach to obtain the task head while using fixed w is to optimize only v in federated manner over all the data. In the case that g is given as a linear model and we absorb the softmax into L this defaults to Linear Probing (LP) (Nguenn al. 2023 Ren al. 2023a
3.2 FedNCM Algorithm
An alternative approach to derive an efficient g is through the use of NCM. We note that FedNCM, the federated version of NCM, can be derived exactly in a federated setting. In FedNCM, outlined in Algo. D. the server only communicates pre-trained weights once to each of the clients and the clients only communicate once with the server to send back their weighted class means. The server
3

---

Algorithm FedNCM. K is the total number of clients, C' is the number of classes in the training dataset, De is the total number of samples of class
Require: (X1, Y1), (X2, Y2)...., (XK. Yx)= Local datasets, Wpt pre-trained model
Server Executes:
1: for each client k E K in parallel do
2: (mElee LocalCCientSttts(X,, YE, wpr) Send to all clients, receive weighted class means
3: end for
4: for each class c E C do
R 5: 1e 4 D. -k=1 m
Di lc can be used in NCM classifier
6: end for
Client Side:
7:
function LOCAL.CLIENTSTTTSSE
Y, w)
8:
for each class C N do
9: Let X. = {or E X, = c}
10: me f Lx fu(II
11: end for
12: return (meleEc
13: end function
can then use each client's class means compute the NCM exactly and use them either to perform classification directly using the class centroids, or to initialize a linear task head for further fine-tuning. FedNCM allows efficient classifier approximation for pre-trained models and addresses many critical concerns in the FL setting including:
(a) communication and computation: FedNCM is negligible in both compute and communi- cation, it requires one communication from the server to each client and back. Used as an initialization for further FT, FedNCM speeds up convergence and reduces the communication and computation burden
(b) client statistical heterogeneity: Robust to typical non-iid distribution shifts (not the case for LP or FT). A FedNCM initialization also makes further FT more robust to heterogeneity,
Notably, FedNCM can be computed using secure aggregation methods. Furthermore, the lack of update to the base model parameters naturally improves differential privacy guarantees [Cattan et al. 2022).
To use NCM an initialization, consider the cross-entropy loss and (g f)() vf w) + b. We can set the matrix corresponding to the class c logit with the normalized class centroid 1./ ||lell and the bias term This allows us to initialize the task head with FedNCM and obtain further improvement through fine-tuning f.
3.3 Two-stage Approach for Transfer Learning in FL (Head + FineTTun)
FL algorithms are often unstable due to the mismatch in client objectives which can lead to large changes during local training causing significant deviations amongst the different client models. When using pre-trained model which allows us a powerful initial representation, we argue that a two-stage procedure will improve training stability and converge (HeadTune) we perform HeadTuning where the parameters of are updated e.g. by linear probing in federated fashion by using FedNCM. As stated in Sec. 3.2 FedNCM is highly efficient, imposing a negligible cost in compute and communication with respect to any typical fine-tuning stage. In the second stage (FineTune)) both and the classifier initialized in stage one, federated setting according to the FL objective function specified in Eq. [I] Taking the negligible cost of communication and compute provided by FedNCM into account, our two-stage approach can have a substantial advantage in convergence when compared to simply a fine-tuning stage [gguyen et al. 2023 Chen et al. 2023
quickly. In the first stage
fine tuned together in
We now give an intuitive interpretation of the advantages of method using the framework of Ren ct al. [003a]]. Assume that the k-th worker is initialized via Wo. and trained locally with SGD for several steps until it reaches the parameter wa. Writing w the optimal parameter, via triangular inequality, obtain the following inequality:

---

Exg (wx: (w"::""; ;X>)ll & [xx, [tllll//woo Xx)-s(w'3 (w: [(wo: Xx)II
(2)
In the neural tangent kernel (NTK) [Ren et al.. 2023a Jacot al. 2018 regime, for sufficiently
small step size, Ren al. 2023a] showed that the second term depends the approximation quality of the head go at initialization, which is bounded (where a is the sigmoid activation and {er the canonical basis) for some 0, by:
II (wa; Xx) f(wo: Xx)Il S c- E(x,, Yajlev, (fwo; Xx))ll
This suggests in particular that good choice of linear head will lead smaller right hand side term in Eq. 2 and thus reduce the distance to the optimum. Consequently, FedNCM or LP derived (compared to random v) may be expected lead to more rapid convergence. Thanks to the initial consensus on the classifier, we may also expect less client drift to occur, at least in the first round of training, when v initialized by HeadTuning, compared to a random initialization.
4 Experiments
In this section we will experimentally demonstrate the advantages of our proposed FedNCM and FedNCM+FT. Additionally, we show that simple LP tuning can times be more stable and commu- nication efficient than undertaking the full fine tuning considered almost exclusively in prior work on FL. with pre-trained models.
Our primary experiments focus on standard image classification tasks. We also provide some NLP clas- sification tasks in Sec. 4.2.1 We consider a setting similar to Nouyen et al. [2023] using the CIFAR 10 dataset [Krizhevsky| [009] and expand our setting to include four additional standard computer vision datasets shown in Tab. [J Following the method Hsu et al. [2019], data is distributed between clients using Dirichlet distribution parameterized by a 0.1 for our primary experiments. We set the num- ber of clients to 100, train for local epoch per round, and client participation to 30% for CIFAR (as in Nguyen et al. 2023) ). For all other datasets we use full client participation for simplicity.
Like Nguyen et al. [2023], we use SqueezeNet [landola et al. 2016]. we also consider ResNet18 [He 2016
Num. Classes
Num.
Dataset
Images
for experiments in Appendix E When performing fine-
CIFAR-10
10
50000
tuning and evaluation for all datasets, we resize images to
Flowers 102
102
1020
224 x 224, the training input size ImageNet. We run all experiments for three seeds using the FL.Sim library described in [gguyen et al. 22023].
Stanford Cars
196
8144
CUB
200
5994
EuroSAT-Sub
10
5000
Baseline methods We compare our methods to the fol- lowing approaches as per Nguyen et al. [2023]: : (a) Ran- dom: the model is initialized at random with no use of pre-trained model or NCM initialization. This setting cor- responds to the standard FL paradigm of McMahan et al. [2017]. (b) LP: Given a pre-trained model, we freeze the base and train only the linear head using standard FL optimizer for training. (c) FT: A pre-trained model is used to initialize the global model weights and then standard FL optimization algorithm is applied. (d) LP and FT Oracles: These equivalent baselines trained in the centralized setting that provide an upper bound to the expected performance.
Table 1: Summary of datasets used in our experiments.
All of the above baseline methods as well as our FedNCM and FedNCM+FT can be combined with any core FL optimization algorithm such as FedAvg and FedAdam (Reddi et al., 2020). Our experiments, we focus on the high-performing FedAvg, FedProx and FedAAdm which have been shown to do well these settings in prior art [guyen et al. 2023
Hyperpameters We follow the approach of Nguyen et al. [2023] Reddi et al. [2020] to select the
learning rate for each method on the various datasets. For CIFAR and SqueezeNet experiments we take the hyperparameters already derived in Nguyen et al. 2023 Additional details of selected
hyperparameters are provided in Appendix
Communication and Computation Budget We evaluate the communication and computation costs of each proposed method. Costs are considered both in total and given a fixed budget for either communication or computation. For the communication costs, we assume that each model parameter that needs to be transmitted is transmitted via a 32-bit floating point number. This assumption allows us to compute the total expected communication between clients and server. It is important to
5

---

Dataset
Method
Accuracy (%)
Total Compute (xFF
Total Comm. (GB)
Random
67.8 0.6
4.5 10a
1803,71
FT
85.4 0.4
3.0 10"
120.25
CIFAR
FcINCM-FT
87.2 0.2
3.0 107
120.25
LP
82.5 3: 0.2
1.0 10"
0.82
FodNCM
61.8 =
4.1 x 10
Random Nauven et
2023
34.2
1.5 x 10
601.24
gguyen
A223
1.5 10
CIFAR-10 x32
FodNCM-F
67.9 0.4
7.5 x 107
300.62
Ngueen el
2023
44.7
5.0 10"
4.10
F
40.02 0.04
4.10 x 10
Random FT
33.2
9.2 10"
1916.76
61.3 1.0
7.7 10
159.73
FLOWERS-102
FoINCM-FT
74.9 0.2
7.7 x 10
159.73
LP
71.1
3.I 10
20.33
FedNCM
71.8 0.03
4.2 x 10
Rankom
1..
5.4 10
2037.18
FT
6.9
1.8 x 10
679.06
FedNCM-FT
0.3
10
LP
0.30
9.0 10
122.88
FedNCM
0.2
8.2 x 10-2
Random
5.6 0.8
8.6 107
2370.97
FI
4.0
2.4
677.42
STANFORD CARS
ReINCMFF
54.8 1.2
2.4 107
677.42
LP
41.2 0.5
2.0 10
200.70
FedNCM
10
Random
85.8 2.7
5.3 x 107
2104.32
FT
95.6 1.0
1.5 10
601.24
EUROSAT-SUR
FedNCM-FT
96.0 1 0.5
1.5 x 107
601.24
LP
02.6 2
5.0 10
FedNCM
81.8 0181:00
4..10 10 10-3
Table Accuracy, total computation and total communication costs pure HeadTuniig methods (below dashed lines) and their full training counterparts. CIFAR-10 x32 indicates CIFAR-10 without re-sizing samples to 224 x 224, we include this setting so we can compare directly with values for FT, Random and LP reported by Nguyen al. (2023). We observe pure HeadTuning approaches, FedNCM and LP can be powerful, especially under compute and communication constraints. The unit, F, used to measure communication is one forward pass of single sample, details of the communication and computation calculations are provided in AppendixA
emphasize that linear probing only requires that we send client updates for the classifier rather than the entire model as is the case in the other settings. Consequently, costs when compared to FT for any given number of rounds. Our proposed FedNCM is one-round algorithm and therefore has even lower communication costs than any other algorithm considered.
has much lower communication
For computation time we consider the total FLOPs executed on the clients. We assume for simplicity that the backward pass of a model is 2x the forward pass. For example, in the case of LP (with data augmentation) each federated round leads to one forward communication on the base model, f. and one forward and one backward (equivalent to two forward passes) on the head, g. Similarly, for FedNCM the communication cost consists only one forward pass through the data.
4.1 Efficiency Pure HeadTuning for FL
As discussed Sec. I tuning the classifier head is at times as effective or more effective than updating the entire model in the context of transfer learning Evci al. 2202). In prior work, this situation was briefly considered as a limited case in Nguyen et al. [203 Appendix C.2] for CIFAR- 10 and suggested that tuning just the linear head (LP) might be weak approach in the heterogencous setting. Here we first revisit this claim and expand the scope of these experiments to highlight where LP can be beneficial in terms of performance, communication costs, and compute time. Subsequently, show another approach for approximating a good classifier, FedNCM, which can be competitive with orders of magnitude less computation and communication cost. We will demonstrate how to get the best of both HeadTuning and fine-tuning in the FL setting.
In Nguyen et al. 2023 the CIFAR-10 fine-tuning is done by feeding the 32 x 32 input image directly
into a pre-trained ImageNet model. Since the architectures are adapted to the 224 x 224 size and trained at this scale originally, such approach can lead to very large distribution shift and may be

---

sub-optimal for transfer learning. Thus we additionally compare to CIFAR- 10 using the traditional approach of resizing the image to the source data [Komblith et al.) 2019 Evei al. 2022]
Tab. 2 shows accuracy, compute, and communication cost results for Pure HeadTTuning Methods NCM and LP) well full tuning approaches including our FedNCM+FT. We note that in
Tab. 2 CIFAR- -1-AR-10-32 x 32 refers to results published in Nguyen et al. (2023). We first point out the difference image input size has on the results and conclusion. Overall accuracy is much higher (highest is 86% vs 63%) and the gap between FT and LP is substantially smaller when using the model's native input size, it shows an absolute improvement of only 4.6% vs 18.4%. For both sizes of CIFAR-10 CUB, Stanford Cars and Eurosat. FedNCM without FT can substantially exceed random performance while maintaining a highly competitive compute and communication budget. For the Flowers102 dataset, FedNCM can already far exceed the difficult-to-rral FT setting and furthermore, IP alone exceeds both FedNCM and FT. Our two-stage method of FedNCM+FT outperforms all other methods in terms of accuracy. In what follows we will show how FedNCM+FT also allows high efficiency given a specific, potentially limited compute and computational budget. When considering the results, we note that CIFAR- 10 contains the same object categories as the original ImageNet dataset but the Flowers 102 and CUB datasets, represent more realistic transfer learning tasks and under these conditions we observe the true effectiveness of HeadTuning methods such as FedNCM and LP.
4.2 FedNCM then FineTTune
We now study in more detail the two-stage approach described in Sec. 3.3 Fig. I shows the
comparison of our baselines and FedNCM+FT with FedAvg. We show both accuracies versus rounds as well as accuracy given a communication and computation cost budget. Firstly, we observe that when going beyond CIFAR-10, LP can converge rather quickly and sometimes to the same accuracy as FT. This result shows the importance of HeadTuning and supports the use LP in federated learning scenarios. Secondly, we can see the clear advantage of FedNCM+FT. After the stage one FedNCM initialization, it is able to achieve strong starting accuracy and in stage two, it converges with better accuracy than FT given the same computation budget. We note that FcdNCM+FT for Cars and EuroSAT-Sun sees an initial performance decrease during the first few rounds fine-tuning before rebounding to reach the best performance. This phenomenon is related to the chosen learning rate and further study is required to determine if this drop behaviour can be eliminated, potentially by using initial warm-up period at lower learning rate.
FeT converges rapidly, allowing it to be highly efficient under most communication budgets compared to other methods. The second column of Tab. ssb.ws that as we impose increasingly severe limitations on communication budgets, the performance gap between FedNCM+FT and all other methods widens significantly in favour of FedNCM+FT. Indeed for all the datasets FedNCM+FT is always optimal early on. For three of the datasets (Flowers, CUB, Cars) it exceeds LP over any communication budget. For CIFAR-10 10 and Eurosat LP can overtake it after the early stage, however, FedNCM+FT remains competitive and ultimately reaches higher performance. Similar trends are observed for computation time. We note overall as compared FT the performance improvement can be drastic when considering the trade-off of accuracy as a function of communication and compute available. We also remark that the variance of LP and FedNCM+FT is lower across runs than the FT and Random counterparts.
We note that the Random baseline, typically requires longer training than others to reach the conver- gence criteria, thus for the purpose our visualization we do not show the fully converged random baseline, which always requires many more communication rounds than the other approaches; however, the full curves are included in Appendix F

---

Accuracy vs Rounds Cilar
Accursay Communication Buige: Cifar
Aocuracy Comoute Budget
00-
20-
Rocees
Toul Conmaniaaiin (K1)
Total Compute (CFLOPs)
Accurasc Rounds Flowers
Accuracy Commun Budget Flowers
Accuracy Compute Budget Flowers
30-
91-
20-
20-
Round's
Torl Convra sicatio (81)
Toxal Cooputr (GF1OR))
Accurscy Bouncs CUB
Accunccy vi Commonicaiin Euégut CUB
Acconsey Compute Budest CUB
200
50-
43
30-
30-
21-
10.0
$30
(2.0)
800
1.000
1500
1:00
10
10
10
Tirsl Convra eicatie
Tetal (GFL
Accuracy
Communicatio Budest
Budget
Acuracy
30-
Hounds
Tora Convucicaion (89)
Toxal Compute (GPLOP))
Accuracy Roands Euroat
Accurazy Communication Budgot Eunsat
100
Accoracy Compurs Budget Euroaat
Acuumty
2010
Toxc Conve pnicatio (M0))
Texal Compune (SFLOR)
Figure 1: comparison of the accuracy hetween models initialized with pre-trained weights and trained on different downstream tasks. The final result for both an LP and an FT oracle are shown and we remark that the NCM initialization allows the model to outperform the results of the LP oracle.
4.2.1 NLP Experiments
We now explore the effectiveness of our approach on NLP tasks. For these experiments, we train a DistillBert model [Sanh et al. 2019 on the AG News dataset [Zhang al., 2015). Fig. shows the comparison of LP, FT and FedNCM+FT where we vary the number of examples per class (15, 90, 1470). We observe the strong performance of FedNCM previous observations around the importance of HeadTuning. As models get larger, the importance of the final layer on the overall size of the model diminishes resulting in larger gaps communication cost between one LP or FedNCM round and an FT round. Additionally, we find that FedNCM's accuracy robust to sample size decreases and FedNCMFFT provides substantial convergence, and by extension communication cost benefits. Indeed it is well known that centroid based classifiers perform robustly in the case of few samples (Snell et al. 2017). We provide additional confirmation of this in Appendix D with experiments using small subsets of C0.
stand alone method, reinforeing
8

---

Accurocy eamplees per
100
Accuracy vs Rounds AG News 90 eumples per dass
Acouracy us Rounds AG News 15 eomple dass
Accenner
40
30
33-
Roanbb
Figure 2: A comparison of the accuracy between models initialized with pre-trained trained using LP, FT and FF. We remark that the FedNCM alone obtains strong performance.
Dataset
Algorithm
FedNCM
FedNCM FT
FT
FEDAVO
64.8
87.2 0.2
85.4 + 0.4
CIFAR-10
FEDPROX
61.8
8.11=0.1
0.08
FEDADAM
04.8:.0.1
89.4
88.2 0.2
FEDAVG
718
71.9002
64.3 1.0
FLOWERS102
FEDPROX
0.08
75.2
65.5 1.5
PEDADAM
71.8 1.8110.00 0.03
76.7 0.2
66.6 1 1.0
Table Model performance with different methods for a variety FL algorithms for FedAvg, FedADAM and FedProx. FedNCM+FT outperforms in all cases.
4.3 Analysis and Ablations
We now focus on demonstrating other advantages of our two-stage method, with focus on using FedNCM as the method of choice for stage one (FedNCM+FT). In particular, we investigate robust- ness to larger number of clients, insensitivity to hyperparameters and compatibility with multiple FL algorithms and architectures. In Appendix we ablate our two-stage method by comparing FedNCM, LP and a three-layer MPL as methods for stage one.
4.3.1 Choice of FL Algorithm
So far have focused on comparisons using FedAvg as the baseline algorithm: however, since our method can be widely applied in FL, we further analyze FedNCM+FT using the FedAdam optimizer and the FedProx algorithm. Tab. 4.3. summarrize the results of using each method with 1.) Fed and 2.) only FT, for the Cifar- 10 and Flowers datasets.
We observe that improved FL. optimizers can complement the two-stage FedNCM+FT which sys- tematically exceeds the performance obtained when only fine-tuning. Regardless of the federated algorithm used, FedNCM alone continues to exceed the performance of FT on Flowers 102. Our results suggest that choice of the FL optimization algorithm [Nguyen et al. 2023] is not always the most critical consideration for optimal performance when using pre-trained models.
Hyperparameter Tuning FL algorithms are known to be challenging for hyperparameter selection Reddi et al.] 2020) and this can affect their practical application. We first note that FedNCM does not have any hyperparameters which already provides a large advantage. In Fig. 5 we observe the final performance for grid search over range of server and client learning rates for FedAdam using both FT and F We observe that FedNCM+FT not only has higher performance but is also more stable over the entire hyperparameter grid on Flowers dataset, and outperforms for all settings on CIF Varyinn Hottroganityy CFAR-10
varying Hoterogenity Fowers152
7002
$7.5
100
61
304
6/00
Figure 3: We vary the heterogeneity (Dirichlet-a)) for CIFAR- 10 and Flowers 102. Methods with HeadTuning: LP and FedNCM+FT are more robust, with the substantial advantage of FedNCM + FT increasing in challenging higher heterogeneity.

---

Segying Loczl MaaceCRAN 141
Dery Lacal Lacir FLOWERS
Uey Fg boca RJ006N
virging Louttspatte CARS
100
g140
Figure 4: We vary the number of local epochs. FedNCM+FT always outperforms FT and nearly always LP in this challenging setting.
PedXCM+FT
CCEPAR- Lc3
(CEPARIOO)
(F Plweres (2)
(Flowres (C)
30.00
30/00
01,00
03.05
10.00
18:00
10:00
$1.00
23:00
18:00
10:00
07.43
11:3
9128
00.76
4::31
3.3
1462
1400
14.1
51.00
12.12
CherlLeameogaate
Clirk Learnin Nute
Clidd Learing Rute
Figure 5: Hyperparameter grids for FedAdam for CIFAR-I0 FT, FedNCMFT (left) and Flowers (right). We observe CIFAR-10 FedNCM-FT tends to do better or equal for all hyperparameters compared to FT. For Flowers it is much easier to tune, achieving strong values over wide range, a noticeable advantage in FL
Heterogeneity Nguyen al. [023] points out that starting from a pre-trained model can reduce the effect of system heterogeneity. This is evaluated by comparing a specific Dirichlet distribution (a 0.1) used to partition data into a non-iid partitioning. Although the effect of heterogeneity is reduced we observe that in highly heterogeneous settings we still see substantial degradation in FT as shown in Fig. Here we consider for CIFAR- 10 the nearly iid a = 100,
0.1 as considered in [gguyen et al. 2023), and a very heterogeneous a = 0.01. Firstly, we observe that FedNCM+FT can provide benefits in the iid setting. As heterogeneity degrades the naive FT setting sees large absolute and relative drop in performance. On the other hand, FedNCM+FT as well as LP are able to degrade more gracefully.
Varying the Local Epoch The number of local epochs can drastically affect FL aglorithms, typi- cally larger amount of local computation between rounds is desired minimize communication.
However, this can often come at a cost of degraded perfor-
mance. We observe in Fig. as in Nguyen et al. 20233 that FT can be relatively robust in some cases (CIFAR- 10) to increasing local epochs. However, we also observe for some datasets that it can degrade, while LP and Fed- NCM+FT are less likely to degrade. Overall FedNCM+FT continues to outperform for larger local epochs.
Accusccy
Karying Humbarr
Charta
Increasing clients Tab. 6 shows that number of clients we observe that the degradation of Fed- NCM+FT is less severe than both LP and FT, suggesting it is stable under large number of workers being averaged. As discussed in Sec. 3.. expected in the same round that representation would shift less from a starting point, and therefore since the starting point is the same for all clients, we expect the client drift within round to be less given a fixed update budget.
increase the
Figure 6: We increase the number of clients on CIFAR- 10. FedNCM+FT de- grades more gracefully than FT and LP.
5 Conclusion and Limitations
We have highlighted the importance of the last layers in FL from pre-trained models. We used this observation to then derive two highly efficient methods FedNCM and FedNCM+FT whose advantages in terms of performance, communication, computation, and robustness to heterogeneity were demonstrated. A limitation of our work is that it focus on image data and models, as this the primary set of data and models studied in prior work particularly in the context of transfer learning.
10

---

Acknowledgements
This research was partially funded by NSERC Discovery Grant RGPIN- 2021-04104, FRONT New Researcher Grant and CÉRC Autonomous AI. E.O. acknowledges support from ANR ADONIS ANR-21-CE23-0030, We acknowledge resources provided by Compute Canada and Calcul Quebec.
References
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. Federated learning based dynamic regularization. arXiv preprint AXXi:2111.00426. 2021.
7aid Alyafeai, Maged Saced AlShaibani, and Irfan Ahmad. A survey on transfer learning in natural language processing. arXiv preprint arXiv:2007.04239 2020.
Yannis Cattan, Christopher A Choquette- Choo, Nicolas Papernot, and Abhradeep Thakurta. Fine- tenintunng with diffcrential privacy necessitates an additional hyperparameter search. arXiv preprint arXiv:2210.02156 2022.
Hong-You You Chen, Cheng-Haa Tu, Ziwei Li, Han Wei Shen, and Wei-Lun Chao. On the importance and applicability of pre-training for federated learning. 2023.
MohammadReza Davari, Nader Asadi, Sudhir Mudur, Rahaf Aljundi, and Eugene Belilovsky. Probing repre- sentation forgetting supervised and unsupervised continual learning. Proceedings of the IEEE/CVF Conference on Compuder Vision and Puttern Recognition, pages 167 12-16721, 2022.
Utku Evci, Vincent Dumoulin, Hugo Larochelle, and Michael C Mozer. Head2toe: Utilizing intermediate representations for better transfer learning. In International Conference on Machine Learning, pages 6009- 6033. PMLR, 2022.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated Icarning with thcoretical guarantees: -noosl-agnostic meta-learning approach. H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 3557- 3568. Curran Associates. Inc., 2020.
Ross Girshick, Jeff Donabue, Trevor Darrell, and Jitendra Malik. Rich feature hierarchies for accurate object detection and semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 580-587,200 2014.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Kaiming He, Ross Girshick, and Piotr Dollar. Rethinking imagenet pre-training. In Proceedings of the IEEEJCVF International Conference on Computer Vision (ICCV), October 2019.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Ges- mundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficien transfer learning for nlp. In International Conference on Machine Learning, pages 2790-2799. PMLR, 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Mcasuring the effects of non identical data distribution for federated visual classification. arXiv preprint arXiv: 96355,,335, 2019.
Forrest N. landola. Song Han, Matthew W. Moskewicz, Khalid Ashraf, William J. Dally, and Kurt Keutzer. Squeezenet: Alexnet-level accuracy with 50x fewer parameters and <0.5mb model size. 2016.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural langent kernel: Convergence and generalization neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Paul Janson, Wenxuan Zhang, Rahaf Aljundi, and Mohamed Elhoseiny. A simple baseline that questions the use of pretrained-models in continual learning. arXiv preprint arXiv:2210.00 10.04428, 2022.
Sai Pranceth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132-5143. PMLR, 2020.
Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Distributed machine learning for on-device intelligence. arXiv preprint 1rX.v.1610.02527, 2016.
11

---

Simon Kornblith, Jonathon Shlens, and Quoe V Le. Do better imagenet models transfer better? In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 2661-2671, 2019.
Alex Krizhevsky. Learning multiple layers of features from tiny images. pages 32-33, 2009. URL https
7/www. cs. wwwotto..eoooo...o weteeeeddu//// wttr:rrrr/ieaannngeeeatu 200---009---00 TR. pdf
Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution arXiv preprint arXiv:2202. 10054, 2022.
Gwen Legate, Lucas Caccia, and Eugene Belilovsky, Re-weighted softmax cross-entropy to control forgetting in federated learning. In Proceedings of The 2nd Conference on Lifelong Learning Agents, 2023.
Tian Li, Anit Kumar Sahu, Manzil Zaheer. Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. 2020.
Zhizhong Li and Derek Hoiem. Learning without forgetting. IEEE transactions on pattern analysis and machine intelligence, (0(12):993522447, 2017.
Dongze Lian, Zhou Daquan, Jiashi Feng, and Xinchao Wang. Scaling & shifting your features: A new baseline for efficient model tuning. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
Xtyeppat6ge
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas, Communication- efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282, PMLR, 2017.
John Nguyen, Jianyu Wang. Kshitiz Malik, Maziar Sanjabi, and Michael Rabbat. Where to hegin? on the impact of pre-training and initialization in federated learning. 2023.
Vishal M Patel, Raghuraman Gopalan, Ruonan Li, and Rama Chellappa. Visual domain adaptation: A survey of recent advances. IEFE Signal Processing Magazine, 32(3):53-69, 2015. doi: 10. 109/MSP2201
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. 2019.
Roger Rateliff. Connectionist models of recognition memory: constraints imposed by learning and forgetting functions. Psychological review, 97 2:285-308, 1990.
Sylvestre-Alvise Rebuffi, Alexander Kolesnikov, Georg Sperl, and Christoph H Lampert. icarl: Incremental classifier and representation learning. In Proceedings of the IEEE conference on Computer Vision and Puttern Recognition, pages 2001-2010, 2017.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00200 00295, 2020.
Yi Ren, Shangmin Guo, Wonho Bac, and Danica J. Sutherland. How to prepare your task head for finetuning. In The Eleventh International Conference on Learning Representations, 2023a. URL https://ppeeerrreee peet/forunnid-gV0X2preee
Yi Ren, Shangmin Guo, Wonho Bae, and Danica J Sutherland. How to prepare your task head for finetuning. arXiv preprint arXiv:2302.0 05779, 2023b.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. 5th Workshop on Energy Efficient Machine Learning and Cognitive Computing NewrlPS, 2019.
Aliaksandra Shysheya, John F Bronskill, Massimiliano Patacchiola, Sebastian Nowozin, and Richard E Turner. Fit: Parameter efficient few-shot transfer learning for personalized and federated image classification. In The Eleventh International Conference on Learning Representations, 2023. httpa://t
forum?id=9aokcc3V1111
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Sebastian U Stich. Local sgd converges fast and communicates little. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
Jianyu Wang and Gauri Joshi. Cooperative sgd: A unified framework for the design and analysis of caamunication-efffcieen sgd algorithms. arXiv preprint 808.1:1888.07576. 2018
12

---

Karl Weiss, Taghi M Khoshgoftaar, and Ding Wang. A survey of transfer learning. Journal of Big data," 3 (1):1-40, 2016.
Moslem Yazdanpanah, Aamer Abdul Rahman, Muawiz Chaudhary, Christian Desrosiers, Mohammad Havaei, Eugene Belilovsky, and Samira Ebrahimi Kahou. Revisiting learnable affines for batch norm in few-shot transfer learning. Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 9109-9118, June 2022.
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 5693-57000 2019.
Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In NIPS, 2015.
Fuzhen Zhuang. Zhiyuan Qi. Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Proceedings of the IEEE, 109(1):43-76, 2020.
13

---

Dataset
FT+FedNCM
FT
LP
Random
CIFAR-10
200
200
200
3000
CIFAR 10 X N:NGUYEN AL.
2023
1000
500
1000
1000
FLOWERS
250
250
500
3000
STANFORD CARS
1000
1000
2500
3500
AT-SSA
1000
1000
1000
3500
Table 4: rounds conducted for each dataset and method combination. For CIFAR-IO 32 x 32 experiments were conducted in Nguyen et al. 22023 with the exception of FedNCM+FT (our method)
Appendix
A
Calculation of Communication and Computation Values in Tab. 2
A.1
Communication
For HeadTuning methods such as LP and FedNCM, communication is calculated according to equation 3 where LLin and Loout are the input and output dimensions of the linear layer, I is the number of federated rounds, is the total number clients and F the fraction F E (0, 1] of clients participating in each federated round:
commi = 2 * (LLout * LLin) * K R F (32 bit)
For methods training the complete model, communication is calculated according to equation where P is the number of parameters present in the base model (excluding the linear layer) and T.J.jn- L.oouts R, K and F have the same meanings as above:
comm 2((Loout LLmm) + P) K R F (32 bit)
(4)
A.2 Computation
For computation, let F be one forward pass of single sample, Let S be the subset of clients, K, selected to participate in each federated round, R. We define one backward pass as 2F therefore each sample contributes 3F per round in full fine-tuning situation. For HeadTuning complete backward pass is not necessary since we only update the linear layer weights and each sample contributes F per round. Eq. 5ad Eq. formalize our computation calculation method in units of F, where E is the number of local epochs performed by each client and N₂ is the sum of the number of samples at each selected client, S.
compy: 3 R R
(5)
comp = R E*
(6)
B Hyperparameter Settings
For CIFAR-10, CUB, Stanford Cars and Eurosat datasets the learning rates for the FedAvg algorithm were tuned via a grid search over learning rates {0.1, 0.07, 0.05, 0.03, 0.01, 0.007,00000 0.005, 0.003, 0.001 For Flowers102, based prelimi- nary analysis we used lower learning rates were tuned over learning rates (0.01, 0.007, 0.005, 0.003, 0.001, 0.0007, 0.0005, 0.0003, 0.0001}. Tab. 4 summarizes the number of rounds conducted for each dataset and method combination.
Prior work on federated learning with pre-trained models has indicated that for FedADAM lower global learning rates and higher client learning rates were more effective. As a result for CIFAR- 10 and Flowers the client learning rate was tuned over 0.1, 0.01, 0.001, 0.0001 and the server learning rate was tuned over (0.001, 0.0001, 0.00001, 0.000001) }. each combination of server and client learning rates were tried.
14

---

Dataset
HeadTuning Method
Accuracy
LP
0.2
CIFAR-10
MLP
84.1
CP
74.1117.1
MLP
7..5
CP
30.0011..3
CUB
MLP
50.3
11.2+0.3
CARS
MLP 41.0
32.61.0.44
SUROSAT-SUR
MLP
91.7
Table 5: LP and MLP HeadTuning results prior to the fine-tuning stage.
Dataset
stage
Method
Accuracy
LP (5 ROUNDS)
85.9 U.4
CIFAR-10 10
ROUNDS)
0.17
FEONCM
87.2 0.2
N/A
85.4
FLOWERS102
LP (5 ROUNDS)
68.6
LP (10 ROUNDS)
08.3 -
FEDNCM
N/A
64.5 1 0.1
Table Outcomes using LP as the Heaing method in our two stage training algorithm, FedNCM results from Tab. are included for comparison.
C
Additional HeadTunnig Analysis
C.1 MLP HeadTuuning
We compare HeadTuning methods by replacing the linear layer in the LP method with a three layer MLP. Tab, aab.[]ishon shows the HeadTuning results obtained using LP and MLP. We observe no advantage to using an MLP instead on a linear layer since most of our results are comparable or in some cases even inferior to the LP case.
C.2 LP and FT
In this section we explore the use of LP as the HeadTuning method in our two-stage algorithm. Tab. indicates performance of our two stage Head + munittunigg method using five rounds of LP. ten rounds of LP and FedNCM as the Head methods. The bottom row in each dataset category where stage Methods denoted n/a, is the result of only fine-tuning, i.e. not using our two stage method, which shows that HeadTuuning prior to fine-tuning is always at least as effective as FT alone and, depending on the HeadTTuning method selected, much more effective.
D Experiments Using Small Subsets of Data
We conduct experiments using only small subset of Cifar- 10 with nine samples of cach class, i.e. 90 samples in total. We distribute these samples between five clients using a Dirichlet distribution with 0.1 each client will have very few samples and some will be entirely missing samples from many classes. Tab. 7 shows a 2 10% improvement in accuracy between FT and FedNCM. which attribute to overfitting of the FT model and the challenge of heterogeneity that FedNCM is not susceptible to. We also observe that LP does quite a bit better than FT indicating the benefit of HeadTuning methods in this setting.
E
ResNet Experiments
We perform experiments for LP, FT, FedNCM and FedNCM+FT using the ResNet18 model using the FedAvg algorithm and the CIFAR-10 10 and Flowers datasets. Results are summarized in Tab. where we observe that FedNCM performance is better by almost 13% compared to squeezenet, while performance degraded compared squeezenet. We hypothesis this is due the challenges of deeper networks in heterogenous federated learning. For the Flowers 102 dataset, FedNCM and
15

---

LP
FT
FedNCM
FeuNCM +IT
53.8 + 0.3
45.2 2.4
55.6 + 0.8
56.7 36.7=00.
Table 7: ResNet18 model performance for FedAvg. As with Squeezenet, FedNCM++T continues to outperforms in all cases.
Dataset
FedNCM
FedNCM FT
FT4Premmin
L.PPretrain
CIFAR 10
77.74 0.05 79.05 1 1.31
77.87
4.07
74.73 3.03 26.35 2.50
FLOWERRIOE 102 734.13
"74. = 2026
3430 10.00
Table 8: ResNet18 model performance for FedAvg- As with Squeezenet, FedNCM+FT continues outperforms all cases.
FedNCM+FT produce the best results far. Additionally, for flowers FedNCM outperformed all other methods. The variance between runs using ResNet18 is much higher than was observed for SqueezeNet, FedNCM appears to help stabilize the results since it provides the most consistency by for both datasets.
Extended Accuracy Comparison Figures
Fig. is the extended version of Fig. in the main body of the paper. In the paper we truncate the number of round displayed for the random setting since random requires many more rounds to converge than the other methods. Fig. shows these same figures with the entirety of the training rounds displayed for the random setting.
G Compute
We use a combination of NVIDIA 1000SXMM440GGB NVIDIA RTX A4500, Tesla V100-SXM2- 32GB and Tesla PIOO-PCIE- 12GB GPUs for total of 1.1 GPU years In addition the experiments reported in the paper, this includes preliminary experiments and hyperparameter searches.
16

---

Accuracy Roands Cilar
Rounds Flowers
Accurzcy Rounds
750
1000
12:0
1900
17.0
2000
1000
2000
121
Roueds
Rounds
Hounds
Accuracy Rounds Cars
Accuracy va Rounds Earosut
30
5401 18881 1000
3000
3000
2000
$590
Figure 7: The full training of Random baseline corresponding to Fig. 1 in the paper is shown. We observer Random is always very far from the other baseliens and converges slowly.
17


Reviews:
[{"review_id": "RQVvgK3xk0", "reviewer": "Reviewer_63Tk", "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.", "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?\n\nCan you put a theoretical convergence analysis of the two-stage FedNCM to show its advantage compared with the one without first-stage head tuning?", "limitations": "Please see above.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Investigating the effect of a linear head on the final FL model tuning is interesting.\n", "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).\n\nThere 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.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful comments and their interest in our research direction. We address the comments in turn:\n\n**Weakness 1: Closed form solution**. \nOur experiments largely consider the cross entropy loss for which there is no closed form solution. Indeed, it is an interesting idea to obtain **V** by solving directly the right-hand side of the equation on the top of page 5. However, we emphasize that FedNCM gives a **V** which is the NCM classifier across all client datasets. While a closed form **V** for each client would lead to a question of how to aggregation a single **V** for all clients, since the **V**’s may heavily disagree. \n\nFinally, we emphasize that the analysis at the end of 3.3 is aimed to provide intuition for improved convergence and the loss we optimize is different.\n\n**Weakness 2: related work**\nWe emphasize that our work is focused on compute/communication efficient FL *in the setting where pretrained models are available* which has different considerations. We thank the reviewer for the relevant reference, we will add it to our citations.  Shysheya et. al. focuses on fine-tuning methods based on adapting only the affine parameters with applications in few shot transfer learning, personalization, and also federated learning. Indeed the federated learning experiments focus on only CIFAR100 and the few shot federated learning setting. Our work, on the other hand, focuses exclusively on the federated learning setting and starts from the observations in very recent literature (ICLR 2023) regarding how to treat pre-trained models in federated learning. The fine-tuning method proposed in Shysheya et. al. can potentially be combined as part of the FT step with our two phase procedure to further reduce communication cost in the FT step; however, we leave this to future work. We also note that Shysheya is specialized to image data while our work is more general and can be applied to any classification problem, as shown in the new NLP experiments (see pdf)\n\n**Question Triangular inequality extension to aggregation and theoretical convergence results**\nThe triangle inequality can be extended to multiple clients, with some additional assumptions. We consider below an extension to FedAvg style aggregation \n\nConsider the aggregation $\\bar{w}=\\frac{1}{K}\\sum_k w_k$ \n\nWe can now write $|f(\\bar{w}; [X_1, ... , X_K])-f(w^*; [X_1, ... , X_K])| \\leq |f(w; X_k)-f(w^*; X_k)|$\n\nThen using an additional triangular inequality:\n$\\leq \\sum_k |f(w^*; X_k)-f(w_0;X_k)| + \\sum_k |f(w_k; X_k)-f(w_0;X_k)| + \\sum_k |f(\\bar{w}; X_k)-f(w_k;X_k)|$\n\nWith additional regularity assumption (namely L smoothness, and boundedness of both the data and the gradient) the last term can be bounded by a term proportional to $\\sum_k |w_k-\\bar{w}|$\n\nThe term $\\sum_k |w_k-\\bar{w}|$ can be seen as a measure of consensus between workers and existing works provide potential upper bounds.  For instance, using Lemma 10 of Stitch et al. [1],  this can be upper bound by several additional assumptions, notably small learning rates. Note that In the NTK regime assumptions of [Ren et al 2023], step sizes are relatively small, which matches \nLemma 10 of Stitch et al. [1]. Bounding some notion of agreement between workers *e.g.* $\\sum_k |w_k-\\bar{w}|$  is important to guarantee convergence (*e.g.* [1,2]).\n\nWe emphasize the above is aimed to provide intuition and the analysis is limited to the first round of training, extending it beyond the first round and to obtain a rate is not obvious without additional theoretical tools. We emphasize there is a limited number of federated learning from pre-trained model references currently available and to our knowledge prior work in this area has focused on empirical analysis. Similarly,  two stage fine tuning procedures have been analyzed in terms of their generalization properties (Kumar et al, Ren et al) but not in terms of convergence rates. A deeper look at the theoretical convergence guarantees of federated transfer learning and two stage finetuning is certainly a valid and interesting research direction but is currently out of scope of our work which focuses on empirical results. \n\n[1] Stich, Sebastian U. \"Local SGD converges fast and communicates little.\" arXiv preprint arXiv:1805.09767 (2018).\n\n[2] Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020).\n\n\n\n\n"}, {"review_id": "ddNDflv20b", "reviewer": "Reviewer_8wWr", "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.", "questions": "1. Line 165: Could you explain NTK or add a reference? \n2. 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.\n3. 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?\n4. Could you add more explanation about Fig 1, such as the different behavior on different datasets?", "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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n2. It provides code and experimental results on several vision datasets, including CIFAR-10, Flowers102, Stanford Cars, CUB and EuroSAT-Sub.\n3. 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.\n2. However, the application of pre-trained models in FL has already been proposed in Nguyen et al. ICLR 2023. \n3. The experimental comparison is not considering Federated Learning baselines, like FedProx, SCAFFOLD, FedYogi, LG-FedAvg [1], FedPer [2], etc.\n4. The experiments only include small scale datasets, in total less than 50,000.\n5. There is not much theory on the communication cost, or convergence rate analysis.\n6. Minor: Line 201 FedAVG -> FedAvg; Line 234: imagenet -> ImageNet;\n\n\n[1] Liang, Paul Pu, et al. \"Think locally, act globally: Federated learning with local and global representations.\" arXiv preprint arXiv:2001.01523 (2020).\n[2] Arivazhagan, Manoj Ghuhan, et al. \"Federated learning with personalization layers.\" arXiv preprint arXiv:1912.00818 (2019).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful comments which help us to improve the paper, we address each comment in turn:\n\n**Clarification:** Our results show substantial improvement over prior work, not just the random baseline.\n\n**W1: Novelty**\nWe emphasize that the simplicity of the proposed approach is one of its strengths. The judgment of novelty is subjective. We argue the observation that NCM is a powerful technique in the pre-trained FL setting and can be implemented under FL constraints. Moreover, the combination with a second stage of fine-tuning is not obvious and therefore quite novel.\n\n**W2: Prior work on pretrained models (Nguyen et al)**\nWhile it is true that Nguyen et al address the application of pre-trained models to FL, the body of work in this domain is still limited. Critically, our work highlights a very important and practically relevant limitation of the considerations of prior work. Specifically, both Nguyen et. al. and Chen et al  use pre-trained models as an initialization point and subsequently adapt **all** model parameters via fine-tuning. They observe for example that the choice of FL algorithm is less critical. Our work shows that there are more important considerations for pre-trained FL (*e.g.* only adapting the classifier, using a 2 phase approach) which can yield drastic improvements and supersede the choice of FL algorithm.   An additional benefit of our work is the use of a greater variety of datasets showing a greater diversity of results, thus exposing the limitations of only fine-tuning.\n   \n**W3: Baseline FL algorithms**\nThe baselines we chose are based on the prior work of Nguyen et al which shows (a) many FL methods achieve similar performance in the pre-trained setting (b) FedAdam and FedAvg were found to be strong performing baselines. Our preliminary experiments found that SCAFFOLD performs very poorly, thus it was excluded from subsequent analysis. Additionally, we note that FedPer is a personalization method while we focus on the standard shared model setting of FL.\n \nWe agree that convincingly showing our method can work with a variety of base FL algorithms is beneficial, thus we now provide the results  for FT and FedNCM+FT using FedProx  (another well performing algorithm from Nguyen et al’) as the baseline method. We thank the reviewer for this suggestion which strengthens our evaluation.\n## FedProx\n|     DS         |       FT        |       FedNCM+FT|\n|----------------|-----------------|-----------------------|\n|Cifar-10     | 87.8 (0.08)    |88.08 (0.1)|\n|Flowers     | 65.5 (1.5)      |75.2 (0.01)|\n\nFedProx performs similarly to FedAdam in Table 3 and our method substantially improves performance. We will include these results in our camera-ready.\n\n**W4: dataset scale**\nWe use a variety of datasets consistent with previous work on pre-trained models. Our use of Cifar-10 allows us to make a direct comparison with Nguyen et. al., the best known work on federated transfer learning. We have also expanded our results to include some text datasets (see pdf), some with up to 120K samples.\n\n**W5: minimal theory on communication and convergence rate**\nWe note that since we use a standard base FL algorithm, we will typically inherit its convergence analysis as a crude upper bound. Similar to the prior work in the pre-trained setting, our work focuses on practical empirical concerns. We leave further theoretical convergence rate analysis of fine-tuning and two stage training to future work. Also see more detailed response to 63Tk\n\n**W6: minor typos**\nThank you for pointing out these inconsistencies, we will fix them in the final version.\n\n**Question 1:**\nNTK refers to the neural tangent kernel (Jacot et al, Ren et al), a recent tool for analyzing neural networks under assumptions on their training dynamics. In our work, we follow Ren et al and use NTK regime assumptions to show that initializing the linear head using FedNCM results in a smaller distance after fixed number of steps to the optimal solution than a random initialization.\n\n[1] Jacot, A., Gabriel, F., & Hongler, C. (2018). Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31.\n\n**Question 2:**\nFedNCM is defined in section 3.2 (L134), it is the federated version of NCM and the algorithm is provided in Algo. 1. Random, LP and FT are defined in L191-L199 under the header Baseline Methods and the terms LP and FT Pretrain are simply LP and FT with “pretrain” added to explicitly indicate that we begin from pre-trained models. FedNCM+FT is our proposed two stage training where phase 1 →  FedNCM (HeadTuning) and phase 2 → FT (fine-tuning). In the final version we will make our explanation of the two stage method clearer and clarify or omit the use of “pretrain” in conjunction with FT and LP.\n\n**Question 3:**\nThere has been limited work using pre-trained model initializations in FL, the best known reference is Nguyen et al with which we compare our work extensively. We have extend the base FL algorithms to include FedProx.\n**Why NCM?:** NCM is a method of initializing the weights of the classifier that is preferable to a random initialization. Compared to training a linear classifier, it can be done in a single round and is not affected by heterogeneity. Furthermore, in the FedNCM+FT approach we use FedNCM to initialize a standard softmax classifier.\n\n**Question 4:**\nFig. 1 is discussed in L255-L272 which emphasizes when the behaviors differ across datasets. We will further briefly expand the discussion of the initial drop in accuracy on EuroSAT (see more detailed reply to reviewer tHeD).\n\n**Regarding Limitations:** We will expand the limitations to mention that we do not provide guarantees on more rapid convergence of the two phase process, note however that we do inherit the guarantees of the base FL method. Regarding large scale datasets, we emphasize above we use realistic and commonly used transfer learning datasets."}, {"review_id": "OnGCKlnHdz", "reviewer": "Reviewer_tHeD", "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.\n\nThis 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.", "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. \n\nI 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.", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Experiments are done with rigor and described in detail the paper.\n2. 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.\n", "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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their overall positive view of the work and their helpful comments which will improve the paper. We address the comments in turn:\n\n**Weakness 1: Comparisons to first stage LP**\nWe have added a comparison below with the suggested baseline of running LP for several rounds. We observe that FedNCM outperforms in the end while being simpler in the sense of not requiring additional communication rounds (and the associated overhead). It is an interesting idea to consider other Head Tuning approaches in our proposed two phase procedure in FL. We note the advantage of FedNCM, besides only taking one round, is that it naturally avoids heterogeneity and hyperparameter search issues which come with using LP or other HeadTuning approaches.  We will include these ablations in the appendix. \n                \n|DS           |        5 rounds LP + FT   |         10 rounds LP + FT      |   FedNCM + FT     |\n|-----------|--------------------------------|------------------------------------|--------------------------|\n|Flowers |        68.6 (1.3)                  |       68.3 (0.6)                       |    74.9 (0.2)              |\n|Cifar-10 |        85.9(0.4)                   |       84.5 (1.7)                       |   87.2 (0.2)                |\n\n**Weakness 2: FedNCM with few samples / NLP Tasks**\n\n**Few Samples:**\t\nFedNCM has substantial benefits in the case of few samples as centroid based classifiers are often used in such settings and can perform robustly [1]. To confirm this we provide the results of an experiment using only a small subset of Cifar-10 with nine samples of each class, *i.e.* 90 samples in total. We distribute these samples between five clients using a Dirichlet distribution with $\\alpha=0.1$ so each client will have very few samples and some will be entirely missing samples from many classes. We observe a $>10$% improvement in accuracy between FT and FedNCM, which we attribute to overfitting of the FT model and the challenge of heterogeneity that FedNCM is not susceptible to. We also observe that LP does quite a bit better than FT indicating the benefit of HeadTuning methods in this setting.\n|LP             |FT             | FedNCM      | FT+FedNCM |\n|--------------|--------------|  ----------------  | ----------------|\n|53.8 (0.3) | 45.2 (2.4) |    55.6 (0.8)  |    56.7 (0.7) | \n\n**NLP:**\nWe also take the reviewer’s suggestion of considering an NLP problem and extend our results to consider several NLP tasks with a distillbert base. The results are shown in the accompanying pdf, and we find that FedNCM’s accuracy is robust as the sample size decreases and FedNCM+FT gives large benefits in convergence and thereby communication costs. \n\n[1] Snell, Jake, Kevin Swersky, and Richard Zemel. \"Prototypical networks for few-shot learning.\" Advances in neural information processing systems 30 (2017).\n\n**Question: Why rapid initial drop in some cases**\nThe behaviour you’ve mentioned in Cars and Eurosat where accuracy will decrease during the first few rounds of fine-tuning before rebounding and continuing to improve is related to the chosen learning rate from the hyperparameter search, which uses the standard setup for hyperparameter search (e.g. Nguyen et al and Reddit et al) that selects the best performing final accuracy on the validation set. Based on our results we are still able to achieve better performance than with any other method, but further study is required to determine if this drop  behaviour can be eliminated using an initial warm-up period at a lower learning rate. For a production system that may have other constraints than final accuracy, a different criteria can be used for the hyperparameter selection; \n"}, {"review_id": "pgnRcqUGs7", "reviewer": "Reviewer_1Br4", "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.", "questions": "1) It seems that only head tuning + full tuning leads to the best results? But then it is not communication cheap?\n\n2) 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?\n\n3) How the approach relates to layerwise personalization federated algorithms?", "limitations": "Evaluation for heterogeneity is done only with very particular non-iid with respect to labels, but it is not highlighted.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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?).\n\n2) 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.\n\n3) 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.\n\nMinor\n- table3 is never referenced\n- figure4 is referenced before figure3\n- figure5 is called table5", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging benefits of our work and their helpful comments which help us improve the clarity of the paper. We address the comments in turn:\n\n**W1: is the final layer updated during fine-tuning** \nWe apologize for the confusion, the final phase also updates the final layer head as in standard finetuning. We will clarify this in sec 3.3 on L156-157.\n\n**W2: Comparisons to methods addressing heterogeneity** \nOur work is based in the pre-trained FL setting, Ngyuen et al have already shown that in this setting, methods that have been developed particularly to address heterogeneity challenges in FL do not outperform FL algorithms not designed to address heterogeneity (*e.g.* FedAvg, FedAdam).  We thus focused on strong baselines for this setting that perform well, *i.e.* FedAvg and FedAdam. However, we have now done further experiments with FedProx, another high performing baseline from Nugyen et al that is designed to tackle heterogeneity. Please see results in the response to tHeD.\n\n**W3: Message**\nOur work has two related messages:\n1) A  two stage method for training pre-trained models in a federated setting works best. The two stage method consists of first tuning the classification head (HeadTune) by any suitable HeadTuning method. In the second phase (FineTune) the entire model is fine tuned in a federated setting. In cases where compute or communication budgets are limited we highlight that using a HeadTuning method alone may be sufficient depending on performance requirements and the complexity of the downstream dataset  \n2) We propose a HeadTuning method, FedNCM that we demonstrate is highly efficient with regard to both communication, compute, and naturally robust to heterogeneity. We focus on the specific two-stage method where FedNCM is our chosen HeadTuning method which we follow by fine-tuning the entire model. We highlight this specific two-stage method since it achieves the best results overall.\n\nWe agree the distinction could be clearer and will clarify this at the end of the introduction, we hope this will also help to address some of the ambiguity regarding the approach in your first point.\n  \n**Q1: FedNCM+FT Not communication cheap** \nDuring full fine-tuning, all model weights are shared; however, in both NCM and LP only features, which represent a small  fraction of parameters, are shared by the clients. This means the cost of HeadTuning using NCM or LP is negligible compared to full fine-tuning. Despite this, in all cases HeadTuning causes the model to converge  faster as shown in Fig. 1, thus leading to a much better accuracy for a given communication budget. \n\n**Q2: NCM privacy/security**\nWe argue that NCM and LP have less risk in terms of data privacy of the final FL trained model than FT trained models. This stems from the lack of updates to the base model parameters. Assuming the model is pre-trained on a public data, not updating the majority of the model leads to naturally improved DP guarantees (e.g. Tang et al, Cattan et al). We will clarify this in the text. \nSimilar to gradient inversion attacks, FedNCM could be susceptible to an attempt by the attacker to invert the feature mean, however it is typically a challenging problem to invert even a single input sample from the last representation layer. For example, Dosovitskiy et al find it difficult to invert even for a shallow AlexNet layer. For FedNCM this inversion is even more challenging since unless a class has  only one sample at the client, a client the mean is transmitted which significantly dilutes individual class information and provides increased privacy.\n\n**Q3: personlized FL/layerwise**\nOur paper focuses on the non-personalized FL setting and starts from pre-trained models which is quite different to Ma et al work on personalized, layerwise FL. However, future work can consider personalization and combinations with layerwise update algorithms. Regarding the non-iid setting, we follow  a standard approach (and alpha value) from Nguyen et al, Reddi et al which selects a fixed alpha for the primary experiments. We also include results for other heterogeneity levels in the ablations. \n\n**References**\n\n[1] Tang, Qiaoyue, Lecuyer, Mathias. “DP-SGD-LF: Improving Utility under Differentially Private Learning via Layer Freezing”. 2023 https://openreview.net/forum?id=coLtCLTHFbW\n\n[2] Cattan, Yannis, et al. \"Fine-tuning with differential privacy necessitates an additional hyperparameter search.\" arXiv preprint arXiv:2210.02156 (2022).\n\n[3] Dosovitskiy, Alexey, and Thomas Brox. \"Inverting visual representations with convolutional networks.\" Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.\n\n[4] Ma, Xiaosong, et al. \"Layer-wised model aggregation for personalized federated learning.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022."}]

Author rebuttal:
We thank the reviewers for their helpful comments and suggestions that help us improve the paper. We have responded to the reviewers in the individual comments. 

We additionally include a pdf with NLP experiments suggested by reviewer tHeD. These show additional results on NLP Datasets using the Distillbert model as in (Nguyen et al). We consider two downstream datasets taken from the HuggingFace datasets repository;  AG News which consists of 4 classes, and a binary task Rotten Tomatoes. AG News has a total of 84K samples and  Rotten Tomatoes a total of 5K samples. We observe that FedNCM provides a powerful initial guess at low communication cost, that improves relative to FT as the dataset size decreases. We also observe the improved convergence performance of FedNCM+FT over FT as in our image experiments.

We also point the reviewers to our additional results with FedProx given in response to reviewer 8wWr


Meta-review: The submission presents a simple recipe of practical relevance that has not appeared in the literature in this form before. The authors’ response clarified a number of reviewer questions regarding the described setting, the relation to prior work and the experimental evaluation. After the discussion phase, a majority of reviewers recommended acceptance.
Decision: Accept (poster)