The role of local steps in local SGD

AbstractWe consider the distributed stochastic optimization problem where n agents want to minimize a global function given by the sum of agents' local functions and focus on the heterogeneous setting when agents' local functions are defined over non-i.i.d. datasets. We study the Local SGD method, where agents perform a number of local stochastic gradient steps and occasionally communicate with a central node to improve their local optimization tasks. We analyze the effect of local steps on the convergence rate and the communication complexity of Local SGD. In particular, instead of assuming a fixed number of local steps across all communication rounds, we allow the number of local steps during the jth communication round, Hj, to be different and arbitrary numbers. Our main contribution is to characterize the convergence rate of Local SGD as a function of {Hj}j=1R under various settings of strongly convex, convex, and nonconvex local functions, where R is the total number of communication rounds. Based on this characterization, we provide sufficient conditions on the sequence {Hj}j=1R such that Local SGD can achieve linear speedup with respect to the number of workers. Furthermore, we propose a new communication strategy with increasing local steps that is superior to constant local steps for strongly convex local functions. On the other hand, for convex and nonconvex local functions, we argue that fixed local steps are the best communication strategy for Local SGD and recover state-of-the-art convergence rate results. Finally, we justify our theoretical results through extensive numerical experiments.Keywords: Federated learninglocal SGDdistributed optimization Disclosure statementNo potential conflict of interest was reported by the author(s).Notes1 Our work was done independently of [Citation26], which was brought to our attention during the review process.2 For constant number of local steps, the results in [Citation11] implies that Local SGD can achieve O(log⁡(T)nT) convergence rate with O(nT) communication rounds without the bounded gradient assumption; [Citation24] showed that linear speedup can be achieved with O(nT) communication rounds, however, their analysis requires the bounded gradient assumption. Recently, [Citation26] proposed a communication strategy with exponentially increasing number of local steps (different from ours) and also showed O(1nT) convergence rate with O(nT) communication rounds.3 It is easy to see that for such choice of parameters conditions (2) and (3) in Corollary 3.1 are satisfied. For the proof of condition (1), we refer to Appendix A.5.4 In fact, Theorem 2 in [Citation35] stated a lower bound of O(1nT)+O(Tn2R3)+min{O(exp⁡(−R2T)),O(TR3)}. The higher order terms also suggest that in order to achieve O(1nT) convergence rate, R needs to be at least max(O(n−1/3T2/3),O((Tlog⁡(nT))1/2)).5 This matches the best-known results in the setting [Citation9].6 This matches the best-known results in the setting [Citation8].7 We have deliberately chosen to train a small neural network to avoid getting an overparameterized model, in which case the convergence rate of Local SGD would be different [Citation23].Additional informationFundingThis material is based upon work supported by the Air Force Office of Scientific Research under award number FA9550-23-1-0107 and the NSF CAREER Award under Grant No. EPCN-1944403. The work of C.A.Uribe was partially supported by the National Science Foundation under Grants No. 2211815 and No. 2213568.Notes on contributorsTiancheng QinTiancheng Qin is a Ph.D student in the department of industrial and systems engineering at the University of Illinois Urbana-Champaign (UIUC) under the supervision of Dr. S. Rasoul Etesami, where his research interests include distributed optimization, federated learning and reinforcement learning. He received a bachelor’s degree from the School of Mathematical Sciences, Peking University, China in 2018.S. Rasoul EtesamiS. Rasoul Etesami is an sssistant professor in the department of industrial and systems engineering at the University of Illinois Urbana-Champaign (UIUC), where he is also affiliated with the department of electrical and computer engineering and the Coordinated Science Laboratory. From 2016 to 2017, he was a postdoctoral research fellow in the department of electrical engineering at Princeton University and WINLAB. He received his Ph.D. in electrical and eomputer Engineering from the University of Illinois Urbana-Champaign in 2015. His research interests include the analysis of complex social, communication and decision-making systems using tools from control and game theory, optimization and learning theory. He was the recipient of the NSF CAREER Award in 2020 and the Air Force Young Investigator Award in 2023.César A. UribeCésar A. Uribe is the Louis Owen Jr. Assistant Professor at the department of electrical and computer engineering at Rice University. He received his M.Sc. degrees in systems and control from Delft University of Technology, in The Netherlands, and applied mathematics from the University of Illinois at Urbana-Champaign, in 2013 and 2016, respectively. He also received a PhD degree in electrical and computer engineering at the University of Illinois at Urbana-Champaign in 2018. He was a postdoctoral associate in the Laboratory for Information and Decision Systems-LIDS at the Massachusetts Institute of Technology-MIT until 2020 and is a visiting professor at the Moscow Institute of Physics and Technology. His research interests include distributed learning and optimization, decentralized control, algorithm analysis and computational optimal transport.