SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Gwen Legate
Concordia University, Mila
gwendolyne.legate@mila.quebec
Nicolas Bernier
Lucas Caccia
McGill University, Mila
Edouard Oyallon
Sorbonne University, ISIR, CNRS
Eugene Belilovsky
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<sup>1</sup>.
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
$^{1}$ https://github.com/GwenLegate/GuidingLastLayerFLPretrain
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 (Head Tuning) in our proposed two stage process which leads to improved accuracy.
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.
In FL, distributed optimization occurs over  $K$  clients with each client  $k \in \{1,\dots,K\}$  having data  $\mathbf{X}_k, \mathbf{Y}_k$  that contains  $n_k$  samples drawn from distribution  $D_{k}$ . We define the total number of samples across all clients as  $n = \sum_{k=1}^{K} n_k$ . The data  $\mathbf{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čný et al., 2016] and aims to find the minimizer of the loss over the sum of the client data:
$$
\mathbf {w} ^ {*}, \mathbf {v} ^ {*} \in \underset {\mathbf {w}, \mathbf {v}} {\arg \min } \sum_ {k = 1} ^ {K} \frac {n _ {k}}{n} \mathcal {L} \left(g \left(f \left(\mathbf {w}, \mathbf {X} _ {k}\right), \mathbf {v}\right)\right). \tag {1}
$$
Here we have split the model prediction into  $f$ , a base parameterized by  $\mathbf{w}$  that produces representations, and  $g$ , a task head parameterized by  $\mathbf{v}$ . In this work we will focus on the case where the task head is a linear model, and the loss function,  $\mathcal{L}$  represents a standard classification or regression loss. The  $\mathbf{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  $\mathbf{w}$  is to optimize only  $\mathbf{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  $\mathcal{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:  $(\mathbf{X}_1,\mathbf{Y}_1),(\mathbf{X}_2,\mathbf{Y}_2),\ldots ,(\mathbf{X}_K,\mathbf{Y}_K)$  - Local datasets,  $w_{pt}$  - pre-trained model Server Executes: 1: for each client  $k\in K$  in parallel do 2:  $[m_c^k ]_{c\in C}\gets \mathrm{LocalClientStats}(X_k,Y_k,w_{pt})$ $\triangleright$  Send to all clients, receive weighted class means 3: end for 4: for each class  $c\in C$  do 5:  $\mathbf{l}_c\gets \frac{1}{D_c}\sum_{k = 1}^{K}m_c^k$ $\triangleright \mathbf{l}_c$  can be used in NCM classifier 6: end for Client Side: 7: function LOCALCLIENTSTATS(X,Y,w) 8: for each class  $c\in N$  do 9: Let  $\mathbf{X}_c = \{x_i\in X,y_i = c\}$  10:  $m_{c}\gets \sum_{x\in X_{c}}f_{w}(x)$  11: end for 12: return  $[m_c]_{c\in C}$  13: end function
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\circ f)(\mathbf{x}) = \mathbf{v}f(\mathbf{x};\mathbf{w}) + \mathbf{b}$ . We can set the matrix  $v$  corresponding to the class  $c$  logit with the normalized class centroid  $\mathbf{l}_c / \| \mathbf{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  $\mathbf{w}_0$ , and trained locally with SGD for several steps until it reaches the parameter  $\mathbf{w}_k$ . Writing  $\mathbf{w}^*$  the optimal parameter, via triangular inequality, we obtain the following inequality:
$$
\mathbb {E} _ {\mathbf {X} _ {k}} \left[ \| f \left(\mathbf {w} _ {k}; \mathbf {X} _ {k}\right) - f \left(\mathbf {w} ^ {*}; \mathbf {X} _ {k}\right) \| \right] \leq \mathbb {E} _ {\mathbf {X} _ {k}} \left[ \| f \left(\mathbf {w} _ {0}; \mathbf {X} _ {k}\right) - f \left(\mathbf {w} ^ {*}; \mathbf {X} _ {k}\right) \| + \| f \left(\mathbf {w} _ {k}; \mathbf {X} _ {k}\right) - f \left(\mathbf {w} _ {0}; \mathbf {X} _ {k}\right) \| \right]. \tag {2}
$$
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  $\sigma$  is the sigmoid activation and  $\{\mathbf{e}_i\}_i$  the canonical basis) for some  $c > 0$ , by:
$$
\mathbb {E} _ {\mathbf {X} _ {k}} \| f (\mathbf {w} _ {k}; \mathbf {X} _ {k}) - f (\mathbf {w} _ {0}; \mathbf {X} _ {k}) \| \leq c \cdot \mathbb {E} _ {(\mathbf {X} _ {k}, \mathbf {Y} _ {k})} \| \mathbf {e} _ {\mathbf {Y} _ {k}} - g _ {\mathbf {V}} (f (\mathbf {w} _ {0}; \mathbf {X} _ {k})) \|.
$$
This suggests in particular that a good choice of linear head  $\mathbf{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  $\mathbf{v}$  (compared to a random  $\mathbf{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  $\mathbf{v}$  it initialized by HeadTuning, compared to a random initialization.
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  $\alpha = 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.
Dataset	Num. Classes	Num. Images
CIFAR-10	10	50000
Flowers102	102	1020
Stanford Cars	196	8144
CUB	200	5994
EuroSAT-Sub	10	5000
Table 1: Summary of datasets used in our experiments.
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 fine-tuning and evaluation for all datasets, we resize images to  $224 \times 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
Dataset	Method	Accuracy (%)	Total Compute (×F)	Total Comm. (GB)
CIFAR-10	Random	67.8 ± 0.6	4.5 × 108	1803.71
FT	85.4 ± 0.4	3.0 × 107	120.25
FedNCM+FT	87.2 ± 0.2	3.0 × 107	120.25
LP	82.5 ± 0.2	1.0 × 106	0.82
FedNCM	64.8 ± 0.1	1	4.1 × 10-3
CIFAR-10 × 32	Random (Nguyen et al. [2023])	34.2	1.5 × 108	601.24
FT (Nguyen et al. [2023])	63.1	1.5 × 108	601.24
FedNCM+FT	67.9 ± 0.4	7.5 × 107	300.62
LP (Nguyen et al. [2023])	44.7	5.0 × 107	4.10
FedNCM	40.02 ± 0.04	1	4.10 × 10-3
FLOWERS-102	Random	33.2 ± 0.7	9.2 × 106	1916.76
FT	64.5 ± 1.0	7.7 × 105	159.73
FedNCM+FT	74.9 ± 0.2	7.7 × 105	159.73
LP	74.1 ± 1.2	5.1 × 105	20.93
FedNCM	71.8 ± 0.03	1	4.2 × 10-2
CUB	Random	15.0 ± 0.7	5.4 × 107	2037.18
FT	52.0 ± 0.9	1.8 × 107	679.06
FedNCM+FT	55.0 ± 0.3	1.8 × 107	679.06
LP	50.0 ± 0.3	9.0 × 106	122.88
FedNCM	37.9 ± 0.2	1	8.2 × 10-2
STANFORD CARS	Random	5.6 ± 0.8	8.6 × 107	2370.97
FT	48.7 ± 2.0	2.4 × 107	677.42
FedNCM+FT	54.8 ± 1.2	2.4 × 107	677.42
LP	41.2 ± 0.5	2.0 × 107	200.70
FedNCM	20.33 ± 0.04	1	8.0 × 10-2
EUROSAT-SUB	Random	85.8 ± 2.7	5.3 × 107	2104.32
FT	95.6 ± 1.0	1.5 × 107	601.24
FedNCM+FT	96.0 ± 0.5	1.5 × 107	601.24
LP	92.6 ± 0.4	5.0 × 106	4.10
FedNCM	81.8 ± 0.6	1	4.10 × 10-3
Table 2: Accuracy, total computation and total communication costs of pure HeadTuning methods (below dashed lines) and their full training counterparts. CIFAR-10  $\times$ 32 indicates CIFAR-10 without re-sizing samples to  $224\times 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 \times$  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