Title: Guiding The Last Layer in Federated Learning with Pre-Trained Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: federated learning; transfer learning; nearest mean classifier; continual learning;
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

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. Code for our experiments is available 1 .

1 Introduction

In recent years an increased focus on data privacy has attracted significant research interest in Federated Learning (FL). FL is an approach in which a common global model is trained by aggregating model updates computed across a set of distributed edge devices whose data is kept private. We desire to train a federated global model capable of equivalent performance as one trained on the same set of centralized data. However, it is worth noting that FedAvg [McMahan et 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 [Li et al., 2020, Acar et al., 2021, Karimireddy et al., 2020].

Transfer learning from pre-trained models that have been trained on 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 et al., 2014], Natural Language Processing (NLP) [Radford et al., 2019], 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 et al., 2016, Patel et al., 2015]. FL literature on the other hand, has been largely focused on models trained from scratch [McMahan et al., 2017, Karimireddy et al., 2020, Li et al., 2020] and the impact

Montreal, Canada

of heterogeneity on algorithmic convergence. Recently, several studies have been conducted on the effect of pre-training on the performance of standard FL algorithms, see e.g., Chen et al. [2023], Nguyen et al. [2023]. 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.

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 [Kornblith et al., 2019], affine parameters [Lian et al., 2022, Yazdanpanah et al., 2022], or adapters Houlsby et al. [2019]. These approaches frequently yield combinations of better performance, computation time and better prevention against over-fitting. The selection is often a function of architecture, task similarity, and dataset size [Evci et al., 2022, Shysheya et al., 2023, Yazdanpanah et al., 2022]. 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 FL is the reduction of communication cost, having largely motivated the design of the FedAvg algorithm (see e.g., McMahan et al. [2017]). Although not studied in the prior works, updating only the linear classifier can be highly efficient in the federated setting when starting from a 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 a 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. [2023, Appendix C].

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-Tune stage). Such two-stage approaches have been applied in practice and studied theoretically in the transfer learning literature [Kumar et al., 2022, Ren et al., 2023b]. They have been shown to give improved performance in both in-distribution and out-of-distribution settings [Kumar et al., 2022]. 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 HeadTuning stage, our work highlights the Nearest Class Mean (NCM), a 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-tuning) 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 or 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 a stand alone technique, or as stage one (HeadTuning) in our proposed two stage process which leads to improved accuracy.

2 Related work

Federated Learning The most well known approach in FL is 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 [Karimireddy et al., 2020, Li et al., 2020, Fallah et al., 2020, Yu et al., 2019]. A commonly

cited problem in the literature is the challenge of heterogeneous or non-iid data and a variety of algorithms have been developed to tackle this [Li et al., 2020, Hsu et al., 2019, Legate et al., 2023, Karimireddy et al., 2020].

Transfer Learning Transfer learning is widely used in many domains where data is scarce [Girshick et al., 2014, Alyafeai et al., 2020, Zhuang et al., 2020, 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 Kornblith et al. [2019] and some more efficient methods such as selecting features Evci et al. [2022], adding affine parameters Lian et al. [2022], Yazdanpanah et al. [2022], and adapters for transformers Houlsby 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-trained 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].

Nearest Class Means Classifier The use of the NCM algorithm in 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 et al. [2017], Li and Hoiem [2017], Davari et al. [2022] who maintain a memory of exemplars used to compute an NCM classifier. Related to our 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 as 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.

3 Methods

3.1 Background and Notation

In FL, distributed optimization occurs over K clients with each client k ∈ { 1 , ..., K } having data X k , Y k that contains n k samples drawn from distribution D k . We define the total number of samples across all clients as n = ∑ K k =1 n k . The data X k 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. 1 [Koneˇ cn` y et al., 2016] and aims to find the minimizer of the loss over the sum of the client data:

Here we have split the model prediction into f , a base parameterized by w that produces representations, 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 a standard classification or regression loss. The w are derived from a pre-trained model and they can be optimized or held fixed.

One approach to obtain the task head while using a fixed w is to optimize only v in a 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)[Nguyen et al., 2023, Ren et 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. 1., 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

Algorithm 1 FedNCM . K is the total number of clients, C is the number of classes in the training dataset, D c is the total number of samples of class c

Require: ( X 1 , Y 1 ) , ( X 2 , Y 2 ) , . . . , ( X K , Y K ) - Local datasets, w pt - pre-trained model Server Executes:

1: for each client k ∈ K in parallel do

2:

[

m

k

c

]

c

∈

C

LocalClientStats(

3: end for

4: for each class c ∈ C do

5: l c ← 1 D c ∑ K k =1 m k c

6: end for

Client Side:

can then use each client's class means to 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 an 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 communication, 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 as an initialization, consider the cross-entropy loss and ( g ◦ f )( x ) = v f ( x ; w ) + b . We can set the matrix v corresponding to the class c logit with the normalized class centroid l c / ∥ l c ∥ and the bias term to 0. 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 (HeadTune + FineTune)

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 a pre-trained model which allows us a powerful initial representation, we argue that a two-stage procedure will improve training stability and converge more quickly. In the first stage (HeadTune) we perform HeadTuning where the parameters of g are updated e.g. by linear probing in federated fashion or 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 f and the classifier initialized in stage one, are fine tuned together in a federated setting according to the FL objective function specified in Eq. 1. 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 [Nguyen et al., 2023, Chen et al., 2023].

We now give an intuitive interpretation of the advantages of our method using the framework of Ren et al. [2023a]. Assume that the k -th worker is initialized via w 0 , and trained locally with SGD for several steps until it reaches the parameter w k . Writing w ∗ the optimal parameter, via triangular inequality, we obtain the following inequality:

X

k

, Y

k

,

w

pt

)

▷

Send to all clients, receive weighted class means

▷ l c can be used in NCM classifier

In the neural tangent kernel (NTK) [Ren et al., 2023a, Jacot et al., 2018] regime, for sufficiently small step size, Ren et al. [2023a] showed that the second term depends on the approximation quality of the head g 0 at initialization, which is bounded (where σ is the sigmoid activation and { e i } i the canonical basis) for some c > 0 , by:

This suggests in particular that a good choice of linear head v will lead to a smaller right hand side term in Eq. 2, and thus reduce the distance to the optimum. Consequently, FedNCM or LP derived v (compared to a random v ) may be expected to lead to a 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 it 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 at times be more stable and communication 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 classification tasks in Sec. 4.2.1. We consider a setting similar to Nguyen et al. [2023] using the CIFAR 10 dataset [Krizhevsky, 2009] and expand our setting to include four additional standard computer vision datasets shown in Tab. 1. Following the method of Hsu et al. [2019], data is distributed between clients using a Dirichlet distribution parameterized by α = 0 . 1 for our primary experiments. We set the number of clients to 100, train for 1 local epoch per round, and set 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 [Iandola et al., 2016], we also consider a ResNet18 [He et al., 2016] for experiments in Appendix E. When performing finetuning and evaluation for all datasets, we resize images to 224 × 224 , the training input size of ImageNet. We run all experiments for three seeds using the FLSim library described in Nguyen et al. [2023].

Baseline methods We compare our methods to the following approaches as per Nguyen et al. [2023]: (a) Random : 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 a standard FL optimization algorithm is applied. (d) LP and FT Oracles : These are equivalent baselines trained in the centralized setting that provide an upper bound to the expected performance.

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 FedAdam which have been shown to do well in these settings in prior art [Nguyen 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-10 and SqueezeNet experiments we take the hyperparameters already derived in Nguyen et al. [2023]. Additional details of selected hyperparameters are provided in Appendix B.

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

Table 2: Accuracy, total computation and total communication costs of pure HeadTuning methods (below dashed lines) and their full training counterparts. CIFAR-10 × 32 indicates CIFAR-10 without re-sizing samples to 224 × 224 , we include this setting so we can compare directly with values for FT, Random and LP reported by Nguyen et 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 a 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, LP has much lower communication costs when compared to FT for any given number of rounds. Our proposed FedNCM is a one-round algorithm and therefore has even lower communication costs than any other algorithm considered.

For computation time we consider the total FLOPs executed on the clients. We assume for simplicity that the backward pass of a model is 2 × 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 of Pure HeadTuning for FL

As discussed in Sec. 1 tuning the classifier head is at times as effective or more effective than updating the entire model in the context of transfer learning Evci et al. [2022]. In prior work, this situation was briefly considered as a limited case in Nguyen et al. [2023, Appendix C.2] for CIFAR-10 and suggested that tuning just the linear head (LP) might be a weak approach in the heterogeneous 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, we 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 × 32 input image directly into a pre-trained ImageNet model. Since the architectures are adapted to the 224 × 224 size and trained at this scale originally, such an approach can lead to a 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 [Kornblith et al., 2019, Evci et al., 2022].

Tab. 2 shows accuracy, compute, and communication cost results for Pure HeadTuning Methods (FedNCM and LP) as well as full tuning approaches including our FedNCM+FT. We note that in Tab. 2, CIFAR-1032 × 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-train FT setting and furthermore, LP 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 Flowers102 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 FineTune

We now study in more detail the two-stage approach described in Sec. 3.3. Fig. 1 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 of 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 a strong starting accuracy and in stage two, it converges with a better accuracy than FT given the same computation budget. We note that FedNCM+FT for Cars and EuroSAT-Sub sees an initial performance decrease during the first few rounds of 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 an initial warm-up period at a lower learning rate.

FedNCM+FT converges rapidly, allowing it to be highly efficient under most communication budgets compared to other methods. The second column of Tab. 2 shows 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 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 to 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 convergence criteria, thus for the purpose of 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.

0

25

50

75

100

125

150

175

200

Rounds

0

20

40

60

80

100

FedNCM:

64

.

8

±

0

.

15

LP oracle

FT oracle

NCM+FT:

87

.

2

±

0

.

2

LP:

82

.

5

±

0

.

2

FT:

85

.

4

±

0

.

4

Random:

67

.

8

±

0

.

6

10

3

10

4

10

5

10

6

Total Communication (Mb)

0

10

20

30

40

50

60

70

80

90

Accuracy vs Communication Budget Cifar

LP

FT

FedNCM+FT

Random

10

1

10

2

10

3

10

4

Total Compute (GFLOPs)

0

10

20

30

40

50

60

70

80

90

Accuracy vs Compute Budget Cifar

LP

FT

FedNCM+FT

Random

0

100

200

300

400

500

Rounds

0

20

40

60

80

100

FedNCM:

71

.

8

±

0

.

03

LP oracle

FT oracle

NCM+FT:

74

.

9

±

0

.

2

LP:

74

.

1

±

1

.

2

FT:

64

.

5

±

1

.

0

Random:

33

.

2

±

0

.

7

10

3

10

4

10

5

10

6

Total Communication (Mb)

0

10

20

30

40

50

60

70

80

Accuracy vs Communication Budget Flowers

LP

FT

FedNCM+FT

Random

10

1

10

2

10

3

10

4

Total Compute (GFLOPs)

0

10

20

30

40

50

60

70

80

Accuracy vs Compute Budget Flowers

LP

FT

FedNCM+FT

Random

0

200

400

600

800

1000

1200

1400

Rounds

0

10

20

30

40

50

60

70

80

FedNCM:

37

.

9

±

0

.

16

LP oracle

FT oracle

NCM+FT:

55

.

0

±

0

.

3

LP:

50

.

0

±

0

.

3

FT:

52

.

0

±

0

.

9

Random:

15

.

0

±

0

.

7

10

3

10

4

10

5

10

6

Total Communication (Mb)

0

10

20

30

40

50

60

Accuracy vs Communication Budget CUB

LP

FT

FedNCM+FT

Random

10

1

10

2

10

3

10

4

Total Compute (GFLOPs)

0

10

20

30

40

50

60

Accuracy vs Compute Budget CUB

LP

FT

FedNCM+FT

Random

0

500

1000

1500

2000

2500

Rounds

0

10

20

30

40

50

60

FedNCM:

20

.

3

±

0

.

04

LP:

41

.

2

±

0

.

5

FT:

48

.

7

±

2

.

0

Random:

5

.

6

±

0

.

8

NCM+FT:

54

.

8

±

1

.

2

10

3

10

4

10

5

10

6

Total Communication (Mb)

0

10

20

30

40

50

Accuracy vs Communication Budget Cars

LP

FT

FedNCM+FT

Random

10

1

10

2

10

3

10

4

Total Compute (GFLOPs)

0

10

20

30

40

50

Accuracy vs Compute Budget Cars

LP

FT

FedNCM+FT

Random

0

200

400

600

800

1000

Rounds

0

20

40

60

80

100

FedNCM:

81

.

8

±

0

.

64

LP oracle

FT oracle

NCM+FT:

96

.

0

±

0

.

5

LP:

92

.

8

±

0

.

3

FT:

95

.

6

±

1

.

0

Random:

84

.

5

±

2

.

1

10

3

10

4

10

5

10

6

Total Communication (Mb)

0

20

40

60

80

100

Accuracy vs Communication Budget Eurosat

LP

FT

FedNCM+FT

Random

10

1

10

2

10

3

10

4

Total Compute (GFLOPs)

0

20

40

60

80

100

Accuracy vs Compute Budget Eurosat

LP

FT

FedNCM+FT

Random

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 et al., 2015]. Fig. 2 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 as a stand alone method, reinforcing 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 in communication cost between one LP or FedNCM round and an FT round. Additionally, we find that FedNCM's accuracy is robust to sample size decreases and FedNCM+FT 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 Cifar-10.

100

100

100

08

80

08

60

60

60

40

40

40

20

20

20

FT 85.2±2.7

LP 66.4±11.5

NCM+FT 87.4±1.0

NCM 79.6±0.9

FT 84.2±0.7

LP 33.5±3.6

NCM 81.0±1.4

NCM+FT 84.0±2.2

FT 81.1±3.4

LP 42.2±3.4

NCM 78.3±4.9

NCM+FT 84.4±5.7

0

0

50

100

150

200

250

300

0

0

50

100

150

200

250

300

0

0

50

100

150

200

250

300

Rounds

Rounds

Rounds

4.3 Analysis and Ablations

We now focus on demonstrating other advantages of our two-stage method, with a focus on using FedNCM as the method of choice for stage one (FedNCM+FT). In particular, we investigate robustness to larger number of clients, insensitivity to hyperparameters and compatibility with multiple FL algorithms and architectures. In Appendix C 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 we 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.1 summarizes the results of using each method with 1.) FedNCM+FT 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 systematically exceeds the performance obtained when only fine-tuning. Regardless of the federated algorithm used, FedNCM alone continues to exceed the performance of FT on Flowers102. 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 a grid search over a range of server and client learning rates for FedAdam using both FT and FedNCM+FT. We observe that FedNCM+FT not only has higher performance but it is also more stable over the entire hyperparameter grid on Flowers dataset, and outperforms for all settings on CIFAR-10.

100

0.1

0.01

Alpha

76

78

80

82

84

86

88

90

88.3

81.8

89.2

85.3

79.4

87.1

80.2

77.0

84.3

Varying Heterogenity CIFAR-10

FT

LP

FedNCM + FT

100

0.1

0.01

Alpha

55.0

57.5

60.0

62.5

65.0

67.5

70.0

72.5

75.0

64.6

75

74.2

64.5

74.9

74.1

56.3

74.9

73.7

Varying Heterogenity Flowers102

FT

LP

FedNCM + FT

1

3

10

Local Epochs

79.0

80.0

81.0

82.0

83.0

84.0

85.0

86.0

87.0

85.3

79.3

87.1

86.1

81.5

87.2

86.7

82.8

87.2

Varying Local Epochs CIFAR-10

FT

LP

FedNCM + FT

1

3

10

Local Epochs

64.0

66.0

68.0

70.0

72.0

74.0

64.5

74

74.9

65.8

74.9

75.3

64.5

75.1

75

Varying Local Epochs FLOWERS

FT

LP

FedNCM + FT

1

3

10

Local Epochs

92.5

93.0

93.5

94.0

94.5

95.0

95.5

96.0

95

93

96

95.1

92.6

94.6

92.7

92.5

94.8

Varying Local Epochs EUROSAT

FT

LP

FedNCM + FT

1

3

10

Local Epochs

42.0

44.0

46.0

48.0

50.0

52.0

54.0

48.7

41.1

54.6

48.4

41.5

54.4

45.7

42.2

53.8

Varying Local Epochs CARS

FT

LP

FedNCM + FT

1

0.1

0.01

0.001

0.0001

Client Learning Rate

0.001

0.0001

0.00001

0.000001

Server Learning Rate

10.00

10.00

10.00

79.98

47.35

10.00

10.00

10.00

88.00

82.25

10.00

10.00

10.00

84.68

81.90

10.00

10.00

10.00

67.24

51.68

Tuning FedAdam - FT

(CIFAR-10)

10

20

30

40

50

60

70

80

1

0.1

0.01

0.001

0.0001

Client Learning Rate

0.001

0.0001

0.00001

0.000001

Server Learning Rate

10.00

10.00

10.00

84.28

50.04

10.00

10.00

10.00

89.46

84.60

10.00

10.00

10.00

88.12

84.65

10.00

10.00

10.00

79.93

75.10

(CIFAR-10)

10

20

30

40

50

60

70

80

0.1

0.01

0.001

0.0001

0.00001

Client Learning Rate

0.001

0.0001

0.00001

0.000001

Server Learning Rate

22.82

21.56

30.15

5.71

4.15

54.29

54.04

59.49

66.63

61.80

47.13

47.78

49.76

44.28

5.53

1.40

1.15

1.27

1.30

0.33

Tuning FedAdam - FT

(Flowers102)

10

20

30

40

50

60

0.1

0.01

0.001

0.0001

0.00001

Client Learning Rate

0.001

0.0001

0.00001

0.000001

Server Learning Rate

3.01

15.43

21.16

18.43

64.56

75.85

76.50

75.85

74.99

73.26

74.63

74.66

74.91

74.40

73.48

74.68

74.45

74.26

73.78

72.43

(Flowers102)

10

20

30

40

50

60

70

Heterogeneity Nguyen et al. [2023] 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 ( α = 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. 3. Here we consider for CIFAR-10 the nearly iid α = 100 , α = 0 . 1 as considered in Nguyen et al. [2023], and a very heterogeneous α = 0 . 01 . Firstly, we observe that FedNCM+FT can provide benefits in the iid setting. As heterogeneity degrades the naive FT setting sees a 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, typically a larger amount of local computation between rounds is desired to minimize communication.

However, this can often come at a cost of degraded performance. We observe in Fig. 4 as in Nguyen et al. [2023] that FT can be relatively robust in some cases (CIFAR10) to increasing local epochs. However, we also observe for some datasets that it can degrade, while LP and FedNCM+FT are less likely to degrade. Overall FedNCM+FT continues to outperform for larger local epochs.

Increasing clients Tab. 6 shows that as we increase the number of clients we observe that the degradation of FedNCM+FT is less severe than both LP and FT, suggesting it is stable under a large number of workers being averaged. As discussed in Sec. 3.3 it is expected in the same round that a 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 a round to be less given a fixed update budget.

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.

100

500

5000

70

75

80

85

90

Accuracy for Varying Number of Clients

FT

LP

NCM + FT

References

Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. arXiv preprint arXiv:2111.04263 , 2021.

Zaid Alyafeai, Maged Saeed 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-tuning with differential privacy necessitates an additional hyperparameter search. arXiv preprint arXiv:2210.02156 , 2022.

Hong-You Chen, Cheng-Hao 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 representation forgetting in supervised and unsupervised continual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16712-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 60096033. PMLR, 2022.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 35573568. Curran Associates, Inc., 2020.

Ross Girshick, Jeff Donahue, 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, 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 IEEE/CVF International Conference on Computer Vision (ICCV) , October 2019.

Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning , pages 2790-2799. PMLR, 2019.

Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335 , 2019.

Forrest N. Iandola, 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 tangent kernel: Convergence and generalization in 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.04428 , 2022.

Sai Praneeth 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ˇ cn` y, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Distributed machine learning for on-device intelligence. arXiv preprint arXiv:1610.02527 , 2016.

Simon Kornblith, Jonathon Shlens, and Quoc 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: //www.cs.toronto.edu/~kriz/learning-features-2009-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 , 40(12):2935-2947, 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, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id= XtyeppctGgc .

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communicationefficient 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 begin? 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. IEEE Signal Processing Magazine , 32(3):53-69, 2015. doi: 10.1109/MSP.2014.2347059.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Roger Ratcliff. 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 Pattern Recognition , pages 2001-2010, 2017.

Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koneˇ cn` y, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295 , 2020.

Yi Ren, Shangmin Guo, Wonho Bae, and Danica J. Sutherland. How to prepare your task head for finetuning. In The Eleventh International Conference on Learning Representations , 2023a. URL https://openreview. net/forum?id=gVOXZproe-e .

Yi Ren, Shangmin Guo, Wonho Bae, and Danica J Sutherland. How to prepare your task head for finetuning. arXiv preprint arXiv:2302.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 NeurIPS , 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. URL https://openreview.net/ forum?id=9aokcgBVIj1 .

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 communication-efficient sgd algorithms. arXiv preprint arXiv:1808.07576 , 2018.

Karl Weiss, Taghi M Khoshgoftaar, and DingDing 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. In Proceedings of 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-5700, 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.

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 LL in and LL out are the input and output dimensions of the linear layer, R is the number of federated rounds, K is the total number of clients and F is the fraction F ∈ (0 , 1] of clients participating in each federated round:

For methods training the complete model, communication is calculated according to equation 4 where P is the number of parameters present in the base model (excluding the linear layer) and LL in , LL out , R , K and F have the same meanings as above:

A.2 Computation

For computation, let F be one forward pass of a single sample, Let S be the subset of clients, K , selected to participate in each federated round, R . We define one backward pass as 2 F therefore each sample contributes 3 F per round in a full fine-tuning situation. For HeadTuning a complete backward pass is not necessary since we only update the linear layer weights and each sample contributes F per round. Eq. 5 and Eq. 6 formalize our computation calculation method in units of F, where E is the number of local epochs performed by each client and N s is the sum of the number of samples at each selected client, S .

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 , 0 . 005 , 0 . 003 , 0 . 001 } . For Flowers102, based on preliminary 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 CIFAR10 and Flowers the client learning rate was tuned over { 1 , 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.

C Additional HeadTuning Analysis

C.1 MLP HeadTuning

We compare HeadTuning methods by replacing the linear layer in the LP method with a three layer MLP. Tab. 5 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. 6 indicates performance of our two stage HeadTuning + fine-tuning method using five rounds of LP, ten rounds of LP and FedNCM as the HeadTuning methods. The bottom row in each dataset category where stage 1 Methods is denoted as n/a, is the result of only fine-tuning, i.e. not using our two stage method, which shows that HeadTuning prior to fine-tuning is always at least as effective as FT alone and, depending on the HeadTuning method selected, much more effective.

D Experiments Using Small Subsets of Data

We conduct experiments 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 α = 0 . 1 so each client will have very few samples and some will be entirely missing samples from many classes. Tab. 7 shows 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.

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 and Flowers102 datasets. Results are summarized in Tab. 8 where we observe that FedNCM performance is better by almost 13% compared to squeezenet, while FT performance is degraded compared to squeezenet. We hypothesis this is due to the challenges of deeper networks in heterogenous federated learning. For the Flowers102 dataset, FedNCM and

Table 8: ResNet18 model performance for FedAvg. As with Squeezenet, FedNCM+FT continues to outperforms in all cases.

FedNCM+FT produce the best results by 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.

F Extended Accuracy Comparison Figures

Fig. 7 is the extended version of Fig. 1 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. 7 shows these same figures with the entirety of the training rounds displayed for the random setting.

G Compute

We use a combination of NVIDIA A100-SXM4-40GB, NVIDIA RTX A4500, Tesla V100-SXM232GB and Tesla P100-PCIE-12GB GPUs for a total of 1.1 GPU years . In addition to the experiments reported in the paper, this includes preliminary experiments and hyperparameter searches.

0

250

500

750

1000

1250

1500

1750

2000

Rounds

0

20

40

60

80

100

FedNCM:

64

.

8

±

0

.

15

LP oracle

FT oracle

NCM+FT:

87

.

2

±

0

.

2

LP:

82

.

5

±

0

.

2

FT:

85

.

4

±

0

.

4

Random:

67

.

8

±

0

.

6

0

500

1000

1500

2000

2500

3000

Rounds

0

20

40

60

80

100

FedNCM:

71

.

8

±

0

.

03

LP oracle

FT oracle

NCM+FT:

74

.

9

±

0

.

2

LP:

74

.

1

±

1

.

2

FT:

64

.

5

±

1

.

0

Random:

33

.

2

±

0

.

7

0

500

1000

1500

2000

2500

3000

Rounds

0

10

20

30

40

50

60

70

80

FedNCM:

37

.

9

±

0

.

16

LP oracle

FT oracle

NCM+FT:

55

.

0

±

0

.

3

LP:

50

.

0

±

0

.

3

FT:

52

.

0

±

0

.

9

Random:

15

.

0

±

0

.

7

0

500

1000

1500

2000

2500

3000

3500

Rounds

0

10

20

30

40

50

60

FedNCM:

20

.

3

±

0

.

04

LP:

41

.

2

±

0

.

5

FT:

48

.

7

±

2

.

0

Random:

5

.

6

±

0

.

8

NCM+FT:

54

.

8

±

1

.

2

0

500

1000

1500

2000

2500

3000

3500

Rounds

0

20

40

60

80

100

FedNCM:

81

.

8

±

0

.

64

LP oracle

FT oracle

NCM+FT:

96

.

0

±

0

.

5

LP:

92

.

8

±

0

.

3

FT:

95

.

6

±

1

.

0

Random:

84

.

5

±

2

.

1

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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).\n\nWEAKNESSES:\n1) 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\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\nEvaluation for heterogeneity is done only with very particular non-iid with respect to labels, but it is not highlighted.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nIn 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.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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).\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}

comments: 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?

LIMITATIONS:
Please see above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-08-01