SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 \times 32$  input image directly into a pre-trained ImageNet model. Since the architectures are adapted to the  $224 \times 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-10-32 × 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.
Figure 1: A comparison of the accuracy between 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 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.
Figure 2: A comparison of the accuracy between models initialized with pre-trained trained using LP, FT and FedNCM+FT. We remark that the FedNCM alone obtains strong performance.
Dataset	Algorithm	FedNCM	FedNCM + FT	FT
CIFAR-10	FEDAVG	64.8 ± 0.1	87.2 ± 0.2	85.4 ± 0.4
FEDPROX	64.8 ± 0.1	88.1 ± 0.1	87.8 ± 0.08
FEDADAM	64.8 ± 0.1	89.4 ± 1.1	88.2 ± 0.2
FLOWERS102	FEDAVG	71.8 ± 0.03	74.9 ± 0.2	64.5 ± 1.0
FEDPROX	71.8 ± 0.03	75.2 ± 0.1	65.5 ± 1.5
FEDADAM	71.8 ± 0.03	76.7 ± 0.2	66.6 ± 1.0
Table 3: Model performance with different methods for a variety of 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 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.
Figure 3: We vary the heterogeneity (Dirichlet-  $\alpha$ ) for CIFAR-10 and Flowers102. Methods with HeadTuning: LP and FedNCM+FT are more robust, with the substantial advantage of FedNCM + FT increasing in challenging higher heterogeneity.
Figure 4: We vary the number of local epochs. FedNCM+FT always outperforms FT and nearly always LP in this challenging setting.
Figure 5: Hyperparameter grids for FedAdam for CIFAR-10 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 a wide range, a noticeable advantage in FL
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  $(\alpha = 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  $\alpha = 100$ ,  $\alpha = 0.1$  as considered in Nguyen et al. [2023], and a very heterogeneous  $\alpha = 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 agorthms, 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 (CIFAR-10) 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 Fed-NCM+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.
Figure 6: We increase the number of clients on CIFAR-10. FedNCM+FT degrades more gracefully than FT and LP.
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 focuses 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.
This research was partially funded by NSERC Discovery Grant RGPIN- 2021-04104, FRQNT New Researcher Grant and CERC Autonomous AI. E.O. acknowledges support from ANR ADONIS - ANR-21-CE23-0030. We acknowledge resources provided by Compute Canada and Calcul Quebec.
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.
Dataset	FT+FedNCM	FT	LP	Random
CIFAR-10	200	200	200	3000
CIFAR-10 32 × 32NGUYEN ET AL. [2023]	1000	500	1000	1000
FLOWERS-102	250	250	500	3000
STANFORD CARS	1000	1000	2500	3500
CUB	1000	1000	1500	3000
EUROSAT-SUB	1000	1000	1000	3500
Table 4: rounds conducted for each dataset and method combination. For CIFAR-10  $32 \times 32$  experiments were conducted in Nguyen et al. [2023] with the exception of FedNCM+FT (our method)
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 \in (0,1]$  of clients participating in each federated round:
$$
c o m m _ {l p} = 2 * \left(L L _ {\text {o u t}} * L L _ {\text {i n}}\right) * K * R * F * (3 2 \text {b i t}) \tag {3}
$$
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:
$$
\operatorname {c o m m} _ {f t} = 2 \left(\left(L L _ {\text {o u t}} * L L _ {\text {i n}}\right) + P\right) * K * R * F * (3 2 \text {b i t}) \tag {4}
$$
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  $2F$  therefore each sample contributes  $3F$  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$ .
$$
\operatorname {c o m p} _ {f t} = 3 * R * E * N _ {s} \tag {5}
$$
$$
c o m p _ {f t} = R * E * N _ {s} \tag {6}
$$
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 CIFAR-10 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.
Dataset	HeadTuning Method	Accuracy
CIFAR-10	LP	82.5 ± 0.2
MLP	84.1 ± 0.2
FLOWERS102	LP	74.1 ± 1.2
MLP	72.5 ± 0.3
CUB	LP	50.0 ± 0.3
MLP	50.3 ± 0.3
CARS	LP	41.2 ± 0.5
MLP	41.0 ± 0.4
EUROSAT-SUB	LP	92.6 ± 0.4
MLP	91.7 ± 0.8
Table 5: LP and MLP HeadTuning results prior to the fine-tuning stage.
Dataset	stage 1 Method	Accuracy
CIFAR-10	LP (5 ROUND)	85.9 ± 0.4
LP (10 ROUND)	84.5 ± 0.17
FEDNCM	87.2 ± 0.2
N/A	85.4 ± 0.4
FLOWERS102	LP (5 ROUND)	68.6 ± 1.3
LP (10 ROUND)	68.3 ± 0.6
FEDNCM	74.9 ± 0.2
N/A	64.5 ± 0.1
Table 6: Outcomes using LP as the HeadTuning method in our two stage training algorithm, FedNCM results from Tab. 2 are included for comparison.
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  $\alpha = 0.1$  so each client will have very few samples and some will be entirely missing samples from many classes. Tab. 7 shows a  $\geq 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
LP	FT	FedNCM	FedNCM + FT
53.8 ± 0.3	45.2 ± 2.4	55.6 ± 0.8	56.7 ± 0.7
Table 7: ResNet18 model performance for FedAvg. As with SqueezeNet, FedNCM+FT continues to outperforms in all cases.
Dataset	FedNCM	FedNCM + FT	FT+Pretrain	LP+Pretrain
CIFAR-10	77.74 ± 0.05	79.05 ± 1.31	77.87 ± 4.07	74.73 ± 3.03
FLOWERS102	-74.13 ± 0.31	-74.1 ± 0.26	-34.41 ± 10.16	-25.35 ± 2.59
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-SXM2-32GB 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.
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.