On Convergence of Federated Averaging Langevin Dynamics

1 INTRODUCTION . Federated learning ( FL ) allows multiple parties to jointly train a consensus model without sharing user data . Compared to the classical centralized learning regime , federated learning keeps training data on local clients , such as mobile devices or hospitals , where data privacy , security , and access rights are a matter of vital interest . This aggregation of various data resources heeding privacy concerns yields promising potential in areas of internet of things Chen et al . ( 2020 ) , healthcare Li et al . ( 2020d ; 2019b ) , text data Huang et al . ( 2020 ) , and fraud detection Zheng et al . ( 2020 ) . A standard formulation of federated learning is a distributed optimization framework that tackles communication costs , client robustness , and data heterogeneity across different clients Li et al . ( 2020a ) . Central to the formulation is the efficiency of the communication , which directly motivates the communication-efficient federated averaging ( FedAvg ) McMahan et al . ( 2017 ) . FedAvg introduces a global model to synchronously aggregate multi-step local updates on the available clients and yields distinctive properties in communication . However , FedAvg often stagnates at inferior local modes empirically due to the data heterogeneity across the different clients Charles & Konečnỳ ( 2020 ) ; Woodworth et al . ( 2020 ) . To tackle this issue , Karimireddy et al . ( 2020 ) ; Pathaky & Wainwright ( 2020 ) proposed stateful clients to avoid the unstable convergence , which are , however , not scalable with respect to the number of clients in applications with mobile devices Al-Shedivat et al . ( 2021 ) . In addition , the optimization framework often fails to quantify the uncertainty accurately for the parameters of interest , which are crucial for building estimators , hypothesis tests , and credible intervals . Such a problem leads to unreliable statistical inference and casts doubts on the credibility of the prediction tasks or diagnoses in medical applications . To unify optimization and uncertainty quantification in federated learning , we resort to a Bayesian treatment by sampling from a global posterior distribution , where the latter is aggregated by infrequent communications from local posterior distributions . We adopt a popular approach for inferring posterior distributions for large datasets , the stochastic gradient Markov chain Monte Carlo ( SGMCMC ) method Welling & Teh ( 2011 ) ; Vollmer et al . ( 2016 ) ; Teh et al . ( 2016 ) ; Chen et al . ( 2014 ) ; Ma et al . ( 2015 ) , which enjoys theoretical guarantees beyond convex scenarios Raginsky et al . ( 2017 ) ; Zhang et al . ( 2017 ) ; Mangoubi & Vishnoi ( 2018 ) ; Ma et al . ( 2019 ) . In particular , we examine in the federated learning setting the efficacy of the stochastic gradient Langevin dynamics ( SGLD ) algorithm , which differs from stochastic gradient descent ( SGD ) in an additionally injected noise for exploring the posterior . The close resemblance naturally inspires us to adapt the optimization-based FedAvg to a distributed sampling framework . Similar ideas have been proposed in federated posterior averaging Al-Shedivat et al . ( 2021 ) , where empirical study and analyses on Gaussian posteriors have shown promising potential of this approach . Compared to the appealing theoretical guarantees of optimization-based algorithms in federated learning Pathaky & Wainwright ( 2020 ) ; Al-Shedivat et al . ( 2021 ) , the convergence properties of approximate sampling algorithms in federated learning is far less understood . To fill this gap , we proceed by asking the following question : Can we build a unified algorithm with convergence guarantees for sampling in FL ? In this paper , we make a first step in answering this question in the affirmative . We propose the federated averaging Langevin dynamics ( FA-LD ) for posterior inference beyond the Gaussian distribution . We list our contributions as follows : • We present a novel non-asymptotic convergence analysis for FA-LD from simulating strongly log-concave distributions on non-i.i.d data when the learning rate is fixed . The frequently used bounded gradient assumption of ` 2 norm in FedAvg optimization is not required . • The convergence analysis indicates that the injected noise , the data heterogeneity , and the stochastic-gradient noise are all driving factors that affect the convergence . Such an analysis also provides a concrete guidance on the optimal selection of the number of local updates . • We present a convergence result for FA-LD with decaying learning rates . This strategy accelerates the computation by a logarithmic factor to achieve the precision . • The algorithm yields appealing extensions : ( 1 ) we can choose to inject either independent or correlated noise across local clients , yielding a trade-off between accuracy and efficacy of federation ; ( 2 ) we can choose whether to activate all the devices to avoid the straggler ’ s effect in real-world applications . Roadmap . In Section 2 , we discuss the related work and literature . In Section 3 , we present the preliminary knowledge . In Section 4 , we propose the federated averaging Langevin dynamics algorithm for posterior inference . In Section 5 , we lay out the required assumptions , sketch the proof , and show the theoretical convergence results . In Section 6 , we conclude our work . 2 RELATED WORK . Federated Learning Current federated learning follows two paradigms . The first paradigm asks every client to learn the model using private data and communicate in model parameters . The second one uses encryption techniques to guarantee secure communication between clients . In this paper , we focus on the first paradigms Dean et al . ( 2012 ) ; Shokri & Shmatikov ( 2015 ) ; McMahan et al . ( 2016 ; 2017 ) ; Huang et al . ( 2021 ) . There is a long list of works showing provable convergence algorithm for FedAvg types of algorithms in the field of optimization Li et al . ( 2020c ; 2021 ) ; Huang et al . ( 2021 ) ; Khaled et al . ( 2019 ) ; Yu et al . ( 2019 ) ; Wang et al . ( 2019 ) ; Karimireddy et al . ( 2020 ) . One line of research Li et al . ( 2020c ) ; Khaled et al . ( 2019 ) ; Yu et al . ( 2019 ) ; Wang et al . ( 2019 ) ; Karimireddy et al . ( 2020 ) focuses on standard assumptions in optimization ( such as , convex , smooth , strongly-convex , bounded gradient ) . The other line of work Li et al . ( 2021 ) ; Huang et al . ( 2021 ) proves the convergence in the regime where the model of interest is an over-parameterized neural network ( also called NTK regime Jacot et al . ( 2018 ) ) . Scalable Monte Carlo methods SGLD Welling & Teh ( 2011 ) is the first stochastic gradient Monte Carlo method that tackles the scalability issue in big data problems . Ever since , variants of stochastic gradient Monte Carlo methods were proposed to accelerate the simulations by utilizing more general Markov dynamics Ma et al . ( 2015 ; 2018 ) ; Chen et al . ( 2014 ) , Hessian approximation Ahn et al . ( 2012 ) , parallel tempering Deng et al . ( 2020 ) , as well as higher-order numerical schemes Chen et al . ( 2015 ) ; Li et al . ( 2019c ) ; Cheng et al . ( 2018 ) ; Ma et al . ( 2021 ) ; Mou et al . ( 2021 ) ; Shen & Lee ( 2019 ) . Distributed Monte Carlo methods Sub-posterior aggregation was initially proposed in Neiswanger et al . ( 2013 ) ; Wang & Dunson ; Minsker et al . ( 2014 ) to accelerate MCMC methods to cope with large datasets . Other parallel MCMC algorithms Nishihara et al . ( 2014 ) ; Ahn et al . ( 2014 ) ; Chen et al . ( 2016 ) ; Chowdhury & Jermaine ( 2018 ) ; Li et al . ( 2019a ) propose to improve the efficiency of Monte Carlo computation in distributed or asynchronous systems . Gürbüzbalaban et al . ( 2021 ) proposed stochastic gradient Monte Carlo methods in decentralized systems . Al-Shedivat et al . ( 2021 ) introduced empirical studies of posterior averaging in federated learning . Notation For any positive integer n , we use [ n ] to denote the set { 1 , 2 , · · · , n } . Let N denote the number of clients . For each c ∈ [ N ] , we use f c and ∇f c as the loss function and gradient of the function f c in client c. ∇f̃ c ( · ) is the unbiased stochastic gradient of ∇f c. In addition , we denote pc as the weight of the c-th client such that pc = nc∑N i=1 ni ∈ ( 0 , 1 ) , where nc > 0 is the number of data points in the c-th client . Let T denote the number of global steps to achieve the precision . Let K denote the number of local steps and hence T /K denotes the number of communications . 3 PRELIMINARIES . 3.1 AN OPTIMIZATION PERSPECTIVE ON FEDERATED AVERAGING . Federated averaging ( FedAvg ) is a standard algorithm in federated learning and is typically formulated into a distributed optimization framework as follows min θ ` ( θ ) : = ∑N c=1 ` c ( θ ) ∑N c=1 nc , ` c ( θ ) : = nc∑ i=1 l ( θ ; xc , i ) , ( 1 ) where θ ∈ Rd , l ( θ ; xc , j ) is a certain loss function based on θ and the data point xc , j . One iterate of the FedAvg algorithm requires the following three steps : • Broadcast : The center server broadcasts the latest model , θk , to all local clients . • Local updates : For any c ∈ [ N ] , the c-the client first sets θck = θk and then conducts K ≥ 1 local steps : βck+1 = θ c k − η∇˜̀c ( θck ) , where η is the learning rate and∇˜̀c is the unbiased estimate of the exact gradient∇ ` c . • Synchronization : The local models are aggregated into a unique model θk+K : =∑N c=1 pcβ c k+K and sent to the center server . From the optimization perspective , Li et al . ( 2020c ) proved the convergence of the FedAvg algorithm on non-i.i.d data such that a larger number of local steps K and a higher order of data heterogeneity slows down the convergence . Notably , Eq . ( 1 ) can be interpreted as maximizing the likelihood function , which is a special case of maximum a posteriori estimation ( MAP ) given a uniform prior . 3.2 STOCHASTIC GRADIENT LANGEVIN DYNAMICS . Posterior inference offers the exact uncertainty quantification ability of the predictions . A popular method for posterior inference with large dataset is the stochastic gradient Langevin dynamics ( SGLD ) Welling & Teh ( 2011 ) , which injects additional noise into the stochastic gradient and adapts an optimization algorithm to a sampling one θk+1 = θk − η∇f̃ ( θk ) + √ 2τηξk , where τ is the temperature and ξk is a standard d-dimensional Gaussian vector . f ( θ ) : = ∑N c=1 ` c ( θ ) is a energy function . f̃ ( θ ) is a unbiased estimate of f ( θ ) . In the longtime limit , a well known result is that θk converges weakly to the distribution π ( θ ) ∝ exp ( −f ( θ ) /τ ) Teh et al . ( 2016 ) as η → 0 .

Summary: In this work, the authors:  1.  propose a federated learning version of Langevin diffusion sampling.   2. develop theoretical guarantees for FA-LD for strongly log-concave distributions with non-i.i.d data and study how the injected noise and the stochastic gradient noise, the heterogeneity of data, and the varying learning rates affect the convergence.  3. analyze the partial participation setting.