Federated Learning of a Mixture of Global and Local Models

1 INTRODUCTION . With the proliferation of mobile phones , wearable devices , tablets , and smart home devices comes an increase in the volume of data captured and stored on them . This data contains a wealth of potentially useful information to the owners of these devices , and more so if appropriate machine learning models could be trained on the heterogeneous data stored across the network of such devices . The traditional approach involves moving the relevant data to a data center where centralized machine learning techniques can be efficiently applied ( Dean et al. , 2012 ; Reddi et al. , 2016 ) . However , this approach is not without issues . First , many device users are increasingly sensitive to privacy concerns and prefer their data to never leave their devices . Second , moving data from their place of origin to a centralized location is very inefficient in terms of energy and time . 1.1 FEDERATED LEARNING . Federated learning ( FL ) ( McMahan et al. , 2016 ; Konečný et al. , 2016b ; a ; McMahan et al. , 2017 ) has emerged as an interdisciplinary field focused on addressing these issues by training machine learning models directly on edge devices . The currently prevalent paradigm ( Li et al. , 2019 ; Kairouz et al. , 2019 ) casts supervised FL as an empirical risk minimization problem of the form min x∈Rd 1 n n∑ i=1 fi ( x ) , ( 1 ) where n is the number of devices participating in training , x ∈ Rd encodes the d parameters of a global model ( e.g. , weights of a neural network ) and fi ( x ) : = Eξ∼Di [ f ( x , ξ ) ] represents the aggregate loss of model x on the local data represented by distribution Di stored on device i . One of the defining characteristics of FL is that the data distributions Di may possess very different properties across the devices . Hence , any potential FL method is explicitly required to be able to work under the heterogeneous data setting . The most popular method for solving ( 1 ) in the context of FL is the FedAvg algorithm ( McMahan et al. , 2016 ) . In its most simple form , when one does not employ partial participation , model compression , or stochastic approximation , FedAvg reduces to Local Gradient Descent ( LGD ) ( Khaled et al. , 2019 ; 2020 ) , which is an extension of GD performing more than a single gradient step on each device before aggregation . FedAvg has been shown to work well empirically , particularly for non-convex problems , but comes with poor convergence guarantees compared to the non-local counterparts when data are heterogeneous . Some issues with current approaches to FL The first motivation for our research comes from the appreciation that data heterogeneity does not merely present challenges to the design of new provably efficient training methods for solving ( 1 ) , but also inevitably raises questions about the utility of such a global solution to individual users . Indeed , a global model trained across all the data from all devices might be so removed from the typical data and usage patterns experienced by an individual user as to render it virtually useless . This issue has been observed before , and various approaches have been proposed to address it . For instance , the MOCHA ( Smith et al. , 2017 ) framework uses a multi-task learning approach to allow for personalization . Next , ( Khodak et al. , 2019 ) propose a generic online algorithm for gradientbased parameter-transfer meta-learning and demonstrate improved practical performance over FedAvg ( McMahan et al. , 2017 ) . Approaches based on variational inference ( Corinzia & Buhmann , 2019 ) , cyclic patterns in practical FL data sampling ( Eichner et al. , 2019 ) transfer learning ( Zhao et al. , 2018 ) and explicit model mixing ( Peterson et al. , 2019 ) have been proposed . The second motivation for our work is the realization that even very simple variants of FedAvg , such as LGD , which should be easier to analyze , fail to provide theoretical improvements in communication complexity over their non-local cousins , in this case , GD ( Khaled et al. , 2019 ; 2020 ) .1 This observation is at odds with the practical success of local methods in FL . This leads us to ask the question : if LGD does not theoretically improve upon GD as a solver for the traditional global problem ( 1 ) , perhaps LGD should not be seen as a method for solving ( 1 ) at all . In such a case , what problem does LGD solve ? A good answer to this question would shed light on the workings of LGD , and by analogy , on the role local steps play in more elaborate FL methods such as local SGD ( Stich , 2020 ; Khaled et al. , 2020 ) and FedAvg . 2 CONTRIBUTIONS . In our work we argue that the two motivations mentioned in the introduction point in the same direction , i.e. , we show that a single solution can be devised addressing both problems at the same time . Our main contributions are : New formulation of FL which seeks an implicit mixture of global and local models . We propose a new optimization formulation of FL . Instead of learning a single global model by solving ( 1 ) , we propose to learn a mixture of the global model and the purely local models which can be trained by each device i using its data Di only . Our formulation ( see Sec . 3 ) lifts the problem from Rd to Rnd , allowing each device i to learn a personalized model xi ∈ Rd . These personalized models are encouraged to not depart too much from their mean by the inclusion of a quadratic penalty ψ multiplied by a parameter λ ≥ 0 . Admittedly , the idea of softly-enforced similarity of the local models was already introduced in the domain of the multi-task relationship learning ( Zhang & Yeung , 2010 ; Liu et al. , 2017 ; Wang et al. , 2018 ) and distributed optimization ( Lan et al. , 2018 ; Gorbunov et al. , 2019 ; Zhang et al. , 2015 ) . The mixture objective we propose ( see ( 2 ) ) is a special case of their setup , which justifies our approach from the modeling perspective . Note that Zhang et al . ( 2015 ) ; Liu et al . ( 2017 ) ; Wang et al . ( 2018 ) provide efficient algorithms to solve the mixture objective already . However , none of the mentioned papers consider the FL application , nor they shed a light on the communication complexity of LGD algorithms , which we do in our work . Theoretical properties of the new formulation . We study the properties of the optimal solution of our formulation , thus developing an algorithmic-free theory . When the penalty parameter is set to zero , then obviously , each device is allowed to train their own model without any dependence on the data stored on other devices . Such purely local models are rarely useful . We prove that the optimal 1After our paper was completed , a lower bound on the performance of local SGD was presented that is worse than the known minibatch SGD guarantee ( Woodworth et al. , 2020a ) , confirming that the local methods do not outperform their non-local counterparts in the heterogeneous setup . Similarly , the benefit of local methods in the non-heterogeneous scenario was questioned in ( Woodworth et al. , 2020b ) . local models converge to the traditional global model characterized by ( 1 ) at the rate O ( 1/λ ) . We also show that the total loss evaluated at the local models is never higher than the total loss evaluated at the global model ( see Thm . 3.1 ) . Moreover , we prove an insightful structural result for the optimal local models : the optimal model learned by device i arises by subtracting the gradient of the loss function stored on that device evaluated at the same point ( i.e. , a local model ) from the average of the optimal local models ( see Thm . 3.2 ) . As a byproduct , this theoretical result sheds new light on the key update step in the model agnostic meta-learning ( MAML ) method ( Finn et al. , 2017 ) , which has a similar but subtly different structure.2 The subtle difference is that the MAML update obtains the local model by subtracting the gradient evaluated at the global model . While MAML was originally proposed as a heuristic , we provide rigorous theoretical guarantees . Loopless LGD : non-uniform SGD applied to our formulation . We then propose a randomized gradient-based method—Loopless Local Gradient Descent ( L2GD ) —for solving our new formulation ( Algorithm 1 ) . This method is , in fact , a non-standard application of SGD to our problem , and can be seen as an instance of SGD with non-uniform sampling applied to the problem of minimizing the sum of two convex functions ( Zhao & Zhang , 2015 ; Gower et al. , 2019 ) : the average loss , and the penalty . When the loss function is selected by the randomness in our SGD method , the stochastic gradient step can be interpreted as the execution of a single local GD step on each device . Since we set the probability of the loss being sampled to be high , this step is typically repeated multiple times , resulting in multiple local GD steps . In contrast to standard LGD , the number of local steps is not fixed , but random , and follows a geometric distribution . This mechanism is similar in spirit to how the recently proposed loopless variants of SVRG ( Hofmann et al. , 2015 ; Kovalev et al. , 2020 ) work in comparison with the original SVRG ( Johnson & Zhang , 2013a ; Xiao & Zhang , 2014 ) . Once the penalty is sampled by our method , the resultant SGD step can be interpreted as the execution of an aggregation step . In contrast with standard aggregation , which performs full averaging of the local models , our method merely takes a step towards averaging . However , the step is relatively large . Convergence theory . By adapting the general theory from ( Gower et al. , 2019 ) to our setting , we obtain theoretical convergence guarantees assuming that each fi is L-smooth and µ-strongly convex ( see Thm . 4.2 ) . Interestingly , by optimizing the probability of sampling the penalty ( we get p ? = λλ+L ) , which is an indirect way of fixing the expected number of local steps to 1+ L λ , we prove an 2λλ+L L µ log 1 ε bound on the expected number of communication rounds ( see Cor . 4.3 ) . We believe that this is remarkable in several ways . By choosing λ small , we tilt our goal towards pure local models : the number of communication rounds is tending to 0 as λ→ 0 . If λ→∞ , the solution our formulation converges to is the optimal global model , and L2GD obtains the communication bound O ( L µ log 1 ε ) , which matches the efficiency of GD . What problem do local methods solve ? Noting that L2GD is a ( mildly nonstandard ) version of LGD,3 which is a key method most local methods for FL are based on , and noting that , as we show , L2GD solves our new formulation of FL , we offer a new and surprising interpretation of the role of local steps in FL . In particular , the role of local steps in gradient type methods , such as GD , is not to reduce communication complexity , as is generally believed . Indeed , there is no theoretical result supporting this claim in the key heterogeneous data regime . Instead , their role is to steer the method towards finding a mixture of the traditional global and the purely local models . Given that the stepsize is fixed , the more local steps are taken , the more we bias the method towards the purely local models . Our new optimization formulation of FL formalizes this as it defines the problem that local methods , in this case L2GD , solve . There is an added benefit here : the more we want our formulation to be biased towards purely local models ( i.e. , the smaller the penalty parameter λ is ) , the more local steps does L2GD take , and the better the total communication complexity of L2GD becomes . Hence , despite a lot of research on this topic , our paper provides the first proof that a local method ( e.g. , L2GD ) can be better than its non-local counterpart ( e.g. , GD ) in terms of total communication complexity in the heterogeneous data setting . We are able to do this by noting that local methods should better be seen as methods for solving the new FL formulation proposed here . Generalizations : partial participation , local SGD and variance reduction . We further generalize and improve our method by allowing for ( i ) stochastic partial participation of devices in each communication round , ( ii ) subsampling on each device which means we can perform local SGD steps 2The connection of FL and multi-task meta learning is discussed in ( Kairouz et al. , 2019 ) , for example . 3To be specific , L2GD is equivalent to Overlap LGD ( Wang et al. , 2020 ) with random local loop size . instead of local GD steps , and ( iii ) total variance reduction mechanism to tackle the variance coming from three sources : locality of the updates induced by non-uniform sampling ( already present in L2GD ) , partial participation and local subsampling . Due to its level of generality , this method , which we call L2SGD++ , is presented in the Appendix only , alongside the associated complexity results . In the main body of this paper , we instead present a simplified version thereof , which we call L2SGD+ ( Algorithm 3 ) . The convergence theory for it is presented in Thm . 5.1 and Cor . 5.2 . Heterogeneous data . All our methods and convergence results allow for fully heterogeneous data and do not depend on any assumptions on data similarity across the devices . Superior empirical performance . We show through ample numerical experiments that our theoretical predictions can be observed in practice .

Summary: of the paper: The paper proposes a new formulation for the federated learning problem, in which each agent has its local model, and a penalty term is added to the objective function to control the deviation of these local models from their average. Next, the authors develop a randomized algorithm to tackle this problem and characterize its convergence under several assumptions, such as smoothness and strong convexity. They also discuss variants of their algorithm, which uses variance reduction techniques or considers users' partial participation.