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: This paper studies federated learning algorithms, where a collection of computing nodes seek to collectively optimize a global objective through parallelized updates executed at a server. Different from prior works which execute local-only stochastic gradient updates on non-convex losses, here a stochastic gradient Langevin update is developed, which additionally incorporates randomized Gaussian perturbations. Convergence theory is presented for the proposed scheme, which establishes point-wise  convergence in mean, as well as convergence in distribution according to the Wasserstein metric. Variations of the noise generating process which introduce heterogeneity and privacy preservation are also proposed and analyzed.