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 presents a new method for federated learning with a nonconvex but smooth objective function. The proposed method is based on Douglas Rachford Splitting to decompose the objective and randomized coordinate methodology to allow a random selection of users to update at each iteration. An extension to asynchronous settings is presented, allowing modeling of communication delays and similar practical considerations. The convergence of these methods (potential using inexact proximal calculations) is bounded by the typical O(\epsilon^-2) rate. Reasonable numerics are given showing the potential effectiveness of such a scheme on synthetic and real datasets.