Where to Begin? On the Impact of Pre-Training and Initialization in Federated Learning
Authors: John Nguyen
Jianyu Wang
Kshitiz Malik
Maziar Sanjabi
Michael Rabbat
Venue: ICLR 2023 notable top 25% (2023)
Keywords: federated learning
optimization

Abstract: An oft-cited challenge of federated learning is the presence of heterogeneity. \emph{Data heterogeneity} refers to the fact that data from different clients may follow very different distributions. \emph{System heterogeneity} refers to client devices having different system capabilities. A considerable number of federated optimization methods address this challenge. In the literature, empirical evaluations usually start federated training from random initialization. However, in many practical applications of federated learning, the server has access to proxy data for the training task that can be used to pre-train a model before starting federated training. Using four standard federated learning benchmark datasets, we empirically study the impact of starting from a pre-trained model in federated learning. Unsurprisingly, starting from a pre-trained model reduces the training time required to reach a target error rate and enables the training of more accurate models (up to 40\%) than is possible when starting from random initialization. Surprisingly, we also find that starting federated learning from a pre-trained initialization reduces the effect of both data and system heterogeneity. We recommend future work proposing and evaluating federated optimization methods to evaluate the performance when starting from random and pre-trained initializations. This study raises several questions for further work on understanding the role of heterogeneity in federated optimization.

The predominant approach to federated training builds on local update methods such as FE- DAvG (McMahan et al., 2016), where a device performs several local updates (e.g., one epoch of SGD their local training set) before transmitting communication overhead, it can also exacerbate data heterogeneity. Several approaches have been Published as a conference paper at ICLR 2023
WHERE TO BEGIN? ON THE IMPACT OF PRE- AND INITIALIZATION IN FEDERATED LEARNING
John Nguyen
Jianyu Wang
Kshitiz Malik
Maziar Sanjabi
Michael Rabbat
Meta AI
(ng jnn, ji piayuwang,, kmal ik2, maziara, weriiirrrmmikeebbbtt.... e @meta.........mmm........... com
ABSTRACT
An oft-cited challenge of federated learning the presence heterogeneity. Data heterogeneity refers the fact that data from different clients may follow very different distributions. System heterogeneity refers to client devices having different system capabilities. considerable number of federated optimization methods address this challenge. In the literature, empirical evaluations usually start federated training from random initialization. However, in many practical applications of federated learning, the server has access to proxy data for the training task that can used to pre-train a model before starting federated training. Using four standard federated learning benchmark datasets, we empirically study the impact of starting from a pre-trained model federated learning. Unsurprisingly, starting from a pre- tre-trained model reduces the training time required to reach target error rate and enables the training of more accurate models (up to 40%) than is possible when starting from random initialization. Surprisingly, we also find that starting federated learning from a pre-trained initialization reduces the effect of both data and system heterogeneity. We recommend future work proposing and evaluating federated optimization methods to evaluate the performance when starting from random and pre-trained initializations. This study raises several questions for further work on understanding the role of heterogeneity in federated optimization.
INTRODUCTION
Federated learning (FL) has emerged as a popular distributed machine learning paradigm for privately training shared model across many participants. At the same time, the training data never leaves the participant's devices. This paper empirically investigates the impact model initialization on federated optimization methods. Previous empirical evaluations of FL methods start federated training from a randomly initialized model. Transfer learning from pre-trained models has become common practice in natural language processing Radford al. (2019); Devlin al. (2018) and computer vision He et al. (2019); Dosovitskiy et al. (2020), yielding state-of-the-an results on many tasks and enabling faster model convergence the centralized setting. Although public proxy data is available at the server in many applications, few prior works studied the impact of starting federated training from a pre-trainee pre-trained model.
In cross-devin FL (Kairouz et al., 2019), the primary setting considered in this paper, central server coordinates many client devices (possibly in hundreds of millions). Each device possesses local dataset, and the data at different devices follow different distributions, leading to the data heterogeneity challenge (Kairouz et al., 2019). Moreover, client devices have different system capabilities, leading to system heterogeneity. Finally, devices communicate with the server over low-bandwidth communication links making the performance bottleneck.
update the server. Although this reduces
proposed to address this challenge (Li et al., 2018; Hsu et al., 2019; Reddi al., 2020; Wang et al.,

---

Published as a conference paper at ICLR 2023
CIFAR-10
Overtiow
20.0
1500
18.0
400
17.0
650
26.0
800
Foldan Froonil
Fadreg Froae
Fategee Friiinl
Filhog Froinal
Fedhoghh Mortc
Mec
Preimined
FEWNIST
Randam
Pre-trlnsd
Reddit
Random
23
24.0
05.0
760
Feciig
Fadenges Poorra
Figure 1: We tested the accuracy of four datasets using random and pre-trained weights. We used solid lines for SGD, dashed lines for PROXIMAL, and dotted lines for MIMELITE. The graph shows how the algorithm rankings changed between the random and pre-trained models. Although no single method is the best for all tasks, FEDADAM with SGD for CLIENTOPT performed consistently well when starting from pre-trained model, especially for the larger language modeling tasks of Stack Overflow and Reddit.
2020; Karimireddy et al., 2020; 2021; Zhang et al., 2021; Acar et al., 2021). However, few prior works examine the impact of initialization on federated training.
Contributions. In this work, consider the question: How does model initialization (random or pre-trained) impact the behavior of federated optimization methods? ? We perform an extensive empirical study, comparing 15 variations of federated optimization methods four commonly-used FL benchmark datasets. Our study reveals three key findings:
1. Although optimizers designed address heterogeneity typically lead to better performance when starting from a random initialization, when starting from pre-trained model, we observe that (cf. Fig. 1): (i) there not as big a difference between optimizers terms of accuracy after fixed number of rounds, and (ii) using an adaptive optimizer at the server, such as FEDADAM, more important than using any method for addressing heterogeneity.
2. Starting from pre-trained model significantly reduces the difference between having non- IID vs. IID data for clients. Furthermore, when starting from a pre-trained model, the number of local epochs per round can be significantly increased without degrading the final accuracy.
3. The initial loss is sometimes lower when starting from random model. However, the largest Hessian eigenvalue (i.c., local Lipshitz constant or smoothness) is consistently smaller initialization when starting from pre-trained model compared to when starting from random initialization.
Some of our empirical observations are consistent with existing FL theoretical convergence guarantees. Our findings also highlight that some aspects of FL not captured with the existing theory, suggesting directions for future work.
Initializing FL with pre-trained model can increase final model accuracy and reduce the number of rounds required to achieve a target accuracy. Pre-training leads to communication savings and reduces the overall training time. Figure demonstrates the benefit of pre-training across several datasets thyperparameters were tuned separately for each dataset-initialiiaee initialization pair; see Section 3 for details of the experimental setup).
2

---

Published as a conference paper at ICLR 2023
Algorithm 1 FedOpt framework
1: Input: initial global model x², server and client step sizes 7, f/c. local epochs E, rounds
2: for each round = 1........... do
3: Server sends x' to all clients E St.
4: for each client S in parallel do
5: Initialize local model v + x(-1
6: Each client performs E epochs of local updates via v, f+t1 = CLIENTOPT(Y!. Fe,mc) Let
vf denote the result after performing epochs of local updates.
7: After local training, client sends A = =x---- yy to the server.
8: end for
9: Server computes aggregate update A
37 Eies' P.
Server updates global model x SERVERÖPT(2 at - 1),
11: end for
Our findings are reproducible using the open-source federated learning framework FLSim (FLSim Authors, 2022). Informed by these findings, we present several recommendations for future research on federated optimization.
2 PROBLEM FORMULATION AND THE FEDOPT FRAMEWORK
We consider the following standard optimization formulation of federated training. We seek to find model parameters that solve the problem.
min (w) P. F((w)
WERd
1=1
where m is the total number of clients, the function F. measures the average loss of a model with parameters the ith client's training data, and p; > 0 is the weight given client i. Usually p; is taken to be proportional to the number of samples at client so that the optimization problem gives equal weight all training samples. The goal is to find model that fits all clients' data well on (weighted) average. In FL, only client can evaluate F; and its gradient.
All of the methods we consider in this study can be expressed in the general FEDOPT framework introduced in Reddi et al. (2020); see Algorithm 1. At round I, the server sends its last model x--1 to a cohort of clients. Each client in the cohort performs E' epochs of local training starting from x² using CLIENTOPT with client learning rate T). producing local model yf" Then each client communicates the difference A between their local model and the server model, where A 111 yf. The server computes weighted average A of the client updates (line 9 in Alg. and updates its own model via 1t+1 = SERVEROPT(Ree A., Ma, l), where SERVEROPT(C;; A., Us. !) a first-order optimizer, 77, is the server learning rate, and t the round number.
EXPERIMENTAL SETUP
3.1 DATASETS, MODELS, AND TASKS
We experiment on four FL benchmark datasets: CIFAR-10, 10, FEMNIST, Stack Overflow, and Reddit. All datasets have natural non-IID partitioning of the data except for CIFAR-10, which we partition using Dirichlet distribution with E [0.1, 1.0, 10). We train Squeezenet and ResNet18 models for CIFAR-I0 and FEMNIST, and DistilGPT2 and CharLM models for Stackoverflow and Reddit pushift.io. See Appendix A. for additional information. We run each experiment with three different seeds and report the average. We tune client and server learning rates and ng. and the proximal penalty parameter for FEDPROX with hyperparameter sweep. Each client update runs one local epoch with fixed batch size per task, and we perform 1050 rounds of training for Stackoverflow, 1000 rounds of training for CIFAR-10, and 1082 training rounds for FEMNIST. For additional
3

---

Published as a conference paper at ICLR 2023
implementation details, see Appendix A.2. We use the open-source federated learning simulation framework FLSim FLSim Authors (2022).
3.2 INITIALIZATION STRATEGIES
We consider two initialization strategies: random initialization and supervised pre-training.
Random initialization. Most prior federated optimization works use random weights to initialize the model. We can use the same random initialization strategies used in the standard (centralized) training of deep networks for each model (Iandola et al., 2016; HuggingFace, 2019; He et al., 2016; Kim et al., 2016).
Supervised pre-training. In many FL applications, pre-training be done a large non-pnn non-private proxy dataset available at the server. To facilitate easily reproducing our results, we use publicly available pre-trained models pre-train on public data. For tasks using Squeezenet and ResNet1 18, we use the version of the model pre-trained on ImageNet, available in the Py Torchvision library.' For tasks using DistilGPT2, we use the model weights provided in the HuggingFace library that has been distilled from a pre-trained GPT2, 2 and for tasks using CharLM, we pre-train the model on WikiText- 103 (Merity et al., 2016) (see Appendix B.I for details).
3.3 ALGORITHMS
We compare federated training with five different CLIENTÖPT strategies:
SGD clients perform standard stochastic gradient descent updates; Proximal (Li et al., 2018) clients perform FEDPROX-style local updates; FEDPROX was originally proposed to reduce client drift due to heterogeneity;
Normalized Averaging (Wang et al., 2020) clients use FEDNOVA-Style updates and aggregation to compensate for data imbalance across clients;
MIMELITE (Karimireddy et al., 2021) clients make of optimizer state (e.g., momentum buffer) from the server during local updates to reduce drift due to data heterogeneity:
GD clients perform full-batch gradient updates; in this case, the update A returned to the server is a full-batch gradient client i's local training evaluated at model parameters z'-1
At the server, we consider three strategies for SERVERÖPT. In all strategies, the server treats the averaged update A as a gradient.
SGD the server updates the global model using stochastic gradient descent; when CLIENTÖPT is also SGD, this is equivalent to FEDAVG (McMahan et al., 2016).
SGD with momentum the server updates the global model using SGD with momentum; when CLIENTÖPT is SGD, this equivalent to FEDAVGM (Hsu al., 2019).
Adam the server updates the global model using the Adam optimizer; when CLIENTOPT SGD, this is equivalent to FEDADAM (Reddi et al., 2020).
The method commonly referred to as FEDSGD (MeMahan et al., 2017) is obtained when CLIENTÖPT full-batch gradient descent (GD) and SERVEROPT is SGD, with 7/e and E 1.
We focus on the above choices for CLIENTOPT and SERVEROPT because they are reflective of the most widely-cited federated optimization methods, and they also represent diverse set of possible choices available to the practitioner seeking to deploy cross-device federated training at scale.
3.4 IMPLEMENTATION AND TUNING
We use three different seeds and report the average of each experiment. Hyperparameter tuning is done for each algorithm. model, and dataset, with parameters including client and server learning
httpp:://gitthbb.com//oott/ hhttess//huggiingfaccc.ooo
4

---

Published as a conference paper at ICLR 2023
Stack Overtow
Stack Dwaffow CharMM
Sharle
GFAR-C0 am0
Pro-trilned
Randor
Fundon
Randon
200
403
Roand
Rount
Round
Round
a118-12 0.1 11
PEOMEST
TEMMIET RasNeb-11
40
Po ne-tiined
20
kandee
Aa1sen
510
Bound
Round
Round
Figure 2: While prior works ignore the importance of initialization, using pre- pre-trained models should be the first step for any practical deployment to save on communication bandwidth and achieve high model accuracy. This figure shows the advantage of using pre-trained model for four tasks. For Stack Overflow and Reddit, we use DistilGPT2. For CIFAR- 10 and FEMNIST, we use SqueezeNet.
rates and Tg., and the proximal penalty parameter j for FEDPROX. Unless otherwise specified, we perform one local epoch with fixed batch size per task for each client update. For Stackoverflow, perform 1050 training rounds. For CIFAR-10, 10, we perform 1000 training rounds. For FEMNIST, we perform 1082 training rounds. Our experiments are implemented using the open-source federated learning simulation framework FL.im FL.im Authors (2022). Further implementation details can be found in Appendix A.2.
4 THE IMPACT OF PRE- TRAINING IN FL
In this section, we illustrate the benefits of pre-training in the federated setting and how pre-training can impact federated optimization algorithms behavior.
Pre-training changes the ranking of federated optimization algorithms. If one sorts federated optimization methods based on their performance when starting from a random initialization, the order is substantially different from when using a pre-trained initialization. We focus nine combinations of SERVEROPT and CLIENTOPT, only using local update methods for CLIENTOPT and excluding full-batch gradient descent. We show the change in performance in Figure 1.
First, observe that the span of final accuracies is much smaller when starting from a pre-trained model. Second, all methods starting with pre-trained models achieve a better accuracy after the same number of steps compared to random models. Lastly, observe that the order of methods changes depending on the initialization. Although no particular method dominates across all workloads in Figure FEDADAM with SGD for CLIENTOPT performs consistently well when starting from pre-trained model, especially on the two larger language modeling workloads, Stack Overflow and Reddit, and so we focus on studying FEDADAM-SGD below.
Faster convergence to better accuracy when starting from a pre-trained model. Figure shows that, as one would hope, when starting from a pre-trained model, it is possible to achieve much better accuracy after fixed number steps than when starting federated training from a random initialization. Note that the initial accuracy is not always substantially higher than a random initialization (See Table 2).
Pre-training closes the accuracy gap between non-III and IID. We study how pre-training and data heterogeneity affect convergence without system heterogeneity by fixing the number of local epochs to E, 1. We compare FEDADAM-SGD under IID and Non-IID data splits. In Figure 4, we report the average accuracy for FedAdam (Reddi et al., 2020) on the four datasets. As expected, randomly initialized models perform much worse than their pre-trained counterparts, and IID partitions yield better quality than non-IID. Surprisingly, the gap between models trained on IID data and models trained on non-IID data is significantly smaller when starting with pre-trained weights Moreover,
5

---

Published as a conference paper at ICLR 2023
30
22
92
V1
20
09 50.8
10
Figure 4: The average accuracy on different seeds for FEDADAM trained on IID and non-lID data. For CIFAR-10 Non-III, we generate 100 non-IID clients using Dirichler(O.1). For other three datasets, we use the natural non-IID client partitions.
IR23
125
712
264
553
00.0 903 68.3
62.3
08.0
752
Figure 5: System heterogeneity results comparing FedAvg, FedAvgM and FedAdam with various client optimizers. We simulate system heterogeneity by randomly select 30% of clients per round to perform time-varying local epochs E; ~ (1, the same approach as in Wang et al. (2020). FedProx and FedNova correspond to FedAvg with Proximal client optimizer and normalized averaging (NovA), respectively. We repeat each experiment for 3 different seeds and report the average.
pre-training reduces the negative effects of data heterogeneity (i.e., client drift). As result, observe that (see Figure 3) when training from pre-trained model, increasing the number of local updates does not degrade the final accuracy, in contrast to training from random model
FEDADAM GD is as effective as FEDADAM SGD with pre-training. The seminal work of McMahan et al. (2016) shows that taking local SGD steps before server averaging reduces communication by 10-100 (0-100x compared to taking a full batch gradient step. To understand how pre-training impacts this comparison, we compare FEDADAM with SGD and FEDADAM with GD. While local SGD can reduce communication, the saving much less when the models are initialized with pre-trained weights compared to random weights. Figure the Appendix shows that with pre-trained initialization, using GD at the client can yield almost the same result as taking local SGD steps.
Pre-trained Random
Accurary
70
Pre-training reduces the impact of system heterogene- ity. To study the impact of pre-training on system hetero- geneity, we follow the setup described in (Li et al., 2018; 18; Wang et al., 2020). We sample 30% of clients uniformly random, and client i performs E, local epochs where E, ~ ((1,5) 5) while the remaining 70% of the clients perform E, = 1 epochs; this models the setting where clients have different processing capabilities and they perform as much work as they can within a given time limit. Figure 5 shows that FEDADAM-SGD consistently outperforms other methods specifically designed for system heterogeneity (NoVA, PROXIMAL, MIMELITE) when starting from a pre-trained model. Apparently using an adaptive optimizer at the server is sufficient to correct for the negative effects systems heterogeneity when starting from a pre-trained model. On the other hand, when starting from a random initialization, optimizers specifically designed for system heterogeneity (i.e., FEDNOVA) outperform SGD (Figure 5 right). Moreover, the accuracy gap between algorithms is more pronounced in the random initialization setting, whereas in the pre-trate pre-trained setting, all algorithms converge to more similar accuracies. Our results suggest that pre- training may reduce the need for algorithms that try correct system heterogeneity.
Figure 3: The accuracy for CIFAR-10 using ResNet-18 18 with increasing number of local epochs.

---

Published as a conference paper at ICLR 2023
Random Pretrained
70
2.5
Initial
Evel Accurary
60
2.0
6000
1.5
1.0
ALP
Random Protrained
Random Fretrained
0.5
Vigll
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Round
Round
Round
L2 Distance from Finall Model
0.6
Random
6000
Random
Gradient Diversity
Pretrained
Pretra ned
=000
Casinets
0.2
2000
Random Protrained
200
205
300
4000
500
100
200
300
400
500
200
200
3:00
400
500
Round
Round
Round
Figure 6: Training and gradient statistics of a Resnet18 on CIFAR- with Dirichlet distribution with parameter 0.1. Top row: Train loss of global model; train accuracy of global model; evaluation accuracy of global model; evaluation loss of global model. Bottom row: Gradient diversity of client updates; cosine similarity between client updates; 1.2 distance of server weights from their final values at the end of training.
UNDERSTANDING WHY PRE-TRAINING HELPS FEDERATED OPTIMIZATION
While pre-training unsurprisingly speeds up convergence, the reason for the speedup is less apparent. In this section, we examine why pre-training beneficial to federated learning-
Pre-training helps align client updates. To better understand why pre-training alleviates the heterogeneity challenge, we first investigate the gradient diversity of the updates received from different clients. We adopt the notion introduced in Yin al. (2017), adapted here to apply to client updates A; (whereas Yin et al. (2017) focus specifically on gradients g:):
= EiEst || E A?ll
CradiennDiivrsityy((:::
In Figure 6, we plot the gradient diversity client updates A; at each round for FEDADAM. In the pre-trained setting, client updates have significantly lower gradient diversity (see the bottom left plot in Figure 6). This suggests that when starting from a pre-trained model, the client local model changes randomly initialized weights are almost orthogonal, suffering more from the client drift problem. In addition, when looking at the cosine similarity of consecutive aggregated update vectors time (bottom middle), we see that consecutive updates point more consistently in a similar direction at the beginning training when starting from pre-trained model.
similar each other. On the other hand, clients local model changes from
From the top middle and bottom right plots in Figure 6, we see that the pre-trained model starts closer the final result. We also examine the largest eigenvalue the Hessian matrix (i.e., local Lipshitz constant smoothness) at the beginning of training, a larger value of which suggests harder-to-optimizer loss surface. compute the Hessian matrix, We examine the largest eigenvalue of the Hessian matrix at initialization, round 0. using Power Iteration from PyHessian by Yao et al. (2020). In Table 1, one can observe that pre-trained models always lead to smaller eigenvalues on different datasets.
Initial loss for pre-trained versus random models. Table 2 shows that pre-traine pre-training does not always lead to lower initial loss. For Squeezenet 1.0 CIFAR- 10 and ResNet- FEMNIST, the initial loss of the randomly initialized models are lower pre-trained models. However, pre-trained model still converges faster as illustrated in Figure 2.
7

---

Published as a conference paper at ICLR 2023
CIFAR-10
FEMNIST
Stack Overflow
Reddit
Pre-trained
661.99
26.29
151.05 185.02
647.19 1309.68
Random
4843.13
355.51
Table 1: The top eigenvalue of the Hessian matrix for each dataset between the pre-trained and random initialized models.
CIFAR-10
FEMNIST
Squeezenet 1.0
ResNet-18 1.07
Squeeze 1.0
ResNet-18
Pre-trained 2.71
3.99
6.58
Random
1.90
1.11
4.31
4.17
Stack Overflow
Reddit
CharLM
DistilGPT2
CharLM
DistilGPT2
Pre-trained 6.78
4.93
5.15
6.34
Random 7.71
9.82
8.61
9.99
Table 2: Loss beginning of training for various model architectures and datasets. The initial loss of the pretrained model is not always lower than that of a random initialization.
Connection to theory. Here, present the existing optimization theory for FEDAVG and discuss how pre-training helps to improve the model convergence. Following the formulation in Section 3, suppose that there total m clients, jointly optimizing global objective function ʃ(w) Σ -==1 p; (:(), and that each client's local loss function F((w) is T--Lipschitz smooth. For ease of presentation, we assume that all clients participate in training and perform K local SGD updates each round. Then, under standard assumptions, one can show that after R communication rounds, the expected gradient norm satisfies (see Theorem in Karimireddy al. (2020)):
E < so
VF VRKm
F2/3(2/3
R2/3
where ã represents weighted sampled model from all previous rounds, Ç is measure of data heterogeneity, and = f(xx) f denotes the gap between the initial loss value and the optimal loss value.
In addition, in order for FEDAVG to achieve the O(1/vRKm asymptotic convergence rate, previous works (Wang al., 2021; Woodworth et al., 2020; Karimireddy al., 2020; Wang and Joshi, 2021) showed that the number of local updates K should be upper bounded as follows:
R//3
K < O
F1/354/3m
Clearly, if F becomes smaller starting from a pre- model, one can use larger number of local updates. This corroborates our empirical observations in Figure 3.
sometimes reduced, as observed in Table 2. As result, the optimization error upper bound (2) will be smaller, i.e., we get better worst-case performance. However a lower initial loss is not always observed in our experiments, so this does not fully explain our observations, suggesting that we may need to re-think the convergence theory of local update methods.
When starting from a pre-trained model, the initial gap
RECOMMENDATIONS
In this work, study the effects of pre-training federated optimization methods. Our results inform the following recommendations:
1. When evaluating FL algorithms, researchers should experiment with both pre-trained (if available) and random weights, as the initialization can clearly impact the relative performance of different methods, and both initialization schemes may be relevant to practice.
8

---

Published as a conference paper at ICLR 2023
2. When deploying FL to a production environment, using adaptive server optimizers such as FEDADAM together with SGD at the client is a simple and competitive approach when it is possible to start from pre-trained model.
When there is public data to pre-train train a model, the impact of heterogeneity can be reduced. Thus, when focusing heterogeneity, it may be worth considering whether not proxy data is available for pre-training to motivate the application considered.
7
RELATED WORK
Transfer learning. Model initialization can significantly impact training and final performance. Previous work studying the loss landscape of deep networks observed significant differences between the landscape around a random initialization and the landscape later in training. In particular, later in training, the loss be much more "well-behaved" (Li et al., 2017; Frankle al., 2020; 2019). Fine-tuning from pre-trained models is common practice in natural language processing and computer vision, yielding strong performance many tasks (Radford al., 2019; Dosovitskiy al., 2020; Devlin et al., 2018; He et al., 2019).
Federated Optimization. While significant amount of research focused on various aspects of FL, including communication-effiieee
McMahan et al. (2016), data and systems heterogencity Li al. (2018); Wang et al. (2020), and faster convergence rate Nearly all previous work in this field neglect the importance of initialization. In our work, we study the impact of initialization on federated optimization in the cross-device setting. We defer the interested reader to surveys of Kairouz et al. (2019) and Wang et al. (2021) for additional background.
Pre-training in Federated Learning. Few studies have investigated pre-trained models federated learning, including Pillutla et al. (2022); Hsu et al. (2020); Zhao et al. (2018); Lin et al. (2021); Stremmel and Singh (2021); Tan et al. (2022). While Zhao et al. (2018) found pre-training does not alleviate the effect of heterogeneity. Tan et al. (2022) focuses on effectively learning from pre-trained models, and Chen et al. (2022); Weller et al. (2022); Chen et al. (2022) focused on synthetically partitioned data and proposed pre-training methods with synthetic data. Our work fills gap by comparing random initialization and pre-training, which other studies Pillutla et al. (2022); Hsu et al. (2020); Lin et al. (2021); Stremmel and Singh (2021) did not address. We systematically investigate both forms of heterogeneity and evaluate 15 SOTA federated optimization algorithms across visual and language tasks. Our study offers theoretical and empirical evidence supporting the benefits of pre-training in FL
8 CONCLUSION AND LIMITATIONS
Limitations. Depending on the application, may not be possible to get public data, in which case random initialization may be the only option. Nevertheless, we believe there is sufficient prevalence and importance of applications where public data is available for this study to be of broad interest. When public data available, it may not necessarily reflect the distribution of all users the population. Consequently, pre-training using public data may introduce bias, which warrants further study, including methods to detect and mitigate such bias. Moreover, we only consider one warm-star initialization strategy: supervised pre-training. Several other possibilities are worth investigating, including meta-learning the warm-start initialization and self-supprvised pre-training.
Conclusion. In this paper, we present a thorough empirical analysis of initialization on federated learning by evaluating on twelve federated learning algorithms across four vision and text tasks. We find that pre-training on public data can recover most of the accuracy drop from heterogeneity. We show that client updates starting from pre-trained weights have higher cosine similarity, which explains why initializing with pre-trained weights can speed up convergence and achieve high accuracy even in heterogeneous settings. We further show that using simple SGD locally can be as good as other local optimizers.

---

Published as a conference paper at ICLR 2023
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regret. International Conference on Learning Representations, 2021.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: benchmark for federated settings. arXiv preprint arXiv:
2018.
Hong-You You Chen, Cheng-Hao Tu, Ziwei Li, Han-Wei Shen, and Wei-Lun Chao. On pre-training for federated learning. arXiv preprint arXiv:2206. /488, 2022.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly. et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 19929, 2020.
FI.Sim Authors. Federated learning simulator (flsim). 2022.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. CoRR, abs/1912.056711, 2019. URL http://
wrxiv.org// abs / 1912... 05671.
Jonathan Frankle, David J. Schwab, and Ari S. Morcos. The early phase of neural network training. CoRR, abs/2002. 2020. URL https: /trr//ixivvv.///////ee////eeeeva org. /2002 10365.
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 Dollár. Rethinking ImageNet pre-training. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non- con-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06:0 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Federated visual classification with real- world data distribution. In European Conference Computer Vision, pages 76 92. Springer, 2020.
HuggingFace. Distilgpt2. 2019. URL hatlhtttt//hhggigggffeeeee co wotteetttttts.s.........
Forrest N landola, Song Han, Matthew W Moskewiez, Khalid Ashraf, William Dally, and Kurt Keutzer. Squeezenet: Alexnet-level accuracy with 50x fewer parameters and< 0.5 mb model size. arXiv preprint arXiv:1602.07360 2016.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv: (ete::9/220099 2019.
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.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Shashank Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. 2021.
Yoon Kim, Yacine Jernite, David Sontag, and Alexander M Rush. Character-aware neural language models. In Thirtieth AAAI conference artificial intelligence, 2016.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
10

---

Published as a conference paper at ICLR 2023
Hao Li, Zheng Xu, Gavin Taylor, and Tom Goldstein. Visualizing the loss landscape of neural nets. CoRR, abs/1712.09913,, 2017. URL httt::tpp/::/ia http://aaxxxxv.x.xaaa/aa wwtg/bb.............../ httpg//bb/1112 1712. 09913.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. 2018.
Federated optimization in heterogeneous networks, arXiv preprini
Bill Yuchen Lin, Chaoyang He, Zihang Zeng, Hulin Wang. Yufen Huang, Mahdi Soltanolkotabi, Xiang Ren, and Salman Avestimehr. Fednip: Benchmarking federated learning methods for natural language processing tasks. arXiv preprint WrXivv2204.000 08815, 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, and Blaise Agüera Arcas. Federated learning of deep networks using model averaging. arXiv preprint arXiv: 05669, 2016.
H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Communication-effiee efficient learning of deep networks from decentralized data. In International Conference on Artificial Intelligence and Statistics, 2017.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. arXiv preprint arXiv:2204.03809 2022.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00294 00295, 2020.
Jasper Snock, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. Advances in neural information processing systems, 25:2951-295 -9999, 2012.
Joel Stremmel and Arjun Singh. Pretraining federated text models for next word prediction. In Future of Information and Communication Conference, pages 477-488. Springer, 2021.
Yue Tan, Guodong Long, Jie Ma, Lu Liu, Tianyi Zhou, and Jing Jiang. Federated learning from pre-trained models: A contrastive learning approach. arXiv preprint arXiv:2209. 10083, 2022.
Jianyu Wang and Gauri Joshi. Cooperative sgd: A unified framework for the design and analysis of local-update sgd algorithms. Journal of Machine Learning Research, 22, 2021.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7622 2020.
Jianyu Wang. Zachary Charles, Zheng Xu, Gauri Joshi, Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide to federated optimization. Xiviv preprint arXiv:2 2021.
Orion Weller, Mare Marone, Vladimir Braverman, Dawn Lawrie, and Benjamin Van Durme. Pre- trained models for multilingual federated learning. arXiv preprint arXiv:2206. 02291, 2022.
Blake Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heterogeneous distributed learning. Advances in Neural Information Processing Systems, 33:6281-6299 2020.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pages 581 581-590. IEEF, 2020.
11

---

Published as a conference paper at ICLR 2023
Dong Yin, Ashwin Pananjady, Maximilian Lam, Dimitris S. Papailiopoulos, Kannan Ramchandran, and Peter L. Bartlett. Gradient diversity empowers distributed learning. CoRR, abs/1706.05699, 2017.
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. FedPD: federated learning framework with adaptivity to non-iid data. IEEE Transactions on Signal Processing, 69:6055-6070, 2021.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data, arXiv preprint arXiv:1806000588 2018.
12


Reviews:
[{"review_id": "qj4j9Eg9k1b", "reviewer": "Reviewer_aMPP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the effect of using pre-trained neural network weights on federated optimization. Several combinations of server and client optimizers are considered in combination with 2 different models for vision and language tasks. It was shown that pre-trained initialization reduces the negative effects of data and systems heterogeneity on training convergence and changes the ranking of optimization methods.", "summary_of_the_review": "Overall, the paper represents a piece of good-quality research. I think that there are enough insights for publication. There is also a potential for an increase in my score if some of the questions and concerns are addressed by the authors.", "strength_and_weaknesses": "## Strengths\n\n**S1.** Thorough and extensive experimental evaluation of different optimization aspects which may be useful to the field of federated learning.\n\n**S2.** Several drawn conclusions are novel and can be interesting for the community.\n\n## Weaknesses\n\n**W1.** The concept of using pre-trained weights (in federated deep learning) is not novel on its own. Some of the ideas were explored in previous works.\n\n**W2.** Some important experimental details are missing. Namely, it is not clear how hyper-parameter tuning was performed. What clients/data were used for tuning step sizes, etc (the ranges also seem very limited)? It would be very helpful to show how pre-trained initialization affects the performance of the methods for different combinations of parameters (e.g., like it was done with heatmaps in the Adaptive FedOpt paper).  What kind of random initialization was used in the experiments? Prior literature showed that various strategies can lead to different results in the centralized case.\n\nPlease correct me, if I missed these details.\n\n**W3.** The thesis \n> Starting from a pre-trained model signiﬁcantly reduces the difference between having non IID vs IID data at clients.\n\nis not well-supported (or inaccurately stated), in my view. It is not fully clear what is meant by IID data splits for naturally partitioned federated datasets in the paper. Was the data somehow reallocated between the same amount of clients or simply centralized training was performed? The latter case would require different wording in my view as it studies not the difference between non-IID and IID but rather federated and centralized cases. To properly show the effect of using the pre-trained model for a level of non IIDness one ideally needs to vary the level of heterogeneity for the same problem.\n\n**W4.** Some of the Section 5 conclusions are questionable and doubtful. For example, the Gradient Diversity (GD) of client updates is lower for the pre-trained model, while the original paper by Yin et al., 2017 argues that better convergence is achieved for higher GD.  As far as I understand the plots show results of single runs which is not very reliable. Cosine similarity behavior looks quite random and similar for different initialization, apart from maybe the start. The connection to theory is not really convincing.\n\nI would also like to ask how exactly the Hessian matrix was computed.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ceM3kYmN7H", "reviewer": "Reviewer_Kusi", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors study federated learning starting from pre-trained models in both the image and text setting. They study a wide range of FL algorithms. Some key observations are that (a) the gap of iid to non-iid closes (b) ranking of FL algorithms is different under pre-training (c) effect of heterogneity are not severed. An investigation of how pre-training helps under heterogeneity is performed. \n", "summary_of_the_review": "The paper has a number of important findings for guiding FL research and practice. The evaluations are thoroughly done across a wide range of datasets and tasks. Insights into the observed effects are also provided. \n\nPost-Rebuttal: I have read the rebuttal and response to other authors. My review and score stand.", "strength_and_weaknesses": "Strengths:\n-The evaluations are extensive covering a wide array of tasks, datasets, and federated optimization settings\n-The observations are potentially important: methods addressing heterogeneity may not be relevant in cases where pre-trained models can be used\n\nNo substantial weakness are noted \n\nSome Comments: \n- The effect of the pre-trained model is not heavily studied. It would be interesting to know if better pre-trained models lead to better results. \n- I would like to see some discussion on how the learning rates selected are different (or same) in the pre-training/random setting\n- It would be interesting to see the effect on more distant combinations of pre-training and target data. For example imagenet->medical/satellite image data\n- In Chen et al it was shown that FedAVG can benefit even from pre-training on synthetic data. It would be interesting to see if such approaches (and the change in ranking of algos) hold for the algorithms evaluated here.\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9w5_FDf2AfU", "reviewer": "Reviewer_EYvh", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the impact of model initialization(random or pre-trained) on federated optimization methods. The experiments show that starting federated learning from a pre-trained initialization reduces the effect of both data and system heterogeneity. The experiments also show that the pre-trained initialization reduces the training time.", "summary_of_the_review": "The paper is significant to model initialization(random or pre-trained) on federated optimization methods. The contribution is clearly presented and the experimental results are thoroughly executed. I would tend to accept this paper.", "strength_and_weaknesses": "Strength:\n\n1)The paper is well written, and the experiments are sufficient.\n\n2)The paper finds that starting federated learning from a pre-trained initialization reduces the effect of both data and system heterogeneity.\n\n3)The experiments show that the pre-trained initialization leads to communication savings and reduces the overall training time in federated learning.\n\n4)The findings in this paper suggest directions for future related work.\n\nWeaknesses:\n\n1)In Figure 1 and Figure 5, what is the representation of the curve radian?\n\n2)The α in Figure 2 is not explained in the main paper.\n\n3)Does the model pre-trained on public data introduce security issues to users? How do we make trade-offs between security and performance?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%