Adaptive Federated Optimization

1 INTRODUCTION . Federated learning ( FL ) is a machine learning paradigm in which multiple clients cooperate to learn a model under the orchestration of a central server ( McMahan et al. , 2017 ) . In FL , raw client data is never shared with the server or other clients . This distinguishes FL from traditional distributed optimization , and requires contending with heterogeneous data . FL has two primary settings , crosssilo ( eg . FL between large institutions ) and cross-device ( eg . FL across edge devices ) ( Kairouz et al. , 2019 , Table 1 ) . In cross-silo FL , most clients participate in every round and can maintain state between rounds . In the more challenging cross-device FL , our primary focus , only a small fraction of clients participate in each round , and clients can not maintain state across rounds . For a more in-depth discussion of FL and the challenges involved , we defer to Kairouz et al . ( 2019 ) and Li et al . ( 2019a ) . Standard optimization methods , such as distributed SGD , are often unsuitable in FL and can incur high communication costs . To remedy this , many federated optimization methods use local client updates , in which clients update their models multiple times before communicating with the server . This can greatly reduce the amount of communication required to train a model . One such method is FEDAVG ( McMahan et al. , 2017 ) , in which clients perform multiple epochs of SGD on their local datasets . The clients communicate their models to the server , which averages them to form a new global model . While FEDAVG has seen great success , recent works have highlighted its convergence issues in some settings ( Karimireddy et al. , 2019 ; Hsu et al. , 2019 ) . This is due to a variety of factors including ( 1 ) client drift ( Karimireddy et al. , 2019 ) , where local client models move away from globally optimal models , and ( 2 ) a lack of adaptivity . FEDAVG is similar in spirit to SGD , and may be unsuitable for settings with heavy-tail stochastic gradient noise distributions , which often arise when training language models ( Zhang et al. , 2019a ) . Such settings benefit from adaptive learning rates , which incorporate knowledge of past iterations to perform more informed optimization . In this paper , we focus on the second issue and present a simple framework for incorporating adaptivity in FL . In particular , we propose a general optimization framework in which ( 1 ) clients perform multiple epochs of training using a client optimizer to minimize loss on their local data and ( 2 ) server updates its global model by applying a gradient-based server optimizer to the average of the clients ’ model updates . We show that FEDAVG is the special case where SGD is used as both client and server optimizer and server learning rate is 1 . This framework can also seamlessly incorporate ∗Authors contributed equally to this work adaptivity by using adaptive optimizers as client or server optimizers . Building upon this , we develop novel adaptive optimization techniques for FL by using per-coordinate methods as server optimizers . By focusing on adaptive server optimization , we enable use of adaptive learning rates without increase in client storage or communication costs , and ensure compatibility with cross-device FL . Main contributions In light of the above , we highlight the main contributions of the paper . • We study a general framework for federated optimization using server and client optimizers . This framework generalizes many existing federated optimization methods , including FEDAVG . • We use this framework to design novel , cross-device compatible , adaptive federated optimization methods , and provide convergence analysis in general nonconvex settings . To the best of our knowledge , these are the first methods for FL using adaptive server optimization . We show an important interplay between the number of local steps and the heterogeneity among clients . • We introduce comprehensive and reproducible empirical benchmarks for comparing federated optimization methods . These benchmarks consist of seven diverse and representative FL tasks involving both image and text data , with varying amounts of heterogeneity and numbers of clients . • We demonstrate strong empirical performance of our adaptive optimizers throughout , improving upon commonly used baselines . Our results show that our methods can be easier to tune , and highlight their utility in cross-device settings . Related work FEDAVG was first introduced by McMahan et al . ( 2017 ) , who showed it can dramatically reduce communication costs . Many variants have since been proposed to tackle issues such as convergence and client drift . Examples include adding a regularization term in the client objectives towards the broadcast model ( Li et al. , 2018 ) , and server momentum ( Hsu et al. , 2019 ) . When clients are homogeneous , FEDAVG reduces to local SGD ( Zinkevich et al. , 2010 ) , which has been analyzed by many works ( Stich , 2019 ; Yu et al. , 2019 ; Wang & Joshi , 2018 ; Stich & Karimireddy , 2019 ; Basu et al. , 2019 ) . In order to analyze FEDAVG in heterogeneous settings , many works derive convergence rates depending on the amount of heterogeneity ( Li et al. , 2018 ; Wang et al. , 2019 ; Khaled et al. , 2019 ; Li et al. , 2019b ) . Typically , the convergence rate of FEDAVG gets worse with client heterogeneity . By using control variates to reduce client drift , the SCAFFOLD method ( Karimireddy et al. , 2019 ) achieves convergence rates that are independent of the amount of heterogeneity . While effective in cross-silo FL , the method is incompatible with cross-device FL as it requires clients to maintain state across rounds . For more detailed comparisons , we defer to Kairouz et al . ( 2019 ) . Adaptive methods have been the subject of significant theoretical and empirical study , in both convex ( McMahan & Streeter , 2010b ; Duchi et al. , 2011 ; Kingma & Ba , 2015 ) and non-convex settings ( Li & Orabona , 2018 ; Ward et al. , 2018 ; Wu et al. , 2019 ) . Reddi et al . ( 2019 ) ; Zaheer et al . ( 2018 ) study convergence failures of ADAM in certain non-convex settings , and develop an adaptive optimizer , YOGI , designed to improve convergence . While most work on adaptive methods focuses on non-FL settings , Xie et al . ( 2019 ) propose ADAALTER , a method for FL using adaptive client optimization . Conceptually , our approach is also related to the LOOKAHEAD optimizer ( Zhang et al. , 2019b ) , which was designed for non-FL settings . Similar to ADAALTER , an adaptive FL variant of LOOKAHEAD entails adaptive client optimization ( see Appendix B.3 for more details ) . We note that both ADAALTER and LOOKAHEAD are , in fact , special cases of our framework ( see Algorithm 1 ) and the primary novelty of our work comes in focusing on adaptive server optimization . This allows us to avoid aggregating optimizer states across clients , making our methods require at most half as much communication and client memory usage per round ( see Appendix B.3 for details ) . Notation For a , b ∈ Rd , we let √ a , a2 and a/b denote the element-wise square root , square , and division of the vectors . For θi ∈ Rd , we use both θi , j and [ θi ] j to denote its jth coordinate . 2 FEDERATED LEARNING AND FEDAVG . In federated learning , we solve an optimization problem of the form : min x∈Rd f ( x ) = 1 m m∑ i=1 Fi ( x ) , ( 1 ) where Fi ( x ) = Ez∼Di [ fi ( x , z ) ] , is the loss function of the ith client , z ∈ Z , and Di is the data distribution for the ith client . For i 6= j , Di and Dj may be very different . The functions Fi ( and therefore f ) may be nonconvex . For each i and x , we assume access to an unbiased stochastic gradient gi ( x ) of the client ’ s true gradient∇Fi ( x ) . In addition , we make the following assumptions . Assumption 1 ( Lipschitz Gradient ) . The function Fi is L-smooth for all i ∈ [ m ] i.e. , ‖∇Fi ( x ) − ∇Fi ( y ) ‖ ≤ L‖x− y‖ , for all x , y ∈ Rd . Assumption 2 ( Bounded Variance ) . The function Fi have σl-bounded ( local ) variance i.e. , E [ ‖∇ [ fi ( x , z ) ] j − [ ∇Fi ( x ) ] j‖2 ] = σ2l , j for all x ∈ Rd , j ∈ [ d ] and i ∈ [ m ] . Furthermore , we assume the ( global ) variance is bounded , ( 1/m ) ∑m i=1 ‖∇ [ Fi ( x ) ] j − [ ∇f ( x ) ] j‖2 ≤ σ2g , j for all x ∈ Rd and j ∈ [ d ] . Assumption 3 ( Bounded Gradients ) . The function fi ( x , z ) have G-bounded gradients i.e. , for any i ∈ [ m ] , x ∈ Rd and z ∈ Z we have | [ ∇fi ( x , z ) ] j | ≤ G for all j ∈ [ d ] . With a slight abuse of notation , we use σ2l and σ 2 g to denote ∑d j=1 σ 2 l , j and ∑d j=1 σ 2 g , j . Assumptions 1 and 3 are fairly standard in nonconvex optimization literature ( Reddi et al. , 2016 ; Ward et al. , 2018 ; Zaheer et al. , 2018 ) . We make no further assumptions regarding the similarity of clients datasets . Assumption 2 is a form of bounded variance , but between the client objective functions and the overall objective function . This assumption has been used throughout various works on federated optimization ( Li et al. , 2018 ; Wang et al. , 2019 ) . Intuitively , the parameter σg quantifies similarity of client objective functions . Note σg = 0 corresponds to the i.i.d . setting . A common approach to solving ( 1 ) in federated settings is FEDAVG ( McMahan et al. , 2017 ) . At each round of FEDAVG , a subset of clients are selected ( typically randomly ) and the server broadcasts its global model to each client . In parallel , the clients run SGD on their own loss function , and send the resulting model to the server . The server then updates its global model as the average of these local models . See Algorithm 3 in the appendix for more details . Suppose that at round t , the server has model xt and samples a set S of clients . Let xti denote the model of each client i ∈ S after local training . We rewrite FEDAVG ’ s update as xt+1 = 1 |S| ∑ i∈S xti = xt − 1 |S| ∑ i∈S ( xt − xti ) . Let ∆ti : = x t i − xt and ∆t : = ( 1/|S| ) ∑ i∈S ∆ t i . Then the server update in FEDAVG is equivalent to applying SGD to the “ pseudo-gradient ” −∆t with learning rate η = 1 . This formulation makes it clear that other choices of η are possible . One could also utilize optimizers other than SGD on the clients , or use an alternative update rule on the server . This family of algorithms , which we refer to collectively as FEDOPT , is formalized in Algorithm 1 . Algorithm 1 FEDOPT 1 : Input : x0 , CLIENTOPT , SERVEROPT 2 : for t = 0 , · · · , T − 1 do 3 : Sample a subset S of clients 4 : xti,0 = xt 5 : for each client i ∈ S in parallel do 6 : for k = 0 , · · · , K − 1 do 7 : Compute an unbiased estimate gti , k of∇Fi ( xti , k ) 8 : xti , k+1 = CLIENTOPT ( x t i , k , g t i , k , ηl , t ) 9 : ∆ti = x t i , K − xt 10 : ∆t = 1 |S| ∑ i∈S ∆ t i 11 : xt+1 = SERVEROPT ( xt , −∆t , η , t ) In Algorithm 1 , CLIENTOPT and SERVEROPT are gradient-based optimizers with learning rates ηl and η respectively . Intuitively , CLIENTOPT aims to minimize ( 1 ) based on each client ’ s local data while SERVEROPT optimizes from a global perspective . FEDOPT naturally allows the use of adaptive optimizers ( eg . ADAM , YOGI , etc . ) , as well as techniques such as server-side momentum ( leading to FEDAVGM , proposed by Hsu et al . ( 2019 ) ) . In its most general form , FEDOPT uses a CLIENTOPT whose updates can depend on globally aggregated statistics ( e.g . server updates in the previous iterations ) . We also allow η and ηl to depend on the round t in order to encompass learning rate schedules . While we focus on specific adaptive optimizers in this work , we can in principle use any adaptive optimizer ( e.g . AMSGRAD ( Reddi et al. , 2019 ) , ADABOUND ( Luo et al. , 2019 ) ) . While FEDOPT has intuitive benefits over FEDAVG , it also raises a fundamental question : Can the negative of the average model difference ∆t be used as a pseudo-gradient in general server optimizer updates ? In this paper , we provide an affirmative answer to this question by establishing a theoretical basis for FEDOPT . We will show that the use of the term SERVEROPT is justified , as we can guarantee convergence across a wide variety of server optimizers , including ADAGRAD , ADAM , and YOGI , thus developing principled adaptive optimizers for FL based on our framework .

Summary: This paper extends the server model averaging step in FedAvg to a more general adaptive optimization step on the global model, specifically, by writing the model averaging as a gradient descent step using a pseudo gradient. Three variants of this scheme (FedOpt) are presented, based on three adaptive optimizers, including AdaGrad, ADAM, and YOGI. While there exist works applying the server-side momentum method, the paper argues that the proposed framework is more general since any adaptive optimizer can be applied to extend FedAvg. Under three assumptions (Lipschitz gradient, bounded gradients, bounded variances), the paper provides a local convergence analysis with nonconvex objectives. The achieved bounds match the best-known convergence rate for FL under reasonable conditions (i.e., T is sufficiently large compared to K). They also provide guidance for how to decay the local learning rate to avoid client drift, and show that increasing local steps can help to reduce the communication rounds (under certain conditions). They also reflect the dependency of these bounds on client heterogeneity measured by global variance and indicates how to mitigate the problem in FedOpt. Experimental comparison to FedAvg, FedAvgM, and SCAFFOLD on seven benchmark FL tasks show that the proposed three adaptive FL methods are better or comparable on early-stage convergence and final validation-set performance. A study about tuning the local and global learning rate is also presented.