Year: 2021

Abstract:
We propose a federated averaging Langevin algorithm (FA-LD) for uncertainty quantification and mean predictions with distributed clients. In particular, we generalize beyond normal posterior distributions and consider a general class of models. We 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 stochasticgradient noise, the heterogeneity of data, and the varying learning rates affect the convergence. Such an analysis sheds light on the optimal choice of local updates to minimize communication cost. Important to our approach is that the communication efficiency does not deteriorate with the injected noise in the Langevin algorithms. In addition, we examine in our FA-LD algorithm both independent and correlated noise used over different clients. We observe that there is also a trade-off between federation and communication cost there. As local devices may become inactive in the federated network, we also show convergence results based on different averaging schemes where only partial device updates are available.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARIES

heading: 3.1 AN OPTIMIZATION PERSPECTIVE ON FEDERATED AVERAGING
text: 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.

heading: 3.2 STOCHASTIC GRADIENT LANGEVIN DYNAMICS
text: 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.

heading: 4 POSTERIOR INFERENCE VIA FEDERATED AVERAGING LANGEVIN DYNAMICS
text: The increasing concern for uncertainty estimation in federated learning motivates us to consider the simulation of the distribution π(θ) ∝ exp(−f(θ)/τ) with distributed clients.
Problem formulation We propose the federated averaging Langevin dynamics (FA-LD) based on the FedAvg framework in section 3.1. We follow the same broadcast step and synchronization step but propose to inject random noises for local updates. In particular, we consider the following scheme: for any c ∈ [N ], the c-the client first sets θck = θk and then conducts K ≥ 1 local steps:
βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητΞck, (2)
where ∇f c(θ) = 1pc∇` c(θ). ∇f̃ c(θ) is the unbiased estimate of ∇f c(θ) and Ξck is an independent Gaussian vector to be defined later.
Summing Eq. (2) from clients c = 1 to N , we have the aggregated stochastic process as follows βk+1 = θk − η∇f̃(θk) + √ 2ητξk,
where
βk = N∑ c=1 pcβ c k, θk = N∑ c=1 pcθ c k, ∇f̃(θk) = N∑ c=1 pc∇f̃ c(θck), ξk = N∑ c=1 pcΞ c k. (3)
By the nature of the synchronization step, we always have βk = θk whether k+ 1 mod K = 0 or not. In what follows, we can write
θk+1 = θk − η∇f̃(θk) + √ 2ητξk, (4)
which resembles the SGLD algorithm except that the construction of stochastic gradients is different and θk is not accessible when k mod K 6= 0. Since our target is to simulate from π(θ) ∝ exp(−f(θ)/τ), we expect that ξk is a standard Gaussian vector. By the concentration property of independent Gaussian variables, it is natural to set Ξck = ξ c k/ √ pc so that
ξk = ∑N c=1 pcΞ c k = ∑N c=1 √ pcξ c k and ξ c k is also a standard Gaussian vector. Now we present it in Algorithm 1.
Algorithm 1 Federated averaging Langevin dynamics algorithm (FA-LD), informal version of Algorithm 4. ηk is the learning rate at iteration k. τ is the temperature. Denote by θck the model parameter in the c-th client at the k-th step. Denote the immediate result of one step SGLD update from θck by β c k. ξ c k is an independent standard d-dimensional Gaussian vector at iteration k for each client c ∈ [N ]. A global synchronization is conducted every K steps. 1:
βck+1 = θ c k − ηk∇f̃ c(θck) + √ 2ηkτ/pcξ c k, (5)
2:
θck+1 =  βck+1 if k + 1 mod K 6= 0∑N c=1 pcβ c k+1 if k + 1 mod K = 0.
(6)
Algorithm 2 Hybrid federated averaging Langevin dynamics algorithm (hFA-LD), informal version of Algorithm 5. ξ̇k is a d-dimensional Gaussian vector shared by all the clients; ξck is an independent standard d-dimensional Gaussian vector at iteration k for each client c ∈ [N ]. ρ denotes the correlation coefficient.
1: βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητρ2ξ̇k + √ 2η(1− ρ2)τ/pcξck,
2:
θck+1 =  βck+1 if k + 1 mod K 6= 0∑N c=1 pcβ c k+1 if k + 1 mod K = 0.
We observe that the local process in Eq. (5) maintains a temperature τ/pc > τ to converge to the stationary distribution π. Such a mechanism may limit the disclosure of individual data and shows a potential to ensure a higher level of privacy.

heading: 5 CONVERGENCE ANALYSIS
text: In this section, we show that FA-LD converges to the stationary distribution π(θ) in the 2-Wasserstein (W2) distance at a rate of O(1/ √ T ) for strongly log-concave and smooth density. The W2 distance is defined between a pair of Borel probability measures µ and ν on Rd as follows
W2(µ, ν) := inf γ2∈Couplings(µ,ν)
(∫ ‖βµ − βν‖22dγ2(βµ,βν) ) 1 2 ,
where ‖ · ‖2 denotes the `2 norm on Rd and the pair of random variables (βµ,βν) ∈ Rd × Rd is a coupling with the marginals following L(βµ) = µ and L(βν) = ν. Note that L(·) denotes a distribution of a random variable. Such a distance is more appealing than the total variation or the Kullback–Leibler divergence in statistical machine learning applications for providing the estimates of the first and second order moments.

heading: 5.1 NOTATION AND ASSUMPTIONS
text: We make standard assumptions on the smoothness and convexity of the functions f1, f2, · · · , fN , which naturally yields appealing tail properties of the stationary measure π. Thus, we no longer require a restrictive assumption on the bounded gradient in `2 norm as in Koloskova et al. (2019); Yu et al. (2019); Li et al. (2020c). In addition, to control the distance between ∇f c and ∇f̃ c, we also assume a bounded variance of the stochastic gradient in assumption 5.3.
Assumption 5.1 (Smoothness). For each c ∈ [N ], we say f c is L-smooth if for some L > 0
f c(y) ≤ f c(x) + 〈∇f c(x), y − x〉+ L 2 ‖y − x‖22 ∀x, y ∈ Rd.
Assumption 5.2 (Strongly convex). For each c ∈ [N ], f c is m-strongly convex if for some m > 0
f c(x) ≥ f c(y) + 〈∇f c(y), x− y〉+ m 2 ‖y − x‖22 ∀x, y ∈ Rd.
Assumption 5.3 (Bounded variance, informal version of Assumption A.3). For each c ∈ [N ], the variance of noise in the stochastic gradient∇f̃ c(x) in each client is upper bounded such that
E[‖∇f̃ c(x)−∇f c(x)‖22] ≤ σ2d, ∀x ∈ Rd.
Quality of non-i.i.d data Denote by θ∗ the global minimum of f . Next, we quantify the degree of the non-i.i.d data by γ := maxc∈[N ] ‖∇f c(θ∗)‖2, which is a non-negative constant and yields a larger scale if the data is less identically distributed.

heading: 5.2 PROOF SKETCH
text: The proof hinges on showing the one-step result in the W2 distance. To facilitate the analysis, we first define an auxiliary continuous-time processes (θ̄t)t≥0 without communication concerns
dθ̄t = −∇f(θ̄t)dt+ √ 2τdW t, (7)
where θ̄t = ∑N c=1 pcθ̄ c t , ∇f(θ̄t) = ∑N c=1 pc∇f c(θ̄ct ), θ̄ct is the continuous-time variable at client c, and W is a d-dimensional Brownian motion. The continuous-time algorithm is known to converge to the stationary distribution π(θ) ∝ e− f(θ) τ , where f(θ) = ∑N c=1 pcf
c(θ). Assume that θ̄0 simulates from the stationary distribution π, then it follows that θ̄t ∼ π for any t ≥ 0.

heading: 5.2.1 DOMINATED CONTRACTION IN FEDERATED LEARNING
text: The first target is to show a certain contraction property of ‖β − θ − η(∇f(β)−∇f(θ))‖22 based on distributed clients with infrequent communications. Consider a standard decomposition
‖β − θ − η(∇f(β)−∇f(θ))‖22
= ‖β − θ‖22 − 2η 〈β − θ,∇f(β)−∇f(θ)〉︸ ︷︷ ︸ I +η2 ‖∇f(β)−∇f(θ)‖22 .
Using Eq.(3), we decompose I and apply Jensen’s inequality to obtain the lower bound of I . In what follows, we have the following lemma. Lemma 5.4 (Dominated contraction property, informal version of Lemma B.1). Assume assumptions 5.1 and 5.2 hold. For any learning rate η ∈ (0, 1L+m ], any {θ c}Nc=1, {βc}Nc=1 ∈ Rd, we have
‖β − θ − η(∇f(β)−∇f(θ))‖22 ≤ (1− ηm) · ‖β − θ‖ 2 2 + 4ηL N∑ c=1 pc · (‖βc − β‖22 + ‖θc − θ‖22)︸ ︷︷ ︸ divergence term ,
where β = ∑N c=1 pcβ c, θ = ∑N c=1 pcθ c, ∇f(θ) = ∑N c=1 pc∇f c(θc), and ∇f(β) =∑N
c=1 pc∇f c(βc). It implies that as long as the local parameters θc, βc and global θ, β don’t differ each other too much, we can guarantee the desired convergence. In a special case when the communication is conducted at every iteration, the divergence term disappears and recovers the standard contraction Dalalyan & Karagulyan (2019).

heading: 5.2.2 BOUNDING DIVERGENCE
text: The following result shows that given a finite number of local steps K, the divergence between θc in local client and θ in the center is bounded in `2 norm. Notably, since the Brownian motion leads to a lower order term O(η) instead of O(η2), a naïve proof framework such as Li et al. (2020c) may lead to a crude upper bound for the final convergence. Lemma 5.5 (Bounded divergence, informal version of Lemma B.3). Assume assumptions 5.1, 5.2, and 5.3 hold. For any learning rate η ∈ (0, 2/m) and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have the `2 upper bound of the divergence between local clients and the center as follows N∑ c=1 pcE‖θck − θk‖22 ≤ O(K2η2d) +O(Kηd).
The result also relies on showing a uniform upper bound in `2 norm, which avoids making extra bounded gradient assumptions.

heading: 5.2.3 COUPLING TO THE STATIONARY PROCESS
text: Note that θ̄t is initialized from the stationary distribution π. The solution to the continuous-time process Eq.(7) follows:
θ̄t = θ̄0 − ∫ t 0 ∇f(θ̄s)ds+ √ 2τ ·W t, ∀t ≥ 0. (8)
Set t→ (k+ 1)η and θ̄0 → θ̄kη for Eq.(8) and consider a synchronous coupling such that W(k+1)η− Wkη := √ ηξk is used to cancel the noise terms, we have
θ̄(k+1)η = θ̄kη − ∫ (k+1)η kη ∇f(θ̄s)ds+ √ 2τηξk. (9)
Subtracting Eq.(4) from Eq.(9) and taking square and expectation on both sides yield that
E‖θ̄(k+1)η − θk+1‖22 ≤ (1− ηm/2) · E‖θ̄kη − θk‖22 + divergence term + time error.
Eventually, we arrive at the one-step error bound for establishing the convergence results. Lemma 5.6 (One step update, informal version of Lemma B.5). Assume assumptions 5.1, 5.2, and 5.3 hold. Consider Algorithm 1 with any learning rate η ∈ (0, 12L ) and ‖θ c 0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], where θ∗ is the global minimum for the function f . Then W 22 (µk+1, π) ≤ (1− ηm/2) ·W 22 (µk, π) +O(η2d(K2 + κ)),
where µk denotes the probability measure of θk and κ = L/m is the condition number.

heading: 5.3 FULL DEVICE PARTICIPATION

heading: 5.3.1 CONVERGENCE BASED ON INDEPENDENT NOISE
text: When the synchronization step is conducted at every iteration k, the FA-LD algorithm is essentially the standard SGLD algorithm Welling & Teh (2011). Theoretical analysis based on the 2-Wasserstein distance has been established in Durmus & Moulines (2016); Dalalyan (2017a); Dalalyan & Karagulyan (2019). However, in scenarios of K > 1 with distributed clients, a divergence between the global variable θk and local variable θck appears and unavoidably affects the performance. The upper bound on the sampling error is presented as follows. Theorem 5.7 (Main result, informal version of Theorem B.6). Assume assumptions 5.1, 5.2, and 5.3 hold. Consider Algorithm 1 with a fixed learning rate η ∈ (0, 12L ] and ‖θ c 0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have †
W2(µk, π) ≤ (1− ηm/4)k · (√ 2d ( D + √ τ/m )) + 30κ √ ηmd · √ (K2 + κ)H0.
where µk denotes the probability measure of θk at iteration k, K denotes the number of local updates, κ := L/m, γ := maxc∈[N ] ‖∇f c(θ∗)‖2, and H0 := D2 + maxc∈[N ] τ mpc + γ 2 m2d + σ2 m2 .
We observe that the initialization, the scale of the injected noise, the heterogeneity of the data, and the noise in the stochastic gradient all affect the convergence. Similar to the result of Li et al. (2020c), FA-LD with K-local steps resembles the behaviour of one-step SGLD with a large learning rate.
Optimal choice of K. To ensure the algorithm to achieve the precision based on the total number of steps T and the learning rate η, we can set
30κ √ ηmd · √ (K2 + κ)H0 ≤ /2, exp ( − ηm 4 T ) · √ 2d(D + √ τ/m) ≤ /2.
This readily leads to ηm ≤ O (
2
dκ2(K2 + κ)H0
) , T ≥ Ω ( log(d/ )
mη
) .
Plugging into the upper bound of ηm, it implies that to reach the precision , it suffices to set
T = Ω( −2dκ2(K2 + κ)H0 · log(d/ )). (10)
It’s obvious that H0 = Ω(D2) = Ω(1), thus we can conclude that the number of communication rounds is around the order
T K = Ω
( K + κ
K
) ,
where the value of T K first decreases and then increases with respect to K, indicating that setting K either too large or too small may lead to high communication costs and hurt the performance. Ideally, K should be selected in the scale of Ω( √ κ). Combining the definition of T in Eq. (10), this suggests an interesting result that the optimal K for FA-LD should be in the order of O( √ T ). Similar results have been achieved by Stich (2019); Li et al. (2020c).

heading: 5.3.2 CONVERGENCE GUARANTEES VIA VARYING LEARNING RATES
text: Theorem 5.8 (Informal version of Theorem B.7). Assume assumptions 5.1, 5.2, and 5.3 hold. Consider Algorithm 1 with an initialization satisfying ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ] and the varying learning rate following
ηk = 1
2L+ (1/12)mk , k = 1, 2, · · · .
Then for any k ≥ 0, we have W2(µk, π) ≤ 45κ √ (K2 + κ)H0 · ( ηkmd )1/2 , ∀k ≥ 0.
†For ease of presentation, we report the result based on K2 instead of (K − 1)2. The upper bound based on (K − 1)2 is detailed in the supplementary file.
Note that the above result implies that to achieve the precision , we require W2(µk, π) ≤ 45κ √ (K2 + κ)H0 · ( md
2L+ (1/12)mk
)1/2 ≤ .
We therefore require Ω( −2d) iterations to achieve the precision , which improves the Ω( −2d log(d/ )) rate for FA-LD with a fixed learning rate by a O(log(d/ )) factor.

heading: 5.3.3 PRIVACY-ACCURACY TRADE-OFF VIA CORRELATED NOISES
text: Note that Algorithm 1 requires all the local clients to generate the independent noise ξck. Such a mechanism enjoys the convenience of the implementation and yields a potential to protect the privacy of data and alleviates the security issue. However, the large scale noise inevitable slows down the convergence. To handle this issue, the independent noise can be generalized to correlated noise based on a correlation coefficient ρ between different clients. Replacing Eq. (5) with
βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητρ2ξ̇k + √ 2η(1− ρ2)τ/pcξck, (11)
where ξ̇k is a d-dimensional standard Gaussian vector shared by all the clients at iteration k and ξ̇k is dependent with ξck for any c ∈ [N ]. Following the synchronization step based on Eq. (6), we have
θk+1 = θk − η∇f̃(θk) + √ 2ητξk, (12)
where ξk = ρξ̇k + √ 1− ρ2 ∑N c=1 √ pcξ c k. Since the variance of i.i.d variables is additive, it is clear that ξk follows the standard d-dimensional Gaussian distribution. The inclusion of the correlated noise implicitly reduces the temperature for each client and naturally yields a trade-off between federation and accuracy. We refer to the algorithm with correlated noise as the hybrid federated averaging Langevin dynamics (hFA-LD) and present it in Algorithm 2.
Since the inclusion of correlated noise doesn’t affect the iterate of Eq. (12), the algorithm property maintains the same except the scale of the temperature τ and efficacy of federation are changed. Based on a target correlation coefficient ρ ≥ 0, Eq. (11) is equivalent to applying a temperature Tc,ρ = τ(ρ
2 +(1−ρ2)/pc). In particular, setting ρ = 0 leads to Tc,0 = τ/pc, which exactly recovers Algorithm 1; however, setting ρ = 1 leads to Tc,1 = τ , where the injected noise in local clients is reduced by 1/pc times. Now we adjust the analysis as follows Theorem 5.9 (Informal version of Theorem B.8). Assume assumptions 5.1, 5.2, and 5.3 hold. Consider Algorithm 2 with a correlation coefficient ρ ∈ [0, 1], η ∈ (0, 12L ] and ‖θ c 0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have
W2(µk, π) ≤ (1− ηm/4)k · (√ 2d ( D + √ τ/m )) + 30κ √ ηmd · √ (K2 + κ)Hρ,
where µk denotes the probability measure of θk, Hρ := D2 + 1m maxc∈[N ] Tc,ρ + γ2 m2d + σ2 m2 .
Such a mechanism leads to a trade-off between the efficacy of federation and accuracy and motivates us to exploit the optimal ρ under the differential-privacy theories Wang et al. (2015).

heading: 5.4 PARTIAL DEVICE PARTICIPATION
text: Full device participation enjoys appealing convergence properties. However, it suffers from the straggler’s effect in real-world applications, where the communication is limited by the slowest device. Partial device participation handles this issue by only allowing a small portion of devices in each communication and greatly increased the communication efficiency in a federated network.
The first device-sampling scheme I Li et al. (2020b) selects a total of S devices, where the c-th device is selected with a probability pc. The first theoretical justification for convex optimization has been proposed by Li et al. (2020c).
(Scheme I: with replacement). Assume Sk = {n1, n2, · · · , nS}, where nj ∈ [N ] is a random number that takes a value of c with a probability pc for any j ∈ {1, 2, · · · , S}. The synchronization step follows that θk = 1S ∑ c∈Sk θ c k.
Another strategy is to uniformly select S devices without replacement. We follow Li et al. (2020c) and assume S indices are selected uniformly without replacement. In addition, the convergence also requires an additional assumption on balanced data Li et al. (2020c).
(Scheme II: without replacement). Assume Sk = {n1, n2, · · · , nS}, where nj ∈ [N ] is a random number that takes a value of c with a probability 1S for any j ∈ {1, 2, · · · , S}. Assume the data is balanced such that p1 = · · · = pN = 1N . The synchronization step follows that θk = N S ∑ c∈Sk pcθ c k = 1 S ∑ c∈Sk θ c k.
Algorithm 3 Hybrid federated Averaging Langevin dynamics Algorithm (FA-LD) with partial device participation, informal version of Algorithm 6. Sk is sampled according to a device-sampling rule based on scheme I or II.
1: βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητρ2ξ̇k + √ 2η(1− ρ2)τ/pcξck,
2:
θck+1 =  βck+1 if k + 1 mod K 6= 0∑ c∈Sk+1 1 Sβ c k+1 if k + 1 mod K = 0.
Theorem 5.10 (Informal version of Theorem C.3). Assume assumptions 5.1, 5.2, and 5.3 hold. Consider Algorithm 3 with a hyperparameter ρ ∈ [0, 1], a fixed learning rate η ∈ (0, 12L ] and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have
W2(µk, π) ≤ (1− ηm/4)k · (√ 2d ( D + √ τ/m )) + 30κ √ ηmd · √ Hρ(K2 + κ) +O (√ d
S (ρ2 +N(1− ρ2))CS
) ,
where CS = 1 for Scheme I and CS = N−SN−1 for Scheme II.
We observe that partial device participation leads to an extra bias regardless of the scale of η. To reduce such a bias, we suggest to consider highly correlated injected noise, such as ρ = 1, to reduce the impact of the injected noise. By setting O( √ d/S) ≤ /3 and following a similar learning rate as in section 5.3.1, we can achieve the precision within Ω( −2d log(d/ )) iterations given a large number of devices satisfying S = Ω( −2d).
The device-sampling scheme I provides a viable solution to handle the straggler’s effect in full device participation and greatly accelerates the communication efficiency. In addition, scheme I is rather robust to the data heterogeneity and doesn’t require the data to be balanced. In other words, this device-sampling scheme is more preferred if a system is free to activate any devices at any time.
In more practical cases where a system can only operate based on the first S messages for the local updates. The device-sampling scheme II proposes a concrete treatment to tackle this issue. Given a balanced data across different clients and each device is uniformly sampled, we can achieve a reasonable approximation. If S = 1, our Scheme II matches the result in the Scheme I. If S = N , then our Scheme II recovers the result in the full device setting. If S = N − o(N), then our Scheme II bound is better than scheme I.

heading: 6 CONCLUSION AND FUTURE WORK
text: We propose a novel convergence analysis for federated averaging Langevin dynamics (FA-LD) with distributed clients. Our results no longer require the bounded gradient assumption in `2 norm as in the optimization-driven literature in federated learning. The theoretical guarantees yield a concrete guidance on the selection of the optimal number of local updates. In addition, the convergence highly depends on the data heterogeneity and the injected noises, where the latter also inspires us to consider correlated injected noise to balance between the efficacy of federation and accuracy.
Our work initiated the theoretical study of standard sampling algorithms in federated learning and paved the way for future works of advanced Monte Carlo methods, such as underdamped Langevin dynamics Cheng et al. (2018), replica exchange Monte Carlo (also known as parallel tempering) Deng et al. (2020) in federated learning. It is also interesting to study the optimal number of local steps under the non-strongly convex Dalalyan (2017b) or non-convex assumptions Raginsky et al. (2017); Ma et al. (2019).

heading: A PRELIMINARIES

heading: A.1 BASIC NOTATIONS AND BACKGROUNDS
text: Let N denote the number of clients. Let T denote the number of global steps to achieve the precision . Let K denote the number of local steps. For each c ∈ [N ] := {1, 2, · · · , N}, we use f c and ∇f c denote the loss function and gradient of the function f c in client c. For the stochastic gradient oracle, we denote by ∇f̃ c(·) the unbiased estimate of the exact gradient ∇f c of client c. In addition, we denote pc as the weight of the c-th client such that pc ≥ 0 and ∑N c=1 pc = 1. ξ c k is an independent standard d-dimensional Gaussian vector at iteration k for each client c ∈ [N ] and ξ̇k is a unique Gaussian vector shared by all the clients.
Algorithm 4 Federated averaging Langevin dynamics algorithm (FA-LD). Denote by θck the model parameter in the c-th client at the k-th step. Denote the immediate result of one step SGLD update from θck by β c k. ξ c k is an independent standard d-dimensional Gaussian vector at iteration k for each client c ∈ [N ]. A global synchronization is conducted every K steps. This is a complete version of Algorithm 1.
1: βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητ/pcξ c k, (13)
2:
θck+1 =  βck+1 if k + 1 mod K 6= 0∑N c=1 pcβ c k+1 if k + 1 mod K = 0.
(14)
Inspired by Li et al. (2020c), we define two virtual sequences
βk = N∑ c=1 pcβ c k, θk = N∑ c=1 pcθ c k, (15)
which are both inaccessible when k mod K 6= 0. For the gradients and injected noise, we also define
∇f(θk) = N∑ c=1 pc∇f c(θck), ∇f̃(θk) = N∑ c=1 pc∇f̃ c(θck), ξk = N∑ c=1 √ pcξ c k. (16)
In what follows, it is clear that E∇f̃(θ) = ∑N c=1 pcE∇f̃ c(θc) = ∇f(θ) for any θc ∈ Rd and any c ∈ [N ]. Summing Eq.(13) from clients c = 1 to N and combining Eq.(15) and Eq.(16), we have
βk+1 = θk − η∇f̃(θk) + √ 2ητξk. (17)
Moreover, we always have βk = θk whether k + 1 mod E = 0 or not by Eq.(14) and Eq.(15). In what follows, we can write θk+1 = θk − η∇f̃(θk) + √
2ητξk, (18) which resembles the SGLD algorithm Welling & Teh (2011) except that the construction of stochastic gradients is different and θk is not accessible when k mod K 6= 0. To facilitate the analysis, we also define an auxiliary continuous-time processes (θ̄t)t≥0
dθ̄t = −∇f(θ̄t) · dt+ √ 2τ · dW t, (19)
where θ̄t = ∑N c=1 pcθ̄ c t , ∇f(θ̄t) = ∑N c=1 pc∇f c(θ̄ct ), θ̄ct is the continuous-time variable at client c, and W is a d-dimensional Brownian motion. The continuous-time algorithm is referred to as Federated Averaging Langevin diffusion and is described as
dβ̄ct = −∇f c(θ̄ct ) · dt+ √ 2τ/pc · dW c t
θ̄ct = N∑ c=1 pcβ̄ c t .
Since the synchronization step is conducted at every time step t, the Federated Averaging Langevin diffusion performs the same as the standard Langevin diffusion with the temperature τ and convergences to the stationary distribution π(θ) ∝ exp(−f(θ)/τ), where f(θ) = ∑N c=1 pcf
c(θ). Assume that θ̄0 simulates from the stationary distribution π, then it follows that θ̄t ∼ π for any t ≥ 0.

heading: A.2 ASSUMPTIONS AND DEFINITIONS
text: Assumption A.1 (Smoothness). For each c ∈ [N ], we say f c is L-smooth if for some L > 0
f c(y) ≤ f c(x) + 〈∇f c(x), y − x〉+ L 2 ‖y − x‖22 ∀x, y ∈ Rd.
Note that the above assumption is equivalent to saying that
‖∇f c(y)−∇f c(x)‖2 ≤ L‖y − x‖2, ∀x, y ∈ Rd.
Assumption A.2 (Strong convexity). For each c ∈ [N ], f c is m-strongly convex if for some m > 0
f c(x) ≥ f c(y) + 〈∇f c(y), x− y〉+ m 2 ‖y − x‖22 ∀x, y ∈ Rd.
An alternative formulation for strong convexity is that
〈∇f c(x)−∇f c(y), x− y〉 ≥ m ‖x− y‖22 ∀x, y ∈ R d.
Assumption A.3 (Bounded variance, restatement of Assumption 5.3). For each c ∈ [N ], the variance of noise in the stochastic gradient∇f̃ c(x) in each client is upper bounded such that
E[‖∇f̃ c(x)−∇f c(x)‖22] ≤ σ2d, ∀x ∈ Rd.
The bounded variance in the stochastic gradient is a rather standard assumption and has been widely used in Cheng et al. (2018); Dalalyan & Karagulyan (2019); Li et al. (2020c). Extension of bounded variance to unbounded cases such as E[‖∇f̃ c(x) −∇f c(x)‖22] ≤ δ(L2x2 + B2) for some M and δ ∈ [0, 1) is quite straightforward and has been adopted in assumption A.4 stated in Raginsky et al. (2017). The proof framework remains the same.
Quality of non-i.i.d data Denote by θ∗ the global minimum of f . Next, we quantify the degree of the non-i.i.d data by γ := maxc∈[N ] ‖∇f c(θ∗)‖2, which is a non-negative constant and yields a smaller scale if the data is more evenly distributed.
Definition A.4. We define parameter Tc,ρ H2ρ , κ and γ2
Tc,ρ := τ(ρ 2 + (1− ρ2)/pc),
Hρ := D2︸︷︷︸ initialization + 1 m max c∈[N ]
Tc,ρ︸ ︷︷ ︸ injected noise
+ γ2
m2d︸ ︷︷ ︸ data heterogeneity
+ σ2
m2︸︷︷︸ stochastic noise ,
κ := L/m,
γ2 := max c∈[N ] ‖∇f c(θ∗)‖22.

heading: B FULL DEVICE PARTICIPATION

heading: B.1 ONE-STEP UPDATE
text: Wasserstein distance We define the 2-Wasserstein distance between a pair of Borel probability measures µ and ν on Rd as follows
W2(µ, ν) := inf γ2∈Couplings(µ,ν)
(∫ ‖βµ − βν‖22dγ2(βµ,βν) ) 1 2 ,
where ‖ · ‖2 denotes the `2 norm on Rd and the pair of random variables (βµ,βν) ∈ Rd × Rd is a coupling with the marginals following L(βµ) = µ and L(βν) = ν. L(·) denotes a distribution of a random variable.
The following result provides a crucial contraction property based on distributed clients with infrequent synchronizations. Lemma B.1 (Dominated contraction property, restatement of Lemma 5.4). Assume assumptions A.1 and A.2 hold. For any learning rate η ∈ (0, 1L+m ], any {θ c}Nc=1, {βc}Nc=1 ∈ Rd, we have
‖β − θ − η(∇f(β)−∇f(θ))‖22 ≤ (1− ηm) · ‖β − θ‖ 2 2 + 4ηL N∑ c=1 pc · (‖βc − β‖22 + ‖θc − θ‖22).
where β = ∑N c=1 pcβ c, θ = ∑N c=1 pcθ c, ∇f(θ) = ∑N c=1 pc∇f c(θc), and ∇f(β) =∑N
c=1 pc∇f c(βc). We postpone the proof into Section D.1. The above result implies that as long as the local parameters θc, βc and global θ, β don’t differ each other too much, we can guarantee the desired convergence.
The following result ensures a bounded gap between θ̄cs and θ̄ c ηb sη c in `2 norm for any s ≥ 0 and c ∈ [N ]. We postpone the proof of Lemma B.2 into Section D.2. Lemma B.2 (Discretization error). Assume assumptions A.1, A.2, and A.3 hold. For any s ≥ 0, any learning rate η ∈ (0, 2/m) and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], the iterates of (θ̄s) based on the continuous dynamics of Eq.(19) satisfy the following estimate
E ∥∥θ̄cs − θ̄cηb sη c∥∥22 ≤ 8η2dκ
( κγ2
d + Lτ
) + 16ηdτ.
The following result shows that given a finite number of local steps K, the divergence between θc in local client and θ in the center is bounded in `2 norm. Notably, since the non-differentiable Brownian motion leads to a lower order term O(η) instead of O(η2) in `2 norm, a naïve proof may lead to a crude upper bound. We delay the proof of Lemma B.3 into Section D.3. Lemma B.3 (Bounded divergence, restatement of Lemma 5.5). Assume assumptions A.1, A.2, and A.3 hold. For any learning rate η ∈ (0, 2/m) and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have the `2 upper bound of the divergence between local clients and the center as follows
N∑ c=1 pcE‖θck − θk‖22 ≤ 112(K − 1)2η2dL2Hρ + 8(K − 1)ηdτ(ρ2 +N(1− ρ2)),
where Hρ, κ and γ2 are defined as Definition A.4.
The following presents a standard result for bounding the gap between∇f(θ) and∇f̃(θ). We delay the proof of Lemma B.4 into Setion D. Lemma B.4 (Bounded variance). Given assumption A.3, we have
E‖∇f(θ)−∇f̃(θ)‖22 ≤ d · σ2, ∀ θ ∈ Rd.
Having all the preliminary results ready, now we present a crucial lemma for proving the convergence of all the algorithms.
Lemma B.5 (One step update, restatement of Lemma 5.6). Assume assumptions A.1, A.2, and A.3 hold. Consider Algorithm 4 with independently injected noise ρ = 0, any learning rate η ∈ (0, 12L ) and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], where θ∗ is the global minimum for the function f . Then
W 22 (µk+1, π) ≤ (
1− ηm 2
) ·W 22 (µk, π) + 400η2dL2H0((K − 1)2 + κ),
where µk denotes the probability measure of θk, H0, κ and γ2 are defined as Definition A.4.
Proof. The solution of the continuous-time process Eq.(19) follows that θ̄t = θ̄0 − ∫ t 0 ∇f(θ̄s)ds+ √ 2τ ·W t, ∀t ≥ 0. (20)
Set t → (k + 1)η and θ̄0 → θ̄kη for Eq.(20) and consider a synchronous coupling such that W(k+1)η −Wkη := √ ηξk
θ̄(k+1)η = θ̄kη − ∫ (k+1)η kη ∇f(θ̄s)ds+ √ 2τ(W(k+1)η −Wkη)
= θ̄kη − ∫ (k+1)η kη ∇f(θ̄s)ds+ √ 2τηξk. (21)
We first denote ζk := ∇f̃(θk)−∇f(θk). Subtracting Eq.(18) from Eq.(21) yields that
θ̄(k+1)η − θk+1 = θ̄kη − θk + η∇f̃(θk)− ∫ (k+1)η kη ∇f(θ̄s)ds
= θ̄kη − θk − η ( ∇f(θk + θ̄kη − θk)−∇f̃(θk) ) − ∫ (k+1)η kη ( ∇f(θ̄s)−∇f(θ̄kη) ) ds (22)
= θ̄kη − θk − η ( ∇f(θk + θ̄kη − θk)−∇f(θk)︸ ︷︷ ︸
:=Xk
) − ∫ (k+1)η kη ( ∇f(θ̄s)−∇f(θ̄kη) ) ds︸ ︷︷ ︸
:=Yk
+ηζk.
Taking square and expectation on both sides, we have
E‖θ̄(k+1)η − θk+1‖22 = E‖θ̄kη − θk − ηXk − Yk‖22 + E‖ηζk‖22 + 2η E〈θ̄kη − θk − ηXk − Yk, ζk〉︸ ︷︷ ︸
Eζk=0 and mutual independence
≤ (1 + q) · E‖θ̄kη − θk − ηXk‖22 + (1 + 1/q) · E‖Yk‖22 + E‖ηζk‖22 ≤ (1 + q) · ( (1− ηm) · E‖θ̄kη − θk‖22 + 4ηL N∑ c=1 pc · ( E‖θ̄ckη − θ̄kη‖22 + E‖θck − θk‖22 ) ) + (1 + 1/q) · E‖Yk‖22 + η2σ2d
≤ (1 + q) · (
(1− ηm)︸ ︷︷ ︸ φ
E‖θ̄kη − θk‖22 + 448η3d(K − 1)2L3H0 + 32(K − 1)η2dLτN )
+ (1 + 1/q) · E‖Yk‖22 + η2σ2d, (23)
where the first inequality follows by the AM-GM inequality for any q > 0, the second inequality follows by Lemma B.1 and Assumption A.3. The third inequality follows by Lemma B.3 with ρ = 0; moreover, the continuous-time process conducts synchronization at any time step, hence θ̄ckη = θ̄kη . Since the learning rate follows 12L ≤ 1 m+L ≤ 2 m , the requirement of the learning rate for Lemma B.1 and Lemma B.3 is clearly satisfied.
Recall that φ = 1− ηm, we get 1+φ2 = 1− 1 2ηm. Choose q = 1+φ 2φ − 1 so that (1 + q)φ = (1+φ) 2 = 1− 12ηm. In addition, we have 1 + 1 q = 1+q q = 1+φ 1−φ ≤ 2 ηm . It follows that
(1 + q) · (1− ηm) ≤ 1− 1 2 ηm, 1 + q ≤ 1− 12ηm 1− ηm ≤ 1.5, (1 + 1/q) ≤ 2 mη , (24)
where the second inequality holds because η ∈ (0, 12L ] ≤ 1 2m .
For the term E‖Yk‖22 in Eq.(23), we have the following estimate
E‖Yk‖22 = E ∥∥∥∥∥ ∫ (k+1)η kη ( ∇f(θ̄s)−∇f(θ̄kη) ) ds ∥∥∥∥∥ 2
2 ≤ η ∫ (k+1)η kη E ∥∥∇f(θ̄s)−∇f(θ̄kη)∥∥22ds
= η ∫ (k+1)η kη E ∥∥∥∥∥ N∑ c=1 pc ( ∇f c(θ̄cs)−∇f c(θ̄ckη) )∥∥∥∥∥ 2
2
ds
≤ η ∫ (k+1)η kη N∑ c=1 pc · E ∥∥∇f c(θ̄cs)−∇f c(θ̄ckη)∥∥22ds
≤ ηL2 ∫ (k+1)η kη N∑ c=1 pc · E ∥∥θ̄cs − θ̄ckη∥∥22ds
≤ ηL2 ∫ (k+1)η kη ( 8η2dκ ( κγ2 d + Lτ ) + 16ηdτ ) ds
= 8η4dL4H0 + 16η 3L2dτ, (25)
where the first inequality follows by Hölder’s inequality, the second inequality follows by Jensen’s inequality, the third inequality follows by Assumption A.1, and the last inequality follows by Lemma B.2. The last equality holds since κdγ
2 + Lτ ≤ LmH0 and κ = L/m. Plugging Eq.(24) and Eq.(25) into Eq.(23), we have
E‖θ̄(k+1)η − θk+1‖22 ≤ (1− ηm
2 ) · E‖θ̄kη − θk‖22
+ 672η3d(K − 1)2L3H0 + 48η2d(K − 1)LτN
+ 16η3dL3κH0 + 32η 2d L2
m τ + η2σ2d.
Choose the specific Langevin diffusion θ̄ in stationary regime, we have W 22 (µk, π) = E‖θ̄kη − θk‖22 and W 22 (µk+1, π) ≤ E‖θ̄(k+1)η − θk+1‖22. Arranging the terms, we have
W 22 (µk+1, π) ≤ (1− ηm
2 ) ·W 22 (µk, π) + 400η2dL2H0((K − 1)2 + κ),
where η ≤ 12L , κ ≥ 1, mτ ≤ Lτ ≤ LτN ≤ Lmaxc∈[N ] Tc,0 ≤ LmH0, and σ 2 ≤ L2H0 are applied to the result.

heading: B.2 CONVERGENCE VIA INDEPENDENT NOISES
text: Theorem B.6 (Restatement of Theorem 5.7). Assume assumptions A.1, A.2, and A.3 hold. Consider Algorithm 4 with a fixed learning rate η ∈ (0, 12L ] and ‖θ c 0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have
W2(µk, π) ≤ (
1− ηm 4
)k · (√ 2d ( D + √ τ/m )) + 30κ √ ηmd · √ ((K − 1)2 + κ)H0.
where µk denotes the probability measure of θk, H0, κ and γ2 are defined as Definition A.4.
Proof. Iteratively applying Theorem B.5 and arranging terms, we have that
W 22 (µk, π) ≤ (
1− ηm 2
)k W 22 (µ0, π) + 1− (1− ηm2 ) k
1− (1− ηm2 )
( 400η2dL2H0((K − 1)2 + κ) ) ≤ (
1− ηm 2
)k W 22 (µ0, π) + 2
ηm
( 400η2dL2H0((K − 1)2 + κ) ) ≤ (
1− ηm 2
)k W 22 (µ0, π) + 800κ 2ηmd((K − 1)2 + κ)H0, (26)
where κ = Lm . By Lemma F.1 and the initialization condition ‖θ c 0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have that
W2(µ0, π) ≤ √ 2d(D + √ τ/m).
Applying the inequality (1− ηm2 ) ≤ (1− ηm 4 ) 2 completes the proof.

heading: Discussions
text: Optimal choice of K. To ensure the algorithm to achieve the precision based on the total number of steps T and the learning rate η, we can set
30κ √ ηmd · (√ ((K − 1)2 + κ)H0 ) ≤
2
e− ηm 4 T · (√ 2d ( D + √ τ/m )) ≤
2 .
This directly leads to ηm ≤ min { m
2L ,O
( 2
dκ2((K − 1)2 + κ)H0
)} , T ≥ Ω ( log ( d ) mη ) .
Plugging into the upper bound of η, it implies that to reach the precision level , it suffices to set
T = Ω
( dκ2((K − 1)2 + κ)H0
2 · log
( d )) . (27)
Since H0 = Ω(D2 + τm ), we observe that the number of communication rounds is around the order
T K = Ω
( K + κ
K
) ,
where the value of T K first decreases and then increases with respect to K, indicating that setting K either too large or too small may lead to high communication costs and hurt the performance. Ideally, K should be selected in the scale of Ω( √ κ). Combining the definition of T in Eq.(27), this suggests an interesting result that the optimal K should be in the order of O( √ T ). Similar results have been achieved by Stich (2019); Li et al. (2020c).

heading: B.3 CONVERGENCE VIA VARYING LEARNING RATES
text: Theorem B.7 (Restatement of Theorem 5.8). Assume assumptions A.1, A.2, and A.3 hold. Consider Algorithm 4 with an initialization satisfying ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ] and varying learning rate following
ηk = 1
2L+ (1/12)mk , k = 1, 2, · · · .
Then for any k ≥ 0, we have W2(µk, π) ≤ 45κ √ ((K − 1)2 + κ)H0 · ( ηkmd )1/2 , ∀k ≥ 0,
Proof. We first denote
Cκ = 30κ √
((K − 1)2 + κ)H0. Next, we proceed to show the following inequality by the induction method
W2(µk, π) ≤ 1.5Cκ (
d
2L+ (1/12)mk
)1/2 = 1.5Cκ ( ηkmd )1/2 , ∀k ≥ 0, (28)
where the decreasing learning rate follows that
ηk = 1
2L+ (1/12)mk .
(i) For the case of k = 0, since
Cκ ≥ 4 √ κ √ H0 ≥ 4 √ κ √ D2 + 1
m max c∈[N ]
Tc,0 ≥ 4 √ κ/d (√ dD2 + √ d
m max c∈[N ] Tc,0 ) ≥ 4 √ κ/dW2(µ0, π), (29)
where the last inequality follows by Lemma F.1 and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ]. It is clear that W2(µ0, π) ≤ 14Cκ √ md L ≤ 1.5Cκ √ η0md by Eq.(29).
(ii) If now that Eq.(28) holds for some k ≥ 0, it follows by Lemma B.5 that W 22 (µk+1, π) ≤ ( 1− ηkm
2
) ·W 22 (µk, π) + 400η2kdL2H0((K − 1)2 + κ)
≤ ( 1− ηkm
2
) ·W 22 (µk, π) + η2km 2
2 C2κd ≤ ( 1− ηkm
2
) · 2.25C2κηkmd+ ηkm
3 2.25C2κηkmd ≤ ( 1− ηkm
6
) · 2.25C2κηkmd.
Since ( 1− ηkm6 ) ≤ ( 1− ηkm12 )2 , we have
W2(µk+1, π) ≤ ( 1− ηkm
12
) · 1.5Cκ ( ηkmd )1/2 .
To prove W2(µk+1, π) ≤ 1.5Cκ ( ηk+1md )1/2 , it suffices to show ( 1− ηkm12 ) η 1/2 k ≤ ηk+1, which is detailed as follows ( 1− ηkm
12
) η 1/2 k = √ 12(24L+mk −m) (24L+mk)3/2
≤ √ 12(24L+mk −m)1/2
24L+mk
≤ √ 12
(24L+m(k + 1))1/2 := ηk+1,
where the last inequality follows since (24L+mk −m)(24L+mk +m)) ≤ (24L+mk)2.
The above result implies that to achieve the precision , we require W2(µk, π) ≤ 1.5Cκ ( md
2L+ (1/12)mk
)1/2 ≤ .
The means that we only require k = Ω( d 2 ) to achieve the precision . By contrast, the fixed learning rate requires T = Ω ( d 2 · log ( d/ )) , which is much slower than the algorithm with varying learning
rate by O ( log ( d/ )) times.

heading: B.4 PRIVACY-ACCURACY TRADE-OFF VIA CORRELATED NOISES
text: Note that Algorithm 4 requires all the local clients to generate the independent noise ξck. Such a mechanism enjoys the convenience of the implementation and yields a potential to protect the privacy of data and alleviates the security issue. However, the scale of noises is maximized and inevitable slows down the convergence. For extensions, it can be naturally generalized to correlated noise based on a hyperparameter, namely the correlation coefficient ρ between different clients. Replacing Eq.(13) with
βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητρ2ξ̇k + √ 2η(1− ρ2)τ/pcξck, (30)
where ξ̇k is a d-dimensional standard Gaussian vector shared by all the clients at iteration k, ξck is a unique d-dimensional Gaussian vector generated by client c ∈ [N ] only. Moreover, ξ̇k is dependent with ξck for any c ∈ [N ]. Following the same synchronization step based Eq.(14), we have
θk+1 = θk − η∇f̃(θk) + √ 2ητξk, (31)
where ξk = ρξk + √ 1− ρ2 ∑N c=1 √ pcξ c k. Since the variance of i.i.d variables is additive, it is clear that ξk follows the standard d-dimensional Gaussian distribution. The inclusion of the correlated noise implicitly reduces the temperature and naturally yields a trade-off between federation and accuracy. We refer to the algorithm with correlated noise as the generalized Federated Averaging Langevin dynamics (gFA-LD) and present it in Algorithm 5.
Since the inclusion of correlated noise doesn’t affect the formulation of Eq.(31), the algorithm property maintains the same except the scale of the temperature τ and federation are changed. Based on a target correlation coefficient ρ ≥ 0, Eq.(30) is equivalent to applying a temperature Tc,ρ = τ(ρ
2 + (1− ρ2)/pc). In particular, setting ρ = 0 leads to Tc,0 = (1− ρ2)/pc, which exactly recovers Algorithm 4; however, setting ρ = 1 leads to Tc,1 = τ , where the injected noise in local clients is reduced by 1/pc times. Now we adjust the analysis as follows
Theorem B.8 (Restatement of Theorem 5.9). Assume assumptions A.1, A.2, and A.3 hold. Consider Algorithm 5 with a correlation coefficient ρ ∈ [0, 1], a fixed learning rate η ∈ (0, 12L ] and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have
W2(µk, π) ≤ (
1− ηm 4
)k · (√ 2d ( D + √ τ/m )) + 30κ √ ηmd · √ ((K − 1)2 + κ)Hρ,
where µk denotes the probability measure of θk, Hρ, κ and γ2 are defined as Definition A.4.
Algorithm 5 Hybrid federated averaging Langevin dynamics algorithm (hFA-LD). Denote by θck the model parameter in the c-th client at the k-th step. Denote the immediate result of one step SGLD update from θck by β c k. ξ c k is an independent standard d-dimensional Gaussian vector at iteration k for each client c ∈ [N ] and ξ̇k is a d-dimensional standard Gaussian vector shared by all the clients. ρ denotes the correlation coefficient of the injected noises. A global synchronization is conducted every K steps. This is a complete version of Algorithm 2.
1: βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητρ2ξ̇k + √ 2η(1− ρ2)τ/pcξck,
2:
θck+1 =  βck+1 if k + 1 mod K 6= 0∑N c=1 pcβ c k+1 if k + 1 mod K = 0.

heading: C PARTIAL DEVICE PARTICIPATION
text: Full device participation enjoys appealing convergence properties. However, it suffers from the straggler’s effect in real-world applications, where the communication is limited by the slowest device. Partial device participation handles this issue by only allowing a small portion of devices in each communication and greatly increased the communication efficiency in a federated network.

heading: C.1 UNBIASED SAMPLING SCHEMES
text: The first device-sampling scheme I Li et al. (2020b) selects a total of S devices, where the c-th device is selected with a probability pc. The first theoretical justification for convex optimization has been proposed by Li et al. (2020c).
(Scheme I: with replacement). Assume Sk = {n1, n2, · · · , nS}, where nj ∈ [N ] is a random number that takes a value of c with a probability pc for any j ∈ {1, 2, · · · , S}. The synchronization step follows that θk = 1S ∑ c∈Sk θ c k.
Another strategy is to uniformly select S devices without replacement. We follow Li et al. (2020c) and assume S indices are selected uniformly without replacement and the synchronization step is the same as before. In addition, the convergence also requires an additional assumption on balanced data Li et al. (2020c).
(Scheme II: without replacement). Assume Sk = {n1, n2, · · · , nS}, where nj ∈ [N ] is a random number that takes a value of c with a probability 1S for any j ∈ {1, 2, · · · , S}. Assume the data is balanced such that p1 = · · · = pN = 1N . The synchronization step follows that θk = N S ∑ c∈Sk pcθ c k = 1 S ∑ c∈Sk θ c k.
Algorithm 6 Hybrid federated averaging Langevin dynamics algorithm (hFA-LD) with partial device participation. ξck is the independent Gaussian vector proposed by each client c ∈ [N ] and ξ̇k is a unique Gaussian vector shared by all the clients. ρ denotes the correlation coefficient. A global synchronization is conducted every K steps. Sk is a subset that contains S indices according to a device-sampling rule based on scheme I or II. This is a complete version of Algorithm 3.
1: βck+1 = θ c k − η∇f̃ c(θck) + √ 2ητρ2ξ̇k + √ 2η(1− ρ2)τ/pcξck,
2:
θck+1 =  βck+1 if k + 1 mod K 6= 0∑ c∈Sk+1 1 Sβ c k+1 if k + 1 mod K = 0.
Lemma C.1 (Unbiased sampling scheme). For any k mod K = 0 based on scheme I or II, we have
Eθk = E ∑ c∈Sk θck = βk := N∑ c=1 pcβ c k.
Proof. According to the definition of scheme I or II, we have θk = 1S ∑ c∈Sk θ c k. In what follows,
Eθk = 1SE ∑ c∈Sk θ c k = 1 S ∑ c0∈Sk ∑N c=1 pcβ c k = ∑N c=1 pcβ c k, where p1 = p2 = · · · = pN for scheme II in particular.

heading: C.2 BOUNDED DIVERGENCE BASED ON PARTIAL DEVICE
text: Lemma C.2 (Bounded divergence based on partial device). Assume assumptions A.1, A.2, and A.3 hold. Consider Algorithm 6 with a correlation coefficient ρ ∈ [0, 1], any learning rate η ∈ (0, 2/m) and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have the following results
For Scheme I, the divergence between θk and βk is upper bounded by
E‖βk − θk‖22 ≤ 112
S K2η2dL2Hρ +
8 S Kηdτ(ρ2 +N(1− ρ2)).
For Scheme II, assuming the data is balanced such that p1 = · · · = pN = 1N , the divergence between θk and βk is upper bounded by
E‖βk − θk‖22 ≤ N − S S(N − 1)
( 112K2η2dL2Hρ + 8Kηdτ(ρ 2 +N(1− ρ2)) ) .
where Hρ, κ and γ2 are defined as Definition A.4.
Proof. We prove the bounded divergence for the two schemes, respectively. For scheme I with replacement, θ̄k = ∑ c∈Sk 1 Sβ c k for a subset of indices Sk. Taking expectation with respect to Sk, we have
E‖θk − βk‖22 = 1
S2 S∑ i=1 E‖βnik − βk‖ 2 2 = 1 S N∑ c=1 pc ‖βck − βk‖ 2 2 , (32)
where the first equality follows by the independence and unbiasedness of θnik for any i ∈ [S]. To further upper bound Eq.(32), we follow the same technique as in Lemma B.3. Since k mod K = 0, k0 = k −K is also the communication time, which yields the same θck0 for any c ∈ [N ]. in what follows,
N∑ c=1 pc ‖βck − βk‖ 2 2 = N∑ c=1 pc ‖βck − θk0 − (βk − θk0)‖ 2 2
≤ N∑ c=1 pc ‖βck − θk0‖ 2 2 , (33)
where the last inequality follows by βk = ∑N c=1 pcβ c k and E‖x− Ex‖ 2 2 ≤ E‖x‖ 2 2. Combining Eq.(32) and Eq.(33), we have
E‖θk − βk‖22 ≤ 1
S N∑ c=1 pc ‖βck − θk0‖ 2 2
≤ 1 S N∑ c=1 pc ∥∥βck − θck0∥∥22
≤ 1 S N∑ c=1 pcE k−1∑ k=k0 2Kη2 ∥∥∥∇f̃ c(θck)∥∥∥2 2 + 4Kηdτ ( ρ2 + (1− ρ2)/pc ) ≤ 1 S N∑ c=1 pc ( k−1∑ k=k0 2Kη2E ∥∥∥∇f̃ c(θck)∥∥∥2 2 + 4Kηdτ ( ρ2 + (1− ρ2)/pc ))
≤ 28 S K2η2dL2Hρ + 4 S Kηdτ(ρ2 +N(1− ρ2))
where the last inequality follows a similar argument as in Lemma B.3.
For scheme II, given p1 = p2 = · · · = pN = 1N , we have θk = 1 S ∑ c∈Sk β c k, which leads to
E‖θk − βk‖22 = E ∥∥∥∥∥ 1S ∑ c∈Sk βck − βk ∥∥∥∥∥ 2
2
= 1
S2 E ∥∥∥∥∥ N∑ c=1 Ic∈Sk(βck − βk) ∥∥∥∥∥ 2
2
,
where IA is an indicator function that equals to 1 if the event A happens.
Plugging the facts that P(c ∈ Sk) = SN and P(c1, c2 ∈ Sk) = S(S−1) N(N−1) for any c1 6= c2 ∈ [N ] into the above equation, we have
E‖θk − βk‖22
= 1
S2 [ ∑ c∈[N ] P(c ∈ Sk) ‖βck − βk‖ 2 2 + ∑ c1 6=c2 P(c1, c2 ∈ Sk)〈βc1k − βk, β c2 k − βk〉 ]
= 1
SN N∑ c=1 ‖βck − βk‖ 2 2 + ∑ c1 6=c2 S − 1 SN(N − 1) 〈βc1k − βk, β c2 k − βk〉
= 1− SN
S(N − 1) N∑ c=1 ‖βck − βk‖ 2 2 ,
where the last equality holds since ∑ c∈[N ] ‖βck − βk‖ 2 2 + ∑ c1 6=c2〈β c1 k − βk, β c2 k − βk〉 = ‖βk − βk‖22 = 0. Eventually, we have
E‖θk − βk‖22 = N − S S(N − 1) E 1 N N∑ c=1 ‖βck − βk‖ 2 2
≤ N − S S(N − 1) E 1 N N∑ c=1 ‖βck − θk0‖ 2 2
≤ N − S S(N − 1)
( 28K2η2dL2Hρ + 4Kηdτ ( ρ2 +N(1− ρ2) )) ,
where the first inequality follows a similar argument as in Eq.(33) and the last inequality follows by Lemma B.3.

heading: C.3 CONVERGENCE VIA PARTIAL DEVICE PARTICIPATION
text: Theorem C.3 (Restatement of Theorem 5.10). Assume assumptions A.1, A.2, and A.3 hold. Consider Algorithm 6 with a correlation coefficient ρ ∈ [0, 1], a fixed learning rate η ∈ (0, 12L ] and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have
W2(µk, π) ≤ (
1− ηm 4
)k · (√ 2d ( D + √ τ/m )) + 30κ √ ηmd · √ Hρ((K − 1)2 + κ) + 2 √ CKdτ
Sm (ρ2 +N(1− ρ2))CS ,
where CK = ηmK 1−e− ηmK 2 , CS = 1 for Scheme I and CS = N−SN−1 for Scheme II.
Proof. Note that E ∥∥θ̄(k+1)η − θk+1∥∥22
= E ∥∥θ̄(k+1)η − βk+1 + βk+1 − θk+1∥∥22
= E ∥∥θ̄(k+1)η − βk+1∥∥22 + E‖βk+1 − θk+1‖22 + E2〈θ̄(k+1)η − βk+1, βk+1 − θk+1〉
= E ∥∥θ̄(k+1)η − βk+1∥∥22 + E‖βk+1 − θk+1‖22,
where the last equality follows by the unbiasedness of the device-sampling scheme in Lemma C.1.
If k+ 1 mod K 6= 0, we always have βk+1 = θk+1 and E‖βk+1 − θk+1‖22 = 0. Following the same argument as in Lemma B.5, both schemes lead to the one-step iterate as follows
W 22 (µk+1, π) ≤ (1− ηm
2 ) ·W 22 (µk, π) + 400η2dL2Hρ((K − 1)2 + κ). (34)
If k + 1 mod K = 0, combining Lemma C.2 and Lemma B.5, we have
W 22 (µk+1, π) ≤ (1− ηm
2 ) ·W 22 (µk, π) + 450η2dL2Hρ(K2 + κ) +
4Kdητ
S (ρ2 +N(1− ρ2))CS ,
(35)
where CS = 1 for Scheme I and CS = N−SN−1 for Scheme II.
Repeatedly applying Eq.(34) and Eq.(35) and arranging terms, we have that W 22 (µk, π) ≤ (
1− ηm 2
)k W 22 (µ0, π) + 2
ηm
( 450η2dL2Hρ(K 2 + κ) )
+ (1− (1− ηm2 ) K)bk/Kc
1− (1− ηm2 )K
( 4Kdητ
S (ρ2 +N(1− ρ2))CS ) ≤ (
1− ηm 2
)k W 22 (µ0, π) + 900ηmdκ 2H0((K − 1)2 + κ)
+ ηmK
1− e− ηmK2︸ ︷︷ ︸ CK
4Kdητ ηmKS (ρ2 +N(1− ρ2))CS ,
= (
1− ηm 2
)k W 22 (µ0, π) + 900ηmdκ 2H0((K − 1)2 + κ)
+ 4CKdτ
Sm (ρ2 +N(1− ρ2))CS ,
where the second inequality follows by (1− r)K ≤ e−rK for any r ≥ 0.

heading: D BOUNDING CONTRACTION, DISCRETIZATION, AND DIVERGENCE

heading: D.1 DOMINATED CONTRACTION PROPERTY
text: Proof of Lemma B.1 . Given a client index c ∈ [N ], applying Theorem 2.1.12 Nesterov (2004) leads to
〈y − x,∇f c(y)−∇f c(x)〉 ≥ mL L+m ‖y − x‖22 + 1 L+m ‖∇f c(y)−∇f c(x)‖22 , ∀x, y ∈ R d.
(36)
In what follows, we have
‖β − θ − η(∇f(β)−∇f(θ))‖22 = ‖β − θ‖22 − 2η 〈β − θ,∇f(β)−∇f(θ)〉︸ ︷︷ ︸
I
+η2 ‖∇f(β)−∇f(θ)‖22 . (37)
For the second item I in the right hand side, we have I = N∑ c=1 pc 〈 β − θ,∇f c(βc)−∇f c(θc) 〉 =
N∑ c=1 pc 〈 β − βc + βc − θc + θc − θ,∇f c(βc)−∇f c(θc) 〉 = −
N∑ c=1 pc (〈 βc − β,∇f c(βc)−∇f c(θc) 〉 + 〈 θ − θc,∇f c(βc)−∇f c(θc) 〉) +
N∑ c=1 pc 〈 βc − θc,∇f c(βc)−∇f c(θc) 〉 ≥ −
N∑ c=1 pc · ( (m+ L) ‖βc − β‖22 + (m+ L) ‖θ c − θ‖22 + 1 2(m+ L) ‖∇f c(βc)−∇f c(θc)‖22 ) +
N∑ c=1 pc · ( mL L+m ‖βc − θc‖22 + 1 L+m ‖∇f c(βc)−∇f c(θc)‖22 ) ≥ −(m+ L)
N∑ c=1 pc ( ‖βc − β‖22 + ‖θ c − θ‖22 ) + mL L+m ‖β − θ‖22
+ 1
2(L+m) ‖∇f(β)−∇f(θ)‖22 , (38)
where the first inequality follows by the AM-GM inequality and Eq.(36), respectively; the last inequality follows by Jensen’s inequality such that
N∑ c=1 pc‖βc − θc‖22 ≥ ∥∥∥∥∥ N∑ c=1 pc(β c − θc) ∥∥∥∥∥ 2
2
= ‖β − θ‖22
N∑ c=1 pc ‖∇f c(βc)−∇f c(θc)‖22 ≥ ∥∥∥∥∥ N∑ c=1 pc ( ∇f c(βc)−∇f c(θc) )∥∥∥∥∥ 2
2
= ‖∇f(β)−∇f(θ)‖22 .
Plugging Eq.(38) into Eq.(37), we have
‖β − θ − η · (∇f(β)−∇f(θ))‖22
≤ ( 1− 2ηmL
m+ L
) · ‖β − θ‖22 + η ( η − 1
m+ L︸ ︷︷ ︸ ≤0 if η≤ 1m+L
) · ‖∇f(β)−∇f(θ)‖22
+ 2η(m+ L) N∑ c=1 pc · (‖βc − β‖22 + ‖θc − θ‖22)
≤ (1− ηm) ‖β − θ‖22 + 4ηL N∑ c=1 pc · ( ‖βc − β‖22 + ‖θc − θ‖22 ) ,
where the last inequality follows by 2Lm+L ≥ 1, m ≤ L, 1 − 2a ≤ (1 − a) 2 for any a, and η ∈ (0, 1m+L ].

heading: D.2 DISCRETIZATION ERROR
text: Proof of Lemma B.2. For any s ∈ [0,∞), there exists a certain k ∈ N+ such that s ∈ [kη, (k+ 1)η). By the continuous dynamics of Eq. (19), we have
θ̄cs = θ̄ c ηb sη c + (s− kη)∇f c(θ̄cηb sη c) + √ 2τ ∫ s kη dW t,
which suggests that
sup s∈[kη,(k+1)η) ∥∥θ̄cs − θ̄cηb sη c∥∥2 ≤ (s− kη)∥∥∇f c(θ̄cηb sη c)∥∥2 + sup s∈[kη,(k+1)η) ∥∥∥∥∫ s kη √ 2τdW t ∥∥∥∥ 2 .
We first square the terms on both sides and take Young’s inequality and expectation
E sup s∈[kη,(k+1)η) ∥∥θ̄cs − θ̄cηb sη c∥∥22 ≤ 2E∥∥(s− kη)∇f c(θ̄cηb sη c)∥∥22 + 2E sup
s∈[kη,(k+1)η) ∥∥∥∥∫ s kη √ 2τdW t ∥∥∥∥2 2 .
Then, by Burkholder-Davis-Gundy inequality (50) and Itô isometry, we have
E sup s∈[kη,(k+1)η) ∥∥θ̄cs − θ̄cηb sη c∥∥22 ≤ 2E∥∥(s− kη)∇f c(θ̄cηb sη c)∥∥22 + 8 d∑ i=1 E ∫ s kη 2τdt
≤ 2η2E ∥∥∇f c(θ̄cηb sη c)∥∥22 + 16ηdτ. (39)
By Young’s inequality and the smoothness assumption A.1, we have
E‖∇f c(θ̄ηb sη c)‖ 2 2 = E‖∇f c(θ̄cηb sη c)−∇f c(θ∗) +∇f c(θ∗)‖22
≤ 2E‖∇f c(θ̄cηb sη c)−∇f c(θ∗)‖22 + 2‖∇f c(θ∗)‖ 2 2
≤ 2L2E‖θ̄cηb sη c − θ∗‖ 2 2 + 2γ 2 ≤ 2L2 ( 1
m
( γ2
m + 2dτ
)) + 2γ2
≤ 4dκ ( κγ2
d + 4Lτ
) , (40)
where the third inequality follows by Lemma E.2, the fourth step holds since κ ≥ 1. Combining Eq. (39) and Eq. (40), we have
E sup s∈[kη,(k+1)η)
∥∥θ̄cs − θ̄cηb sη c∥∥22 ≤ 8η2dκ ( κγ2 d + Lτ ) + 16ηdτ.

heading: D.3 BOUNDED DIVERGENCE
text: Proof of Lemma B.3. For any k ≥ 0, consider k0 = Kb kK c such that k ≤ k0 and θ c k0 = θk0 for any k ≥ 0. It is clear that k − k0 ≤ K − 1 for all k ≥ 0. Consider the non-increasing learning rate such that ηk0 ≤ 2ηk for all k − k0 ≤ K − 1. By the iterate Eq.(18), we have
N∑ c=1 pcE‖θck − θk‖ 2 2
= N∑ c=1 pcE‖θck − θk0 − (θk − θk0)‖ 2 2
≤ N∑ c=1 pcE‖θck − θk0‖ 2 2
≤ N∑ c=1 pcE k−1∑ k=k0 2(K − 1)η2k ∥∥∥∇f̃ c(θck)∥∥∥2 2 + 4(K − 1)ηkdτ(ρ2 + (1− ρ2)/pc)
≤ N∑ c=1 pc ( k−1∑ k=k0 2(K − 1)η2k0E ∥∥∥∇f̃ c(θck)∥∥∥2 2 + 4(K − 1)ηk0dτ(ρ2 + (1− ρ2)/pc) ) ≤ 112(K − 1)2η2kdL2Hρ + 8(K − 1)ηkdτ(ρ2 +N(1− ρ2)),
where the first inequality holds by E‖θ − Eθ‖22 ≤ E‖θ‖22 for a stochastic variable θ; the second inequality follows by ( ∑K−1 i=1 ai) 2 ≤ (K − 1) ∑K−1 i=1 a 2 i ; the last inequality follows by Lemma E.3 and η2k0 ≤ 4η 2 k. Hρ is defined in Definition A.4.

heading: D.4 BOUNDED VARIANCE
text: Proof of Lemma B.4. By assumption A.3, we have
E ∥∥∥∇f(θ)−∇f̃(θ)∥∥∥2
2 = E ∥∥∥∥∥ N∑ c=1 pc ( ∇f c(θc)−∇f̃ c(θc) )∥∥∥∥∥ 2
2
= N∑ c=1 p2cE ∥∥∥∇f c(θc)−∇f̃ c(θc)∥∥∥2 2
≤ dσ2 N∑ c=1 p2c ≤ dσ2 ( N∑ c=1 pc )2 := dσ2.

heading: E UNIFORM UPPER BOUND

heading: E.1 DISCRETE DYNAMICS
text: Lemma E.1 (Discrete dynamics). Assume assumptions A.1, A.2, and A.3 hold. We consider the generalized formulation in Algorithm 5 with the temperature
Tc,ρ = τ(ρ 2 + (1− ρ2)/pc)
given a correlation coefficient ρ. For any learning rate η ∈ (0, 2/m) and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have the `2 norm upper bound as follows
sup k
E‖θck − θ∗‖ 2 2 ≤ dD
2 + 6d
m ( max c∈[N ] Tc,ρ + σ2 m + γ2 md ) ,
where γ := maxc∈[N ] ‖∇f c(θ∗)‖2 and θ∗ denotes the global minimum for the function f .
Proof. First, we consider the k-th iteration, where k ∈ {1, 2, · · · ,K − 2, (K − 1)−} and (K − 1)− denotes the (K − 1)-step before synchronization. Following the iterate of Eq.(13) in a local client of c ∈ [N ], we have
E ∥∥θck+1 − θ∗∥∥22
= E‖θck − θ∗ − η∇f̃ c(θck)‖22 + √
8ηTc,ρE〈θck − θ∗ − η∇f̃ c(θck), ξk〉+ 2ηTc,ρE‖ξk‖22 = E‖θck − θ∗ − η∇f̃ c(θck)‖22 + 2ηdTc,ρ, (41)
where the last equality follows from Eξk = 0 and the conditional independence of θck − θ∗ − f̃ c(θck) and ξk. Note that
E‖θck − θ∗ − ηf̃ c(θck)‖22 = E‖θck − θ∗ − η∇f c(θck)‖ 2 2 + η
2E‖∇f c(θck)−∇f̃ c(θck)‖22 + 2ηE〈θck − θ∗ − η∇f c(θck),∇f c(θck)−∇f̃ c(θck)〉
= E‖θck − θ∗ − η∇f c(θck)‖ 2 2 + η 2E‖∇f c(θck)−∇f̃ c(θck)‖22 ≤ E‖θck − θ∗ − η∇f c(θck)‖22 + η2dσ2, (42)
where the first step follows from simple algebra, the second step follows from the unbiasedness of the stochastic gradient, and the last step follows from Assumption A.3. For any q > 0, we can upper bound the first term of Eq.(42) as follows
E‖θck − θ∗ − η∇f c(θck)‖22 = E‖θck − θ∗ − η(∇f c(θck)−∇f c(θ∗))− η∇f c(θ∗)‖22
≤ (1 + q)E‖θck − θ∗ − η(∇f c(θck)−∇f c(θ∗))‖22 + η2 ( 1 + 1
q
) ‖∇f c(θ∗)‖22
≤ (1 + q) (
1− ηm 2 )2 ︸ ︷︷ ︸
ψ2
E‖θck − θ∗‖ 2 2 + η 2
( 1 + 1
q
) γ2, (43)
where the first inequality follows by the AM-GM inequality; the second inequality is a special case of Lemma B.1 based on Assumption A.2, where no local steps is involved before the synchronization step. Similar results have been achieved in Theorem 3 Dalalyan (2017a). In addition, γ := maxc∈[N ] ‖∇f c(θ∗)‖2.
Choose q = ( 1+ψ2ψ ) 2 − 1 so that (1 + q)ψ2 = (1+ψ)
2
4 . Moreover, since ψ = 1 − ηm 2 , we get
1+ψ 2 = 1− 1 4ηm. In addition, we have 1 + 1 q = 1+q q =
(1+ψ)2
(1−ψ)(1+3ψ) ≤ 2 ηm . It follows that η2 ( 1 + 1
q ) ≤ 2η m . (44)
Combining Eq. (41), Eq. (42), Eq. (43), and Eq. (44), we have the following iterate E‖θck+1 − θ∗‖22 ≤ (
1− ηm 4 )2 ︸ ︷︷ ︸
:=g(η)
E‖θck − θ∗‖22 + 2ηdTc,ρ + η2dσ2 + 2ηγ2
m .
Note that 11−g(η) = 1
ηm 2 (1− ηm 8 ) ≤ 3ηm given η ∈ (0, 2 m ). Recursively applying the above equation k
times, where k ∈ {1, 2, · · · ,K − 1,K−} and K− denotes the K-step without synchronization, it follows that
E‖θck − θ∗‖22 ≤ g(η)k‖θc0 − θ∗‖22 + 1− g(η)k 1− g(η) · ( 2ηdTc,ρ + η 2dσ2 + 2ηγ2 m ) (45)
≤ ‖θc0 − θ∗‖22 + 3 ηm · ( 2ηdTc,ρ + η 2dσ2 + 2ηγ2 m ) ≤ dD2 + 6d
m ( max c∈[N ] Tc,ρ + σ2 m + γ2 md ) ︸ ︷︷ ︸
:=U
,
where the second inequality holds by g(η) ≤ 1, the third inequality holds because ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ] and η < 2m . In particular, the K-th step before synchronization yields that
E‖θcK− − θ∗‖ 2 2 ≤ dD2 + U. (46)
Having all the results ready, for the K-local step after synchronization, applying Jensen’s inequality E‖θcK − θ∗‖22 = E ∥∥∥∥ N∑ c=1 pcθ c K− − θ∗ ∥∥∥∥2 2
≤ N∑ c=1 pcE ∥∥θcK− − θ∗∥∥22
≤ dD2 + U, (47) Now starting from iteration K, we adapt the recursion of Eq.(45) for the k-th step, where k ∈ {K + 1, · · · , 2K − 1, (2K)−} and (2K)− denotes the 2K-step without synchronization, we have
E‖θck − θ∗‖22
≤ g(η)k−K · E‖θcK − θ∗‖22 + 1− g(η)k−K 1− g(η) · ( 2ηd max c∈[N ] Tc,ρ + η 2dσ2 + 2ηγ2 m ) ≤g(η)k−K(dD2 + U) + 1− g(η) k−K
mη/3
mη
3 U
≤dD2 + g(η)k−KU + (1− g(η)k−K)U ≤dD2 + U, (48)
where the second inequality follows by Eq.(47), the fact that 1− g(η) ≥ ηm/3 and η ≤ 2m , and the definition of U . The third one holds since g(η) ≤ 1. By repeating Eq.(47) and (48), we have that for all k ≥ 0, we can obtain the desired uniform upper bound.
Discussions: Since the above result is independent of the learning rate η, it can be naturally applied to the setting with decreasing learning rates. The details are omitted.

heading: E.2 CONTINUOUS DIFFUSION
text: Lemma E.2 (Continuous time). Assume assumption A.2 holds. We have the `2 norm upper bound as follows
sup t E ∥∥θ̄ct − θ∗∥∥22 ≤ 1m
( γ2
m + 2dτ
) ,
where γ := maxc∈[N ] ‖∇f c(θ∗)‖2 and θ∗ denotes the global minimum for the function f .
Proof. Since the synchronization is conducted at every time t, the essential temperature applied to each client is τ . Let q(θ̄ct ) = ∥∥θ̄ct − θ∗∥∥22. For any time t ≥ 0, applying Itô’s lemma leads to dq(θ̄ct ) = −2〈θ̄ct − θ∗,∇f c(θ̄ct )〉dt+ 2dτdt+ √ 8τ〈θ̄ct − θ∗,dW t〉
= −2〈θ̄ct − θ∗,∇f c(θ̄ct )−∇f c(θ∗) +∇f c(θ∗)〉dt+ 2dτdt+ √ 8τ〈θ̄ct − θ∗,dW t〉
≤ −2m ∥∥θ̄ct − θ∗∥∥22 dt− 2〈θ̄ct − θ∗,∇f c(θ∗)〉dt+ 2dτdt+√8τ〈θ̄ct − θ∗,dW t〉
≤ −2m ∥∥θ̄ct − θ∗∥∥22 dt+m∥∥θ̄ct − θ∗∥∥22 dt+ ‖∇f c(θ∗)‖22m dt+ 2dτdt+√8τ〈θ̄ct − θ∗,dW t〉
≤ −mq(θ̄ct )dt+ ( γ2
m + 2dτ
) dt+ √ 8τ〈θ̄ct − θ∗,dW t〉,
where the first inequality follows by Assumption A.2; the second inequality follows by the AM-GM inequality; the third inequality follows by the definition that γ2 = maxc∈[N ] ‖∇f c(θ∗)‖ 2 2.
In other words, we have
d(emtq(θ̄ct )) = me mtq(θ̄ct )dt+ e mtdq(θ̄ct ) ≤ memtq(θ̄ct )dt+ emt ( −mq(θ̄ct )dt+ ( γ2
m + 2dτ
) dt+ √ 8τ〈θ̄ct − θ∗,dW t〉 ) ≤ emt ( γ2
m + 2dτ
) dt+ √ 8τemt〈θ̄ct − θ∗,dW t〉.
The solution is upper bounded by emtq(θ̄ct ) ≤ em·0q(θ̄c0) + ∫ t 0 ( ems ( γ2 m + 2dτ ) ds+ √ 8τems〈θ̄cs − θ∗,dW s〉 ) .
By the martingale property of Itô integral, taking expectations yields Eq(θ̄ct ) ≤ e−mtEq(θ̄c0) + e−mt ( γ2
m + 2dτ )∫ t 0 emsds
= e−mtEq(θ̄c0) + 1− e−mt
m
(γ2 m + 2dτ )
≤ e−mtEq(θ̄c0) + 1− e−mt
m
( γ2 m
+ 2dτ︸ ︷︷ ︸ :=V
) , (49)
where the last inequality follows since the synchronization is conducted at any time step t. Since θ̄c0 is simulated from the stationary distribution π, by Lemma 12 Durmus & Moulines (2016) or Theorem 17 Cheng et al. (2018), we have
Eq(θ̄c0) = E‖θ̄c0 − θ∗‖22 ≤ dτ m ≤ 1 m ( γ2 m + 2dτ) = V m ,
which completes the proof.

heading: E.3 BOUNDED GRADIENT
text: Lemma E.3 (Bounded gradient in `2 norm). Given assumptions A.1, A.2, and A.3 hold, for any client c and any learning rate η ∈ (0, 2/m) and ‖θc0 − θ∗‖ 2 2 ≤ dD2 for any c ∈ [N ], we have the `2 norm upper bound as follows
E‖∇f̃ c(θck)‖22 ≤ 14dL2Hρ,
where Hρ = D2 + 1m maxc∈[N ] Tc,ρ + γ2 m2d + σ2 m2 .
Proof. Decompose the `2 of the gradient as follows E ∥∥∥∇f̃ c(θck)∥∥∥2
2 = E ∥∥∥∇f̃ c(θck)−∇f c(θck) +∇f c(θck)∥∥∥2 2
= E‖∇f c(θck)‖ 2 2 + E ∥∥∥∇f̃ c(θck)−∇f c(θck)∥∥∥2 2
+ 2E 〈 ∇f̃ c(θck)−∇f c(θck),∇f c(θck) 〉 ≤ E‖∇f c(θck)‖ 2 2 + σ 2d = E‖∇f c(θck)−∇f c(θ∗) +∇f c(θ∗)‖ 2 2 + σ 2d
≤ 2E‖∇f c(θck)−∇f c(θ∗)‖ 2 2 + 2E ∥∥∇f c(θ∗)∥∥22 + σ2d ≤ 2L2E‖θck − θ∗‖ 2 2 + 2γ 2 + σ2d ≤ 2dL2D2 + 12dL 2 m · ( max c∈[N ] Tc,ρ + σ2 m + γ2 md ) + 2γ2 + σ2d
≤ 14dL2 · ( D2 + 1
m max c∈[N ]
Tc,ρ + γ2 m2d + σ2 m2
) := 14dL2Hρ,
where the first inequality follows by Assumption A.3; the second inequality follows by Young’s inequality; the third inequality follows by Assumption A.1 and the definition that γ := maxc∈[N ] ‖∇f c(θ∗)‖2; the fourth inequality follows by Lemma E.1; the last inequality follows by κ := Lm ≥ 1.
F INITIAL CONDITION
Lemma F.1 (Initial condition). Let µ0 denote the Dirac delta distribution at θ0. Then, we have
W2(µ0, π) ≤ √ 2(‖θ0 − θ∗‖2 + √ dτ/m).
Proof. By Cheng et al. (2018), there exists an optimal coupling between µ0 and π such that
W 22 (µ0, π) ≤ Eθ∼π[‖θ0 − θ‖22] ≤ 2Eθ∼π[‖θ0 − θ∗‖22] + 2Eθ∼π[‖θ − θ∗‖22] = 2‖θ0 − θ∗‖22 + 2Eθ∼π[‖θ − θ∗‖22] ≤ 2‖θ0 − θ∗‖22 + 2dτ/m,
where the second step follows from triangle inequality, the last step follows from Lemma 12 Durmus & Moulines (2016) and the temperature τ is included to adapt to the time scaling.
Burkholder-Davis-Gundy inequality Let φ : [0,∞) → Rr×d for some positive integers r and d. In addition, we assume E ∫∞ 0 |ψ(s)|2ds < ∞ and let Z(t) = ∫ t 0 ψ(s)dWs, where Ws is a d-dimensional Brownian motion. Then for all t ≥ 0, we have
E sup 0≤s≤t |Z(s)|2 ≤ 4E ∫ t 0 |φ(s)|2ds. (50)

Review summaries:
This paper studies the federated learning problem. In this model, local clients are allowed to jointly train a model without sharing user data, and the central goal is to design communication efficient algorithms. The author gives a federated averaging Langevin algorithm and provide theoretical guarantees for strongly log-concave distributions. These results provide guidance on the choice of learning rates and local updates to minimize communication cost. The authors also consider applying correlated noises and using only partial device updates, which are more applicable in practice.

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. 

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.