FedDR – Randomized Douglas-Rachford Splitting Algorithms for Nonconvex Federated Composite Optimization

1 Introduction . Training machine learning models in a centralized fashion becomes more challenging and marginally inaccessible for a large number of users , especially when the size of datasets and models is growing substantially larger . Consequently , training algorithms using decentralized and distributed approaches comes in as a natural replacement . Among several approaches , federated learning ( FL ) has received tremendous attention in the past few years since it was first introduced in [ 18 , 30 ] . In this setting , a central server coordinates between many local users ( also called agents or devices ) to perform their local updates , then the global model will get updated , e.g. , by averaging or aggregating local models . Challenges . FL provides a promising solution for many machine learning applications such as learning over smartphones or across organizations , and internet of things , where privacy protection is one of the most critical requirements . However , this training mechanism faces a number of fundamental challenges , see , e.g. , [ 31 ] . First , when the number of users gets substantially large , it creates communication bottleneck during model exchange process between server and users . Second , the local data stored in each local user may be different in terms of sizes and distribution which poses a challenge : data or statistical heterogeneity . Third , the variety of users with different local storage , computational power , and network connectivity participating into the system also creates a major challenge , known as system heterogeneity . This challenge also causes unstable connection 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . between server and users , where some users may be disconnected from the server or simply dropped out during training . In practice , we can expect only a subset of users to participate in each round of communication . Another challenge in FL is privacy concern . Accessing and sharing local raw data is not permitted in FL . In addition , distributed methods exchange the objective gradient of local users , and private data can be exposed from the shared model such as the objective gradients [ 51 ] . Therefore , FL methods normally send the global model to each user at the start of each communication round , each user will perform its local update and send back only the necessary update for aggregation . Our goal and approach . Our goal in this paper is to further and simultaneously address these fundamental challenges by proposing two new algorithms to train the underlying common optimization model in FL . Our approach relies on a novel combination between randomized block-coordinate strategy , nonconvex Douglas-Rachford ( DR ) splitting , and asynchronous implementation . While each individual technique or partial combinations is not new , our combination of three as in this paper appears to be the first in the literature . To the best of our knowledge , this is the first work developing randomized block-coordinate DR splitting methods for nonconvex composite FL optimization models , and they are fundamentally different from some works in the convex setting , e.g. , [ 7 , 8 ] . Contribution . Our contribution can be summarized as follows . ( a ) We develop a new FL algorithm , called FedDR ( Federated Douglas-Rachford ) , by combining the well-known DR splitting technique and randomized block-coordinate strategy for the common nonconvex composite optimization problem in FL . Our algorithm can handle nonsmooth convex regularizers and allows inexact evaluation of the underlying proximal operators as in FedProx or FedPD . It also achieves the best known O ( ε−2 ) communication complexity for finding a stationary point under standard assumptions ( Assumptions 2.1- 2.2 ) , where ε is a given accuracy . More importantly , unlike FedSplit [ 33 ] and FedPD [ 49 ] , which require full user participation to achieve convergence , our analysis does allow partial participation by selecting a subset of users to perform update at each communication round . ( b ) Next , we propose an asynchronous algorithm , asyncFedDR , where each user can asynchronously perform local update and periodically send the update to the server for proximal aggregation . We show that asyncFedDR achieves the same communication complexity O ( ε−2 ) as FedDR ( up to a constant factor ) under the same standard assumptions . This algorithm is expected to simultaneously address all challenges discussed above . Let us emphasize some key points of our contribution . First , the best known O ( ε−2 ) communication complexity of our methods matches the lower bound complexity up to a constant factor as shown in [ 49 ] , even with inexact evaluation of the objective proximal operators . Second , our methods rely on a DR splitting technique for nonconvex optimization and can handle possibly nonsmooth convex regularizers , which allows us to deal with a larger class of applications and with constraints [ 47 ] . Furthermore , it can also handle both statistical and system heterogeneity as discussed in FedSplit [ 33 ] and FedPD [ 49 ] . However , FedSplit only considers the convex case , and both FedSplit and FedPD require all users to update at each communication round , making them less practical and applicable in FL . Our methods only require a subset of users or even one user to participate in each communication round as in FedAvg or FedProx . In addition , our aggregation step on the server is different from most existing works due to a proximal step on the regularizer . It is also different from [ 47 ] . Third , as FedProx [ 23 ] , we allow inexact evaluation of users ’ proximal operators with any local solver ( e.g. , local SGD or variance reduced methods ) and with adaptive accuracies . Finally , requiring synchronous aggregation at the end of each communication round may lead to slow-down in training due to the heterogeneity in computing power and communication capability of local users . It is natural to have asynchronous update from local users as in , e.g. , [ 34 , 35 , 39 ] . Our asynchronous variant , asyncFedDR , can fairly address this challenge . Moreover , it uses a general probabilistic model recently introduced in [ 5 ] , which allows us to capture the variety of asynchronous environments and architectures compared to existing methods , e.g. , [ 39 , 44 ] . Related work and comparison . Federated Averaging ( FedAvg ) is perhaps the earliest method used in FL . In FedAvg , users perform stochastic gradient descent ( SGD ) updates for a number of epochs then send updated models to server for aggregation . FedAvg ’ s practical performance has been shown in many early works , e.g. , [ 18 , 29 , 48 ] and tends to become the most popular method for solving FL applications . [ 26 ] show that local SGD where users perform a number of local updates before global communication takes place as in FedAvg may offer benefit over minibatch SGD . Similar comparison between minibatch SGD and local SGD has been done in [ 42 , 43 ] . Analyzing convergence of FedAvg was very challenging at its early time due to the complexity in its update as well as data heterogeneity . One of the early attempt to show the convergence of FedAvg is in [ 39 ] for convex problems under the iid data setting and a set of assumptions . [ 45 ] also considers local SGD in the nonconvex setting . Without using an additional bounded gradient assumption as in [ 39 , 45 ] , [ 41 ] improves the complexity for the general nonconvex setting while [ 11 ] uses a Polyak-Łojasiewicz ( PL ) condition to improve FedAvg ’ s convergence results . In heterogeneous data settings , [ 17 ] analyzes local GD , where users performs gradient descent ( GD ) updates instead of SGD . The analysis of FedAvg for non-iid data is given in [ 24 ] . The analysis of local GD/SGD for nonconvex problems has been studied in [ 13 ] . However , FedAvg might not converge with non-iid data as shown in [ 33 , 49 , 50 ] . FedProx [ 23 ] is an extension of FedAvg , which deals with heterogeneity in federated networks by introducing a proximal term to the objective in local updates to improve stability . FedProx has been shown to achieve better performance than FedAvg in heterogeneous setting . Another method to deal with data heterogeneity is SCAFFOLD [ 16 ] which uses a control variate to correct the “ client-drift '' in local update of FedAvg . MIME [ 15 ] is another framework that uses control variate to improve FedAvg for heterogeneous settings . However , SCAFFOLD and MIME require to communicate extra information apart from local models . Compared to aforementioned works , our methods deal with nonconvex problems under standard assumptions and with composite settings . FedSplit [ 33 ] instead employs a Peaceman-Rachford splitting scheme to solve a constrained reformulation of the original problem . In fact , FedSplit can be viewed as a variant of Tseng ’ s splitting scheme [ 1 ] applied to FL . [ 33 ] show that FedSplit can find a solution of the FL problem under only convexity without imposing any additional assumptions on system or data homogeneity . [ 49 ] proposes FedPD , which is essentially a variant of the standard augmented Lagrangian method in nonlinear optimization . Other algorithms for FL can be found , e.g. , in [ 6 , 10 , 12 , 14 , 25 , 46 ] . Our approach in this paper relies on nonconvex DR splitting method , which can handle the heterogeneity as discussed in [ 33 ] . While the DR method is classical , its nonconvex variants have been recently studied e.g. , in [ 9 , 21 , 40 ] . However , the combination of DR and randomized block-coordinate strategy remains limited [ 7 , 8 ] even in the convex settings . Alternatively , asynchronous algorithms have been extensively studied in the literature , also for FL , see , e.g. , [ 2 , 34 , 35 ] . For instance , a recent work [ 44 ] analyzes an asynchronous variant of FedAvg under bounded delay assumption and constraint on the number of local updates . [ 39 ] proposes an asynchronous local SGD to solve convex problems under iid data . However , to our best knowledge , there exists no asynchronous method using DR splitting techniques with convergence guarantee for FL . In addition , most existing algorithms only focus on non-composite settings . Hence , our work here appears to be the first . Content . The rest of this paper is organized as follows . Section 2 states our FL optimization model and our assumptions . Section 3 develops FedDR and analyzes its convergence . Section 4 considers an asynchronous variant , asyncFedDR . Section 5 is devoted for numerical experiments . Due to space limit , all technical details and proofs can be found in Supplementary Document ( Supp . Doc . ) . 2 Nonconvex Optimization Models in Federated Learning . The underlying optimization model of many FL applications can be written into the following form : min x∈Rp { F ( x ) : = f ( x ) + g ( x ) = 1 n n∑ i=1 fi ( x ) + g ( x ) } , ( 1 ) where n is the number of users , and each fi is a local loss of the i-th user , which is assumed to be nonconvex and L-smooth ( see Assumptions 2.1 and 2.2 below ) , and g is a proper , closed , and convex regularizer . Apart from these assumptions , we will not make any additional assumption on ( 1 ) . We emphasize that the use of regularizers g has been motivated in several works , including [ 47 ] . Let dom ( F ) : = { x ∈ Rp : F ( x ) < +∞ } be the domain of F and ∂g be the subdifferential of g [ 1 ] . Since ( 1 ) is nonconvex , we only expect to find a stationary point , which is characterized by the following optimality condition . Definition 2.1 . If 0 ∈ ∇f ( x∗ ) + ∂g ( x∗ ) , then x∗ is called a [ first-order ] stationary point of ( 1 ) . The algorithms for solving ( 1 ) developed in this paper will rely on the following assumptions . Assumption 2.1 ( Boundedness from below ) . dom ( F ) 6= ∅ and F ? : = infx∈Rp F ( x ) > −∞ . Assumption 2.2 ( L-smoothness ) . All functions fi ( · ) for i ∈ [ n ] : = { 1 , · · · , n } are L-smooth , i.e. , fi is continuously differentiable and there exists L ∈ ( 0 , +∞ ) such that ‖∇fi ( x ) −∇fi ( y ) ‖ ≤ L‖x− y‖ , ∀x , y ∈ dom ( fi ) . ( 2 ) Assumptions 2.1 and 2.2 are very standard in nonconvex optimization . Assumption 2.1 guarantees the well-definedness of ( 1 ) and is independent of algorithms . Assuming the same Lipschitz constant L for all fi is not restrictive since if fi is Li-smooth , then by scaling variables of its constrained formulation ( see ( 11 ) in Supp . Doc . ) , we can get the same Lipschitz constant L of all fi . Proximal operators and evaluation . Our methods make use of the proximal operators of both fi and g. Although fi is L-smooth and nonconvex , we still define its proximal operator as proxηfi ( x ) : = argminy { fi ( y ) + 1 2η‖y − x‖ 2 } , ( 3 ) where η > 0 . Even fi is nonconvex , under Assumption 2.2 , if we choose 0 < η < 1L , then proxηfi is well-defined and single-valued . Evaluating proxηfi requires to solve a strongly convex program . If proxηfi can only be computed approximately up to an accuracy ≥ 0 to obtain z , denoted by x+ : ≈ proxηfi ( x ) , if ‖x+ − proxηfi ( x ) ‖ ≤ i . Note that instead of absolute error , one can also use a relative error as ‖x+ − proxηfi ( x ) ‖ ≤ i‖x+ − x‖ as in [ 37 ] . For the convex function g , its proximal operator proxηg is defined in the same way as ( 3 ) . Evaluating proxηfi can be done by various existing methods , including local SGD and accelerated GD-type algorithms . However , this is not our focus in this paper , and therefore we do not specify the subsolver for evaluating proxηfi . Gradient mapping . As usual , let us define the following gradient mapping of F in ( 1 ) . Gη ( x ) : = 1η ( x− proxηg ( x− η∇f ( x ) ) ) , η > 0 . ( 4 ) Then , the optimality condition 0 ∈ ∇f ( x∗ ) + ∂g ( x∗ ) of ( 1 ) is equivalent to Gη ( x∗ ) = 0 . However , in practice , we often wish to find an ε-approximate stationary point to ( 1 ) defined as follows . Definition 2.2 . If x̃ ∈ dom ( F ) satisfies E [ ‖Gη ( x̃ ) ‖2 ] ≤ ε2 , then x̃ is called an ε-stationary point of ( 1 ) , where the expectation is taken overall the randomness generated by the underlying algorithm . Note that , for Gη ( x̃ ) to be well-defined , we require x̃ ∈ dom ( F ) . In our algorithms below , this requirement is fulfilled if x̃ ∈ dom ( f ) , which is often satisfied in practice as dom ( f ) = Rp .

Summary: This paper proposes FedDR, which extends FedSplit by allowing non-convex objective, partial participation and convex composite regularizers. This work also studies the asynchronous variant of FedDR. Detailed proofs and experiments are provided.