On Convergence of Federated Averaging Langevin Dynamics
Authors: Wei Deng
Yian Ma
Zhao Song
Qian Zhang
Guang Lin
Venue: ICLR 2022 Submitted (2022)
Keywords: federated learning
Langevin dynamics
federated averaging
posterior inference
MCMC

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 stochastic-gradient 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.

Under review as a conference paper at ICLR 2022
ON CONVERGENCE OF FEDERATED AVERAGING LANGEVIN DYNAMICS
Anonymous authors Paper under double-blind review
ABSTRACT
We propose 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 stochastic- gradient 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 commu- nication 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 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.
INTRODUCTION
Federated learning (FL) allows multiple parties to jointly train consensus model without sharing user data. Compared 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 al. (2020), and fraud detection Zheng al. (2020).
A standard formulation of federated learning is distributed optimization framework that tackles communication costs, client robustness, and data heterogeneity across different clients Li al. (2020a). Central to the formulation is the efficiency of the communication, which directly motivates the cntuuunicaaiooeffffiieeeee federated averaging (FedAvg) McMahan et al. (2017). FedAvg introduces global model to synchronously aggregate multi-step local updates the available clients and yields distinctive properties in communication. However, FedAvg often stagnates at inferior local modes empirically due the data heterogeneity across the different clients Charles & Konečný (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 popular approach for inferring posterior distributions for large datasets, the stochastic gradient Markov chain Monte Carlo (SG-
MCMC) method Welling & Teh (2011); Vollmer et al. (2016); Teh et al. (2016); Chen et al. (2014); Ma al. (2015), which enjoys theoretical guarantees beyond convex scenarios Raginsky et al. (2017); Zhang et al. (2017); Mangoubi Vishmoi (2018); Ma et al. (2019). In particular, we examine in the federated learning setting the efficacy of the stochastic gradient Langevin dynamics (SGLD)

---

Under review as a conference paper at ICLR 2022
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-ban based FedAvg distributed sampling framework. Similar ideas have been proposed in federated postenior 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-bede 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 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 non- asymptotiiiic 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 pochasstic-radient 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 convergence result for FA-LD with decaying learning rates. This strategy accelerates the computation by logarithmic factor to achieve the precision e.
The algorithm yields appealing extensions: (1) we can choose to inject either independent correlated noise across local clients, yielding 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.
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), 2); 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. (202 I); 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 al. (2019); Karimireddy et al. (2020) focuses on standard assumptions in optimization (such as, convex, smooth, strongly-conver bounded gradient). The other line of work Li al. (2021); Huang 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 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).
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 al. (2014);
Distributed Monte Carlo methods Sub-posterior aggregation
2

---

Under review as a conference paper at ICLR 2022
Chen et al. (2016); Chowdhury & Jermaine (2018); Li et al. (2019a) propose 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 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, .... Let N denote the number of clients. For each c E [N]. we use I and V as the loss function and gradient of the function f in client V (-) is the unbiased stochastic gradient of f. In addition, The E (0, 1), where ne > O is the number of data as the weight the e-th client such that P. = Σ n.
denote Pe
points in the e-th client. Let T. denote the number of global steps to achieve the precision c. 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 standard algorithm in federated learning and is typically formulated into a distributed optimization framework as follows
min ( = ΣN (0) Ee The
( = Σ 1(0:
(1)
where 0 E IRd 100; Togg is a certain loss function based on and the data point Togi
One iterate of the FedAvg algorithm requires the following three steps:
Broadcast:
The center server broadeasts the latest model, 6s. to all local clients.
Local updates: For any e [N], the e-the client first sets 0f = 0g and then conducts 2 local steps:
Bk 0% (O)),
where 17 is the learning rate and V the unbiased estimate of the exact gradient VeS Synchronization: The local models are aggregated into unique model 0 Ed Po and sent 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 larger number of local steps K and higher order of data heterogeneity slows down the convergence. Notably, Eq. (1) can be interpreted as maximizing the likelihood function, which 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
0k+1 = 02 V ((Ox) + 22TTER.
where r is the temperature and & is a standard d-dimensional Gaussian vector. f(0) := MN (0) a energy function. f(0) unbiased estimate of f()). In the longtime limit, a well known result is that 0₂ converges weakly to the distribution π (0) x exp(-f(0)/5 Teh et al. (2016) as 17 + 0.
4 POSTERIOR INFERENCE VIA FEDERATED AVERAGING LANGEVIN DYNAMICS
The increasing concern for uncertainty estimation in federated learning motivates us to consider the simulation of the distribution (0) x (xp(-((0)/7) with distributed clients.
3

---

Under review as a conference paper at ICLR 2022
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, consider the following scheme: for any c E [N], the e-the client first sets 0% 02: and then conducts K 2 local steps:
0% / (0E) + V2
where V (0) 1 (0). (0) is the unbiased estimate of V (0) and Eq an independent Gaussian vector to be defined later.
Ps
Summing Eq. (2) from clients = 1 to N, we have the aggregated stochastic process as follows
Bkt1 02 -()))( v E1
where
N 0₂ =
N
N
N Σ
3x = Σ
Pe
Σ Petle
Vf(ak) = Σ Pe (()), & c=1
c=1 (3)
P
By the nature of the synchronization step, we always have 3g = 02 whether + 1 mod = 0 or not. In what follows, we can write
0x+1 = 02 (((b)) + 22
(4)
which resembles the SGLD algorithm except that the construction of stochastic gradients is dif- ferent and 0g is not accessible when k mod K * 0. Since our target is to simulate from (0) & exp(-f(0)/r), we expect that & a standard Gaussian vector. By the concentra- tion property of independent Gaussian variables, it is natural to set Ei = /g/ so that
Ek >N Po Σ VPeff and & is also standard Gaussian vector. Now we present it in Algorithm
Algorithm Federated averaging Langevin dynamics algorithm (FA-LD), informal version of Algorithm 4. 77K is the learning rate at iteration k. T the temperature. Denote by 0 the model parameter in the e-th client at the k-th step. Denote the immediate result of one step SGLD update from 0% by 35- E an independent standard d-dimensional Gaussian vector at iteration le for each client E [N]. A global synchronization is conducted every K steps.
1:
+ = 0 T/ ( + 2TkT/poSE.
(5)
2:
if l: + mod K +
(6)
Pe
N
if l: + 1 mod K = 0.
Algorithm Hybrid federated averaging Langevin dynamics algorithm (hFA-LD), informal version
of Algorithm 5. & is a d-dimensional Gaussian vector shared by all the clients; şg is an independent standard d-dimensional Gaussian vector at iteration l for each client [N]. p denotes the correlation coefficient. 1:
3 = 0g ()())(0))) + 2mm + 22 - 12)T/po5E.
2:
if k: + mod K +
EN PoPf.1
N
if + 1 mod K = 0.
We observe that the local process in Eq. (5) maintains a temperature 7 7/p. >> 7 to converge to the stationary distribution æ. Such a mechanism may limit the disclosure of individual data and shows a potential to ensure higher level of privacy.
4

---

Under review as a conference paper at ICLR 2022
5 CONVERGENCE ANALYSIS
In this section, we show that FA-LD converges to the stationary distribution (0) in the 2- 2-Wasserstein (W2) distance at a rate of (11vVT,) for strongly log-concave and smooth density. The W2 distance defined between a pair of Borel probability measures μ and Rd follows
(f11s, Bu
&
(2(,, v) =
inf ((CCuppiing((e..)
where || ll2 denotes the &2 norm on IRd and the pair of random variables (3,, RR Rd is a coupling with the marginals following [(3,)) = j and [(B,) = v. Note that L(.) denotes distribution of random variable. Such distance more appealing than the total variation the Kullback-Leibhee divergence in statistical machine learning applications for providing the estimates of the first and second order moments.
5.1 NOTATION AND ASSUMPTIONS
We make standard assumptions the smoothness and convexity of the functions f f ...... fNN which naturally yields appealing tail properties of the stationary measure π. Thus, we no longer require restrictive assumption the bounded gradient in & norm in Koloskova et al. (2019); Yu al. (2019); Li et (2020c). addition, control the distance between f and Vfs, we also assume bounded variance of the stochastic gradient in assumption 5.3.
Assumption 5.1 (Smoothness). For each E [N], we say f is L-smooth for some L > 0
(0( S (x) + (V - x) + _ flv 212 Vx, y E Rd.
Assumption 5.2 (Strongly convex). For each [N], fe is m-strongly convex for some >> 0
(z) +
4) 11y- Vz, Rd
Assumption 5.3 (Bounded variance, informal version of Assumption A.3). For each e E [N,, the variance of noise in the stochastic gradient V (xx in each client is upper bounded such that
[[IVV)(()) - (=)) S odd Va E
Quality of non-i..d data Denote by 0. the global minimum of f. Next, quantify the degree of the non-i.i.d data by 7 max,g[N] ||V f*(6))ll,, which is non-negative constant and yields larger scale if the data less identically distributed.
5.2 PROOF SKETCH
The proof hinges on showing the one-step result in the W2 distance. To facilitate the analysis, we first define an auxiliary continuous-ting time processes (0, (2)) without communication concerns
dØ - (.. + 22d
where 0, =N P. V=(0,)= = Pa (Q)). oc the continuous-time variable client e, and is a d-dimensional Brownian motion. The continuous-thn time algorithm is known to converge to fi where f(0) ΣN' fef (o)). Assume that dg simulates
the stationary distribution (0) x -
from the stationary distribution n, then it follows that 0g for any t 0.
5.2.1 DOMINATED CONTRACTION IN FEDERATED LEARNING
The first target is show certain contraction property of ||B 0 (VS(3) (V[()) based on distributed clients with infrequent communications, Consider a standard decomposition
3 (V)(3)) ((()))
5

---

Under review as a conference paper at ICLR 2022
||3 011 2n (3 8, -f(() tn [V) V f(0)22
Using E3), 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. Assume assumptions 5.1 and 5.2 hold. For any learning rate E (0, Lim) any [0=) (00) E Rd,, we have
N
|| - 1 (Vf()) - ((()))
2 (1 1pm) ||3 - 011 + L Σ Pe (11 + -11))
c=
divergence term
where 3 Ee Po
E and VS(B) V It implies that as long as the local parameters , and global B don't t differ each other too much, we can guarantee the desired convergence. In special when the communication is conducted at every iteration, the divergence term disappears and recovers the standard contraction Dalalyan & Karagulyan (2019).
Em VS(0)
5.2.2 BOUNDING DIVERGENCE
The following result shows that given finite number of local steps K, the divergence between 6 in local client and in the center is bounded & norm. Notably, since the Brownian motion leads to a lower order term (T)) instead of O(v²), naïve proof framework such Li et al. (2020c) may lead to crude upper bound for 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 r, (0,2/m) and ||6 0z dD2 for any c E [N], have the l2 upper bound of the divergence between local clients and the center as follows
N
P.BE 0ul12 < O(K3))d)+ + OK(nd)
The result also relies on showing a uniform upper bound in & norm, which avoids making extra bounded gradient assumptions.
5.2.3 COUPLING TO THE STATIONARY PROCESS
Note that 0, is initialized from the stationary distribution æ. The solution to the continuous-tine time process (7) follows:
0. 00=00 00
(O.)ddss /2T W.,
14 2 0.
(8)
Sett + (k + 1)n and 0g - Oun for (8) and consider synchronous coupling such that W W(k+1)) Win == Th is used to cancel the noise terms, we have
(k+1))
Uk++1m ORn
vf(0.)ds + V
(9)
km
Subtracting (4) from (9) and taking square and expectation both sides yield that B100 0z+1l 12 < (1 7mm/2) El,, 0z 112 + 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 with any learning rate y E (0, )t) and ||0g 0.112 dD2 for any
c E [N], where O, is the global minimum, for the function f. Then
(2 τ) S (1 (1--rm)) mm//2)- W (1ts T) + O(7² d(2 + k)). where Hi denotes the probability measure of 02 and K T./m is the condition number.

---

Under review as a conference paper at ICLR 2022
5.3 FULL DEVICE PARTICIPATION
5.3.1 CONVERGENCE BASED ON INDEPENDENT NOISE
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 & Karag- ulyan (2019). However, in scenarios of > with distributed clients, a divergence between the global variable 0g and local variable g 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 I with a fixed learning v (0; 2r. and 1,0g 0.1 S dD2 for any E [NV, we have"
W2(uu; x) (1- (1-mm//4) (V2d(D+V/7)
30 300%/mmm V(K2 + K Ko
where Ho denotes the probability measure of 0g at iteration k K denotes the number of local updates, = L/m, Y maXxe[N] [ and Ho == D max,e[M] mpe
o²
m2d
m
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 large learning rate.
Optimal choice of K. To ensure the algorithm to achieve the precision c based on the total number of steps T. and the learning rate 17. set
50xynmml VKK2 Ho S 5/22,
exp (- T (.) v2d(D + vT/m) S 6/22.
This readily leads to
m S 00
T. 2 02 log(d/e)
dK2 1))Ho)
Plugging into the upper bound of rym, it implies that to reach the precision e, it suffices to set
T. = 2dk dxK² + Hu log(d/c)).
(10)
It's obvious that Hg I(D²)) (1). thus we can conclude that the number of communication rounds is around the order
T. 2 (
K
K
where the value of 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 ( / /R). 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 TT). Similar results have been achieved by Stich (2019); Li et al. (2020c).
5.3.2 CONVERGENCE GUARANTEES VIA VARYING LEARNING RATES
Theorem 5.8 (Informal version of Theorem B.7). Assume assumptions 5.1, 5.2, and 5.3 hold. Consider Algorithm with an initialization satisfying 100 0.1 dD2 for any E [N] and the varying learning rate following
Mi
2L + (1/12)mk*
k: = 1,2, ....
Then for any k 2 0, we have
W2(Rs, T) < 45K (K2
Ho
(nmmd) 1/2
V/ 2 0.
For ease of presentation, report the result based on K² instead of (K 1)a. The upper bound based on (K 1)² is detailed in the supplementary file.
7

---

Under review as a conference paper at ICLR 2022
Note that the above result implies that to achieve the precision we require
md
1/2
W2 T) S 45x (K² + Ho
S.
2L + (1/12)mk
We therefore require dd) iterations to achieve the precision 6, which improves the D(E=2dlog(d// 2dlog(d/e)) rate for FA-LD with a fixed learning rate by a ((log(d/e)) factor.
Q(E
5.3.3 PRIVACY-ACCURACY TRADE-OFF VIA CORRELATED NOISES
Note that Algorithm I requires all the local clients to generate the independent noise Ei- Such mechanism enjoys the convenience of the implementation and yields 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 correlation coefficient p between different clients. Replacing Eq. (5) with
+ 0% f(( + p2 + v 22 p²)T/peSk.
(11)
where & is d-dimensional standard Gaussian vector shared by all the clients iteration k and É is dependent with E for any E [M] Following the synchronization step based on Eq. (6), we have
Bx+1 = 02 f(Ox)) + 21
(12)
where =Ex + ~² Lc=l v P- Since the variance of i.id variables is additive, it is clear that εN follows the standard d-dimensional Gaussian distribution. The inclusion of the correlated noise implicitly reduces the temperature for each client and naturally yields 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 correlated noise doesn't '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 2 0, Eq. (11) is equivalent to applying a temperature Te,p = (pp + (1-p2)/pd PP)). In particular, setting p 0 leads to Te,o T/Pe. which exactly recovers Algorithm 1; however, setting
leads to Te,1 = т, where the injected noise in local clients is reduced by 1 1/p. times. Now we adjust the analysis 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 E [0, n E (0, 1 and |00 0_ 12 s dD2 for any E [N]. we have
(T S S (1-1mm) (m-/m//)
(VEd( + (VV(DDVV/m)) T//m
30m nmd V(K2 + k) x)
where μi denotes the probability measure of 02. Hp = D + [NX T m²a
Such mechanism leads to trade-off between the efficacy of federation and accuracy and motivates us to exploit the optimal under the fffferential-privage theories Wang et al. (2015).
5.4 PARTIAL DEVICE PARTICIPATION
Full device participation enjoys appealing convergence properties. However, it suffers from the straggler's 's effect in real-world applications, where the communication limited by the slowest device. Partial device participation handles this issue by only allowing small portion of devices in each communication and greatly increased the communication efficiency in a federated network.
The first device- device-sampling scheme I Li et al. (2020b) selects total of S devices, where the e-th device selected with a probability P.- The first theoretical justification for convex optimization has been proposed by Li et al. (2020c).
(Scheme I: with replacement). Assume Sş = {mI n2, ... ,ns). where n, E [N] is a random number that takes value of with a probability pe for any j (.. S). The synchronization step follows that 0g = Eess. 8%
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 balanced data Li et al. (2020e).
8

---

Under review as a conference paper at ICLR 2022
(Scheme II: without replacement). Assume Sk {n1, n2.... ( random number that takes a value of with probability 3 for any j 21, 2. S}. Assume the data is balanced such that P1
.ns). where n, [N] is
PN N The synchronization step follows that
0g = Eces, Pelf = Eces, 0%.
Algorithm 3 Hybrid federated Averaging Langevin dynamics Algorithm (FA-LD) with partial device participation, informal version of Algorithm 6. Sy is sampled according to device-sampling rule based scheme or II.
1:
0k+1 0 (V) 22r + 22 /²)T/poCk-
2:
Sk+1
if + mod K #0
0
Eceseer
5 5P
if 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 hyperparameter E [0, 1}, fixed learning rate E (0, 2L |10g - 0.112 S dD² for any e E [N], we have
1 and
W2(us, x) S (1 (1-11110) (mm/4)/ (V2d(D- vT/m))
d 5 ((11-2))Cs)
+ 1030% V H22 x)) xk)
where Cs for Scheme and Cs N-S N-1 for Scheme II.
We observe that partial device participation leads to an extra bias regardless of the scale of y). To reduce such bias, suggest consider highly correlated injected noise, such p reduce the impact of the injected noise. By setting O( O/Vd/S) S e/3 and following similar learning rate achieve the precision within Q( logle(d) iterations given large
in section 5, number of devices satisfying 92(e-22 2d))
The device-sampling scheme provides viable solution handle the straggler's effect in full device participation and greatly accelerates the communication efficiency. In addition, scheme is rather robust to the data heterogeneity and doesn't Int require the data 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 system can only operate based the first S messages for the local updates. The device-sampling scheme II proposes concrete treatment to tackle this issue. Given balanced data across different clients and each device uniformly sampled, we can achieve 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 Il bound is better than scheme I.
6 CONCLUSION AND FUTURE WORK
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 & norm as in the optimization-dia drriven 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 underdamped Langevin dynamics Cheng et al. (2018), replica exchange Monte Carlo (also known as parallel tempering) Deng 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).

---

Under review as a conference paper at ICLR 2022
REFERENCES
Sungjin Ahn, Anoop Korattikara, and Max Welling, Bayesian Posterior Sampling via Stochastic Gradient Fisher Scoring. In Proc. of the International Conference on Machine Learning (ICML), 2012.
Sungjin Ahn, Babak Shahbaba, and Max Welling. Distributed Stochastic Gradient MCMC. In International Conference on Machine Learning (ICML), 2014.
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated Learning via Posterior Averaging: New Perspective and Practical Algorithms. In ICLR. arXiv:2010.05273, 2021.
Zachary Charles and Jakub Konečný. On the Outsized Importance of Learning Rates in Local Update Methods. arXiv:2007. 00878, 2020.
Changyou Chen, Nan Ding, and Lawrence Carin. On the Convergence of Stochastic Gradient MCMC Algorithms with High-order Integrators. In Advances in Neural Information Processing Systems (NeurIPS), pp. 2278-2286, 2015.
Changyou Chen, Nan Ding, Chunyuan Li, Yizhe Zhang, and Lawrence Carin. Stochastic Gradient MCMC with Stale Gradients. In Advances in Neural Information Processing Systems (NeurIPS), 2016.
Mingzhe Chen, Zhaohui Yang, Walid Saad, Changchuan Yin, H Vincent Poor, and Shuguang Cui. A Joint Learning and Communications Framework for Federated Learning over Wireless Networks. IEEE Trans. Wireless Communications, 2020.
Tianqi Chen, Emily B. Fox, and Carlos Guestrin. Stochastic Gradient Hamiltonian Monte Carlo. In Proc. of the International Conference on Machine Learning (ICML), 2014.
Xiang Cheng, Niladri S Chatterji, Peter I. Bartlett, and Michael I Jordan. Underdamped Langevin MCMC: A non-asymptoon -smmpttpptoiic analysis. In Conference on Learning Theory (COLT), pp. 300-323.
PMLR, 2018.
Arkabandhu Chowdhury and Chris Jermaine. Parallel and Distributed MCMC via Shepherding Distributions. In AISTAT; 2018.
Arnak S. Dalalyan. Further and Stronger Analogy Between Sampling and Optimization: Langevin Monte Carlo and Gradient Descent. In Conference on Learning Theory (COLT), June 2017a.
Amak S Dalalyan. Theoretical Guarantees for Approximate Sampling from Smooth and Log-concave Densities. Journal of the Royal Statistical Society: Series B, (9(3))651-6766 2017b.
Arnak S Dalalyan and Avetik Karagulyan. User-friendly Guarantees for the Langevin Monte Carlo with Inaccurate Gradient. Stochastic Processes and their Applications, 129(12):5278-5311 2019,
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Mare' aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large Scale Distributed Deep Networks. In Advances in neural information processing systems (NeurIPS), pp. 1223-1231, 2012.
Wei Deng, Qi Feng, Liyao Gao, Faming Liang, and Guang Lin. Non-Convex Learning via Replica Exchange Stochastic Gradient MCMC. In Proc. of the International Conference on Machine Learning (ICML), 2020.
Alain Durmus and Éric Moulines. Sampling from a Strongly Log-concave Distribution with the Unadjusted Langevin Algorithm. 1rXiv.1605.0155, 2016.
Mert Gürrbüzbalahan, Xuefeng Gao, Yuanhan Hu, and Lingjiong Zhu. Decentralized Stochastic Gradient Langevin Dynamics and Hamiltonian Monte Carlo. arXiv:2007.005900: 2021.
Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. FL-NTK: A Neural Tangent Kernel-bace Framework for Federated Learning Convergence Analysis. In International Conference on Machine Learning (ICML), 2021.
10

---

Under review as a conference paper at ICLR 2022
Yangsibo Huang, Zhao Song, Danqi Chen, Kai Li, and Sanjeev Arora. TextHide: Tackling Data Privacy in Language Understanding Tasks. EMNLP, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kemel: Convergence and Gener- alization in Neural Networks. In Advances neural information processing systems (NeurIPS), pp. 8571 -5710,580, 2018.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic Controlled Averaging for Federated Learning. In International Conference on Machine Learning, pp. 5132-5143, PMLR, 2020.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. First Analysis of Local GD on Heterogeneous Data. arXiv:1909.0444 04715, 2019.
Anastasia Koloskova, Sebastian U.Stich, and Martin Jaggi. Decentralized Stochastic Optimization and Gossip Algorithms with Compressed Communication. In Proc. of the International Conference on Machine Learning (ICML), 2019.
Chunyuan Li, Changyou Chen, Yunchen Pu, Ricardo Henao, and Lawrence Carin. Communication Efficient Stochastic Gradient MCMC for Neural Networks. In Proc. of the National Conference on Artificial Intelligence (AAAI), 2019a.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated Learning: Challenges, Methods, and Future Directions. IEEE Signal Processing Magazine, 37(3):00-60, 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smithy. Federated Optimization in Heterogeneous Networks. In Proceedings of the 3rd MLSys Conference, 2020b.
Wenqi Li, Fausto Milletari, Daguang Xu, Nicola Ricke, Jonny Hancox, Wentao Zhu, Maximilian Baust, Yan Cheng, Sébastien Ourselin, M Jorge Cardoso, al. Privacy-preserving Federated Brain Tumour Segmentation. In International Workshop on Machine Learning in Medical Imaging, pp. 131-141. Springer, 2019b.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the Convergence of FedAvg on Non-IID Data. In Proc. of the International Conference on Learning Representation (ICLR), 2020c.
Xiaoxiao Li, Yufeng Gu, Nicha Dvornek, Lawrence Staib, Pamela Ventola, and James $ Duncan. Multi-site fMRI Analysis using Privacy-preserving Federated Learning and Domain Adaptation: ABIDE Results. Medical Image Analysis, 2020d.
Xiaoxiao Li. Meirui Jiang, Xiaofci Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learning on Non-IID Features via Local Batch Normalization. In International Conference on Learning Repre- sentations (ICLR), 2021. URL https: /wwp//eeeeeeeeeeeeeeeeeeeeeeeeee net. / forum? i wtte://oommmmmmuuumuea dd-YiiiuUGGGGIGGGGGGG
Xuechen Li, Denny Wu, Lester Mackey, and Murat A. Erdogdu. Stochastic Runge-Kutta Accelerates Langevin Monte Carlo and Beyond. In Advances in Neural Information Processing Systems (NeurIPS), pp. 7746-7758, 2019c.
Y.-A Ma, E. B. Fox, T. Chen, and L. Wu. Inreversible samplers from jump and continuous Markov processes. Stat. Comput., pp. 1-26, 2018.
Y-A. Ma, N. S. Chatterji, X. Cheng, N. Flammarion, P. L. Bartlett, and M. I. Jordan. Is there an analog of Nesterov acceleration for MCMC? Bernoulli, 1991129--99,, 2021.
Yi-An Ma, Tianqi Chen, and Emily B. Fox. Complete Recipe for Stochastic Gradient MCMC. In Neural Information Processing Systems (NeurIPS), 2015.
Yi-An Ma, Yuansi Chen. Chi Jin, Nicolas Flammarion, and Michael I. Jordan. Sampling Can Be Faster Than Optimization. PNAS, 2019.
Oren Mangoubi and Nisheeth K. Vishnoi. Convex Optimization with Unbounded Nonconvex Oracles using Simulated Annealing. In Proc. of Conference Learning Theory (COLTT, 2018.
11

---

Under review as a conference paper at ICLR 2022
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication efffeiiennn learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282, PMLR, 2017.
H. McMahan, Eider Moore, Daniel Ramage, and Blaise Aguera y Arcas. Federated Learning of Deep Networks using Model Averaging. 2016.
Minsker, S. Srivastava, L. Lin, and D. B. Dunson. Scalable and Robust Bayesian Inference via Median Posterior. In International Conference on Machine Learning (ICML), 2014.
Wenlong Mou, Yi-An Ma, Martin J. Wainwright, Peter Bartlett, and Michael I. Jordan. High-Order Langevin Diffusion Yields an Accelerated MCMC Algorithm. Journal of Machine Learning Research (JMLR), 22:1-41, 2021.
W. Neiswanger, C. Wang, and E. Xing. Asymptotically Exact, Embarrassingly Parallel MCMC. arXiv:1311.4788, 2013.
Y. Nesterov. Introductory Lectures on Convex Optimization, in: Applied Optimization. Kluwer Academic Publishers, Boston, MA, 2004.
R. Nishihara, I. Murray, and R. P. Adams. Parallel MCMC with Generalized Elliptical Slice Sampling. Journal of Machine Learning Research, 15(1):2087-2112 2014.
Reese Pathaky and Martin J. Wainwright. Fedsplit: An Algorithmic Framework for Fast Federated Optimization. arXiv:2005.05238,, 2020.
Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex Learning via Stochastic
Gradient Langevin Dynamics: Nonasymptotic Analysis. In Conference on Learning Theory, June 2017.
R. Shen and Y. T. Lee. The randomized midpoint method for log concvee sampling. In Advances in Neural Information Processing Systems, pp. 2098-2109, 2019.
Reza Shokri and Vitaly Shmatikov. Privacy-preserving Deep Learning. In SIGSAC conference on computer and communications security (CCS). ACM, 2015.
Sebastian U. Stich. Local SGD Converges Fast and Communicates Little.
2019.
Yee Whye Teh, Alexandre Thiery, and Sebastian Vollmer. Consistency and Fluctuations for Stochastic Gradient Langevin Dynamics. Journal of Machine Learning Research, 17:1-33, 2016.
Sebastian J. Vollmer, Konstantinos C. Zygalakis, and Yee Whye Tch. Exploration of the (Non-) Asymptotic Bias and Variance of Stochastic Gradient Langevin Dynamics. Journal of Machine Learning Research, 17(159):1-44 2016.
Shiqiang Wang, Tiffany Tuor, Theodoros Salonidis, Kin K. Leung, Christian Makaya, Ting He, and Kevin Chan. Adaptive Federated Learning in Resource Constrained Edge Computing Systems. IEEE Journal on Selected Areas in Communications, 37(6):2055-13 2019.
Xiangyu Wang and David B. Dunson. Parallelizing MCMC via Weierstrass Sampler.
Yu-Xiang Wang, Stephen Fienberg. and Alex Smola. Privacy for Free: Posterior Sampling and Stochastic Gradient Monte Carlo. In ICMI., pp. 2493-2502, 2015.
Max Welling and Yee Whye Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. In International Conference on Machine Learning. 2011.
Blake Woodworth, Kumar Kshitij Patel, and Nathan Srebro. Minibatch vs Local SGD for Heteroge- neous Distributed Learning. urXiv:2006.047.0 2020.
Hao Yu, Sen Yang. and Shenghuo Zhu. Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning. In Proc. of Conference Artificial Intelligence (AAAI), 2019.
12

---

Under review as a conference paper at ICLR 2022
Yuchen Zhang, Percy Liang, and Moses Charikar. A Hitting Time Analysis of Stochastic Gradient Langevin Dynamics. In Proc. of Conference on Learning Theory (COLT), pp. 1980-2022, 2017.
Wenbo Zheng, Lan Yan, Chao Gou, and Yui Yue Wang. Federated meta-learning for fraudulent credit card detection. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence (IJCAI), 2020.
13

---

Under review as a conference paper at ICLR 2022
Roadmap. In Section A, we layout the formulation of the algorithm, basic notations, and definitions. In Section B, we present the main convergence analysis for full device participation. We discuss the optimal number of local updates based on a fixed learning rate. the acceleration achieved by varying learning rates, and the privacy-accuracy trade-off through correlated noises. In Section C, analyze the convergence of partial device participation through two device- device-sampling schemes. In Section D, we provide lemmas to upper bound the contraction, discretization and divergence for proving the main convergence results. Section E, we include supporting lemmas to prove results in the previous section. In Section F, we establish the initial condition.
A PRELIMINARIES
A. 1 BASIC NOTATIONS AND BACKGROUNDS
Let N denote the number of clients. Let Tg denote the number of global steps to achieve the precision Let K denote the number of local steps. For each E [N] := {1, 2, ... .N. we use fe and V denote the loss function and gradient of the function f in client c. For the stochastic gradient oracle, we denote by (-) the unbiased estimate of the exact gradient f of client c. In addition, we denote Pr. as the weight of the e-th client such that Pe: 2 0 and E Pr. 1. &g is an independent
standard d-dimensional Gaussian vector iteration for each client [N] and &n is unique Gaussian vector shared by all the clients.
Algorithm 4 Federated averaging Langevin dynamics algorithm (FA-LD). Denote by ff the model parameter in the e-th client at the k-th step. Denote the immediate result of one step SGLD update from 0% by 35. 4% is an independent standard d-dimensional Gaussian vector at iteration & for each client E [N]. A global synchronization is conducted every K steps. This is complete version of Algorithm 1.
1:
ph+1 = (0)) v 22tr///es
(13)
2:
if l: + 1 mod K + 0
0
(14)
Dd 1
if + mod K = 0.
k++
Inspired by Li et al. (2020c), we define two virtual sequences
N
N
Bk
Σ
P Ry
0y = Σ P.ff
(15)
which are both inaccessible when ls mod K + O. For the gradients and injected noise, we also define
vf(())) = PPVF Vf(0.) = N ( &
VPelE.
(16)
In what follows, it is clear that EVJ(0) En FV ff()) for any 0 E Rd and any e E [... Summing (13) from clients e = 1 to and combining (15) and (q.(, we have
Bk+1 02 n + 2
(17)
Moreover, we always have 3g = 02 whether l: + mod E = 0 or not by (q. and Eq.() In what follows, we can write
0x+1 02 T(0x)) + 22
(18)
which resembles the SGLD algorithm Welling & Teh (2011) except that the construction of stochastic gradients is different and 02 is not accessible when ls mod K o O. To facilitate the analysis, we also define an auxiliary continuous-timg time processes (0:)
d0, (0.) dt 22 dW, t
(19)
14

---

Under review as a conference paper at ICLR 2022
where 9, Ex= peff. Vf(0.) Pe BE is the continuous-tima variable at client and W is d-dimensional Brownian motion. The continuous-time algorithm is referred as Federated Averaging Langevin diffusion and is described as
d3; = -V dt + 2 dW;
0g = PoBE.
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 T and conver- gences to the stationary distribution (0) X exp(-f(0)TT), where f(0) Σ (0) Assume that Og simulates from the stationary distribution æ, then follows that 02 ~ π for any 2 0.
A.2 ASSUMPTIONS AND DEFINITIONS
Assumption A.1 (Smoothness). For each [N], we say f is T-smooth if for some T. 0
fwv) S Fx) (V) x) + fly- 212
-
Vz, y E R2
Note that the above assumption is equivalent to saying that
[V/ (=)(2)) I L|yy x||2, Vz, y Rd.
Assumption A.2 (Strong convexity). For each E [N], fe m-strongly corvex for some m
m (x)x) 2 (w) + (vf - y) lly 11 Vz, E Rd
An alternative formulation for strong convexity is that
(x) (y), x y) 2 m || 212 Ve, y E IRd.
Assumption A.3 (Bounded variance, restatement of Assumption 5.3). For each [N], the variance of noise the stochastic gradient (x) in each client is upper bounded such that
E V odd,
Vo
Rd
The bounded variance the stochastic gradient is a rather standard assumption and has been widely used in Cheng al. (2018); Dalalyan & Karagulyan (2019); Li et al. (2020c). Extension of bounded variance unbounded cases such as E[Il (x) V f=) < 6(122² 2 B²) for some M and 8 E [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 O. the global minimum of f. Next, we quantify the degree of the non-i.i.d data by 7 max |/V (0(0))ll2, which is non-negative constant and yields smaller scale if the data is more evenly distributed.
Definition A.4.
We define parameter Te.p H2, x and y?
Top:: :o:==(pf T(p² (1-p)))))).
H :=
D2
max Tep cE[N]
+
m²d
initialization
injected noise
data heterogeneity
stochastic
K = L/m,
2 = max [0)9()))
cE[N]
15

---

Under review as a conference paper at ICLR 2022
B FULL DEVICE PARTICIPATION
B.1 ONE-STEP UPDATE
Wasserstein distance We define the Wasserstein distance between pair of Borel probability measures μ and y on R as follows
+
(f119,) Bu Bul
V)(p, =
inf [(CCoppings((e))
where || l| denotes the E2 norm on I2 and the pair of random variables (Sμ, Bb) R x R is coupling with the marginals following ((3,)) μ and [(3,) = v. E(-) denotes distribution of random variable.
The following result provides crucial contraction property based on distributed clients with infre- quent synchronizations.
Lemma B.1 (Dominated contraction property, restatement of Lemma 5.4). Assume assumptions A. and A.2 hold. For any learning rate y) E (0, b any [00) )e=11 [B0) E TRd, have
N
011 2 4nL Σ - S (1 - rpm) 113-013
- - -
Pe (11 - B11 + 110 0112).
where Ee-1
>ー IEPP Vf(0)
Σ (o and Vf(8)
Ee ( We postpone the proof into Section D.1. The above result implies that as long as the local parameters 8 8 and global 0, 3 don't differ each other too much, we can guarantee the desired convergence.
The following result ensures a bounded gap between 0g and Øc n_ l2 norm for any 2 and E [N]. We postpone the proof of Lemma B.2 into Section D.2.
Lemma B.2 (Discretization error). Assume assumplions A.I, A.2, and A.3 hold. For any 2 0, any learning rate (0,2/ m)) and ||eg 00 < dD² for any [N). the iterates of (0.) based the continuous dynamics of (q.(1 satisfy the following estimate
R22 + LTT) 1-1dd.
E 0 0h 0115 81
The following result shows that given finite number of local steps K, the divergence between 0f in local client and in the center is bounded in &g norm. Notably, since the non-differentiable Brownian motion leads to lower order term O(7) instead of O(1²) in L² norm, navve proof may lead to 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. A.2, and A.3 hold. For any learning rate 1) E (0, 2/m) and |00 0. 2 < dD² for any c E [N], we have the l2 upper bound of the divergence between local clients and the center follows
N
Σ
0ul13 < 112(K - 1)3 + 8(K - 1)) + N(1 - -)²)), P.
where Hp, and 72 are defined Definition A.4.
The following presents a standard result for bounding the gap between (f(() and Vī(0). We delay the proof of Lemma B.4 into Setion D.
Lemma B.4 (Bounded variance). Given assumption A.3, we have
(0) (i(0) < d o
V
Having all the preliminary results ready, now we present a crucial lemma for proving the convergence of all the algorithms.
16

---

1 Under review as a conference paper at ICLR 2022
Lemma B.5 (One step update, restatement of Lemma 5.6). Assume assumptions A. A.2, and A.3 hold. Consider Algorithm with independently injected noise p = 0, any learning rate E (0, 2L and |10 0-112 dD2 for any [N], where 0, is the global minimum for the function f. Then
(2(
1pm 2) (2i,, + (o
k)),
where μt denotes the probability measure of 02, Ho. R and 72 are defined Definition A.4.
Proof. The solution of the continuous-time process (1. follows that
8: Au-- L vf(O)dds + 2T W.,
Vt 2 0.
(20)
Set + (k: + 1)n and 0u Okn for (20) and consider synchronous coupling such that W(k+))y - Wkn = ViS
(k+1)n
(k+1)m 0kn k
Vi(0,)ds + (2T(W((+1)) - Way)
(M+1)n
=0kn
f((0)dds 2
(21)
We first denote Co = vf(On) f00)). Subtracting (q. from (2,(2 yields that
((++)) Ox+1
(k+1)) Tf(0))ds
= =0kn 0k + +
kn
(k+1)n
=0k- 02 vf(02 +Okn - 02) vF((w)
Vf(0.) - v)(Okn) ds
(22)
kn
(k+1)n
=0kn 02 (f(kk + 0nn 0n) - v)(0k)
Vf(0.) v((ūun) ds +T/K.
kn
==
Taking square and expectation both sides, have
( l)n 0x+1 12
= km 02- nxx 1x11 + ElI + 2n B(Okn n - YR Sc)
EG-=0 and mutual indepenneence
< (1 + q) Bll On x 12= + (1 1/9) Yu + 11lvi&11
< (1+ 9) ((1 vm) 1E| 0km 0x113 2 + 4nL Σ Pe (E1 0R + 0xll2)
+ (1+ + 1/9) E + 2
$ <1+ (1 7m) B 0ull2
+ +2(K -
(1+1//) E| +
(23)
where the first inequality follows by the AM-GM inequality for any >> 0, the second inequality follows by Lemma B. and Assumption A.3. The third inequality follows by Lemma B.3 with p = 0: moreover, the continuous-time process conducts synchronization at any time step, hence 0i, 0x. 1 Since the learning rate follows 21. < 1 and Lemma B.3 is clearly satisfied.
m+L < 2. the requirement of the learning rate for Lemma B.
m
17

---

where Hi denotes the probability measure of Oz. Ho, K. and 72 defined Definition A.4. Under review as a conference paper at ICLR 2022
Recall that ø 1 rym, get
1t 1 2m.. Choose q
1+0 1 so that (1+ 9)0 (1++)
26
1 1m.. In addition, we have 1 + q 9
1+a
1+0
2 It follows that
I-6
mm
1
mm
(1 + q) (1 rpm) - 1
2 nm,
1 + q -
I 1.5, (1+1/q)= - mm
(24)
1 - nm
where the second inequality holds because n (0, 1 2m
For the term B|| in (q.( have the following estimate
(1+1)n
12
IE
v.(.0) V f(0km)) ds
kn
(R+1)m
17
E
((.)) - f((0km))1
(k+1)n
11 kn
E
Pe
v V -
ds
(k+1)n N
S
"
Pe E V v
( 112
km
(b+1)m N
S TT?
Pc B 0 - 0z 12
knn
c=1
(k+1)M
1n2 d
S 11L2 km
807
LT) + 16tdd ds
dd Ho + 1.,72122r.
(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.I, and the last inequality follows by Lemma B.2. The last equality holds since $ + LT S Lmllg and L/m.
Plugging (24) and (25) into (23), we have 7i
E (1 0x+1l12 <
El 0K 0x
+ 672n) 2 dl K - 1
3 Ho + 2 I)L
-ا + 1202d
L²
+ 161 dL 3xttt= + 32n
m
Choose the specific Langevin diffusion 0 stationary regime, we have W2 (Mk: π) Ell0kk 0x (++
and W22 (px+1. x) & (
0x+1 112. Arranging the terms, we have
(2 T) I (1- 1/m W2(,e, T) + 400n/² dL (o2Hoo(K - 1)2 + k))
where 7) S z 1, mr S LT I L I L maxx To,0 LmHo, and 0² I L Ho applicd to the result.
B.2 CONVERGENCE VIA INDEPENDENT NOISES
Theorem B.6 (Restatement of Theorem 5.7). Assume assumptions A. A.2, and A.3 hold. Consider
Algorithm 4 with a fixed learning rate E (0, + and ||06 0. < dD2 for any c E [N], we have
1mm 1--m)
& (vzd VF/m)) + +30x mpnd- V(KK - 1)2 + K)Hlg
W2(nn x) S
18

---

Under review as a conference paper at ICLR 2022
Proof. Iteratively applying Theorem B.5 and arranging terms, we have that
W2 T) _ (1
7m
& (24, T)
1 - (1 yh +
400r Hol(H - 1)2 + x))
2
- (1
yea 2
<
mm
W.2 (1o, T) + mpm
2 400 d Hol(K - 1)2 + xx)
2
mm
k
W2 (No, т) + 8000="mmd(( 1)² + +k)
(26)
2
where A: m By Lemma F.1 and the initialization condition |006 0. 112 I dD² for any E [N],
L have that
W2(Ho, x) S 2d VT/m).
Applying the inequality (1
mm
(1 )² completes the proof.
Discussions
Optimal choice of K. To ensure the algorithm to achieve the precision based on the total number of steps T. and the learning rate 17. we can set
(VI(K- - 1)2 HO
30k /nmd
I
1H T.
(V2d(D+ + vT/m) $
6
2
This directly leads to
nmm S min
m 22
o
dk 1)2 H) m)Hb))
2
}
T. 2 Ω log (
mn
Plugging into the upper bound of 17, it implies that to reach the precision level e, it suffices to set
T. Ω dk ((K
1) Ho
)
d -
log
(27)
²²
Since Ho = QDD2 + 2),, we observe that the number of communication rounds is around the order
T. K K
= 22
K
where the value of either large K first decreases and then increases with respect to K, indicating that setting K small may lead high communication costs and hurt the performance. Ideally, K should be selected in the scale of ((VR). Combining the definition of T, in (q.)) this suggests an interesting result that the optimal K should be in the order O(VT.) Similar results have been achieved by Stich (2019): Li et al. (2020c).
B.3 CONVERGENCE VIA VARYING LEARNING RATES
Theorem B.7 (Restatement of Theorem 5.8). Assume assumptions A.1, A.2, and A.3 hold. Consider
Algorithm with initialization satisfying |l0| 0.113 < dD² for any [N] and varying learning rate following
1
7e 21 (1/12)mk
k 1,2,
Then for any k have
Wa(mu; RT) S 45kv (5K//(KK 1)2 Ho (m-dd) 1/2
vk: 20 0,
19

---

Under review as a conference paper at ICLR 2022
Proof. We first denote
C₂ 300//00KKKKK 1)2 + N).
Next, we proceed to show the following inequality by the induction method
d
1/2
1.50 (7amd) 1/2
W2(Hu, T) & 1.50
44 2 0,
(28)
2 (1/12)mk
where the decreasing learning rate follows that
T 2L + (1/12)mk
(i) For the case of k 0, since
CR 2 4v AV Ho 2 AVE
D2 1
max Too 2 4v 4V44/// k/dd VdD2 + m cE[N]
d max To.
m [N]
2 2/Kk/W22 T)
(29)
where the last inequality follows by Lemma F.1 and |10| 0.112 < dD² for any e [N].
clear that W2 T) S 1 Cr V md
T. s 1.50 Viomm by Eq.))
(ii) If now that (28) holds for some k 2 0, follows by Lemma B.5 that
(2 S (1 m
W (px: x) 40.
1)² K)
< (1 mmm
W2 π) +
Ogd
2
2
< (1
m
2.250mmmmd +
7k
2.25
2
3
< (1
mh m
2.250 2.250mmmmd.
6
Since (1
mm
s (1
Mu 12
we have
W2 W2 R)) S <1
Mkm
1.50,(( 1/2
12
To prove W2(m-11, T) < 1.5C (1CC((7)=+ (mk+ 1/2 suffices to show (1 -
mem 1/2
12 m 77++1, which is
detailed as follows
mem
1/2
VI2(2AL + mk m)
1/
12
(24I. mk)3/2
/12(24L + mkk - 1/2 m/2
24L + mk
V12
(24L + m(k + 1)) = 78+1
where the last inequality follows since
(24L + mk m2(24L + mk + m)) - (24L + mk))
The above result implies that to achieve the precision require
md
1/2
W2(n: +) S 1.5C, 2L + (1/12)mk
The means that only require k: = (( to achieve the precision By contrast, the fixed learning
rate requires T. = $ log (dfe)) which is much slower than the algorithm with varying learning rate by O( (log (d/e)) times.
20

---

Under review as a conference paper at ICLR 2022
B.4 PRIVACY-ACCURACY TRADE- OFF VIA CORRELATED NOISES
Note that Algorithm 4 requires all the local clients to generate the independent noise 5g- Such mechanism enjoys the convenience of the implementation and yields 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 hyperparameter, namely the correlation coefficient p between different clients. Replacing (13) with
+ 0% - f + 2 + 2n - -PL1-PP)T/PEE
(30)
where & is d-dimensional standard Gaussian vector shared by all the clients at iteration k, & is unique d-dimensional Gaussian vector generated by client [N] only. Moreover, & is dependent with Eg for any E [N]. Following the same synchronization step based Eq.(14), we have
0k+1 02 7) V F(Ox) + 2
(31)
where & = PCE + 12 co= v P Since the variance of iid variables is additive, it is clear that & follows the standard d-dimensional Gaussian distribution. The inclusion of the correlated noise implicitly reduces the temperature and naturally yields 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 p 2 0, (30) is equivalent to applying a temperature Te = Tpp + (1 p²) p/Pe) In particular, setting p = 0 leads to Te,0 (1 -02) / which exactly recovers Algorithm 4; however, setting
1 leads to Te,1 = where the injected noise in local clients is reduced by 1/p. times. Now we adjust the analysis as follows
Theorem B.8 (Restatement of Theorem 5.9). Assume assumptions A. A.2, and A.3 hold. Con- sider Algorithm 5 with correlation coefficient E [0, a fixed learning rate 1) E (0, 27. and |10g 1111 S dDD for any E [M], we have
W2(Un, x) & (1-17)²
3m & (VZi(D + VT/m) + +300g/m and v (((K((1)))))) 1)2 + k)
where μy denotes the probability measure of On. Hpp K and 72 are defined as Definition A.4.
Algorithm 5 Hybrid federated averaging Langevin dynamics algorithm (hFA-LD). Denote by fg the model parameter in the c-th client at the k-th step. Denote the immediate result of one step SGLD update from og by B6. <g an independent standard d-dimensional Gaussian vector at iteration & for each client e E [N] and É is a d-dimensional standard Gaussian vector shared by all the clients. denotes the correlation coefficient of the injected noises. global synchronization is conducted every K steps. This is a complete version of Algorithm 2.
1:
Bh+1 0% (()) + 22 + 2 p²)T/
2:
if ks + mod K + 0
0
Me Pr
if k + mod K = 0.
C PARTIAL DEVICE PARTICIPATION
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 federated network.
21

---

Under review as a conference paper at ICLR 2022
C.I UNBIASED SAMPLING SCHEMES
The first device-sampling scheme I Li et al. (2020b) selects a total of S devices, where the e-th device selected with a probability Pe- The first theoretical justification for convex optimization has been proposed by Li et al. (2020c).
(Scheme I: with replacement). Assume Sk {ni na: ....... .ns). where nj E [N] is random number that takes a value of with a probability for any j E {1,2, (1, The synchronization step follows that 0g $ Eees. Og.
Another strategy is to uniformly select devices without replacement. We follow Li et al. (2020c) and assume S indices are selected uniformly without replacement and the synchronization step the same as before. In addition, the convergence also requires an additional assumption on balanced data Li al. (2020c).
(Scheme II: without replacement). Assume Są = {M1, n2. .nst. where n, E [N] is random number that takes a value of с with probability ş for any j E {1.2. .... S}. Assume the data is balanced such that p1
pn t The synchronization step follows that
0g Eces, pelf- = Eces, Og.
Algorithm 6 Hybrid federated averaging Langevin dynamics algorithm (hFA-LD) with partial device participation. &G is the independent Gaussian vector proposed by each client [N] and & is unique Gaussian vector shared by all the clients. denotes the correlation coefficient. A global synchronization is conducted every K steps. Sg is subset that contains S indices according to device-sampling rule based on scheme or II. This is complete version of Algorithm 3.
1:
s 0%
V (08)
22n + 2 p²)T/poEE.
2:
Pi+1
if k: mod K 00
0
Eees4
8 8
if mod K 0.
Lemma C.1 (Unbiased sampling scheme). For any k mod K 0 based on scheme or II, we have
EA, E
Σ 0
Bx
P.B..
cES²
Proof. According to the definition of scheme or II, we have 02 = Eces, &%. In what follows,
EAg- = EE 0% scheme Il in particular.
3/2 Ecses, ΣN PeBs Em= PeBi. where p1 = P2 = = PN for
C.2 BOUNDED DIVERGENCE BASED ON PARTIAL DEVICE
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 p E [0, 1}, any learning rale 1j (0, 2/m) and |106- 0. < dd2 for any e E [N), we have the following results
For Scheme I, the divergence between 02 and B& is upper bounded by
112
8 + KKnd + N(1 (1)-
ElIBA - 0x113
For Scheme II, assuming the data is balanced such that P1 = ... PN = 1. the divergence between 0% and By is upper bounded by
- BI 0n.13 < S(N-1)
112K2 8KndT(p N ((11-)1))
021))
where H p. and 72 are defined Definition A.4.
22

---

Under review as a conference paper at ICLR 2022
Proof. We prove the bounded divergence for the two schemes, respectively.
For scheme with replacement, 0% = Eees, 3% for a subset of indices Sg. Taking expectation with respect to Sk, we have
1
E 0x 3xll2
BIST Ball2
Σ Pe 1 3x
(32)
where the first equality follows by the independence and unbiasedness of 0h for any i E [S].
To further upper bound (3.) we follow the same technique in Lemma B.3. Since k mod K 0,
kn- = k - is also the communication time, which yields the same 0% for any c E [M]. in what follows,
Pe 11
Bell2
Pe.
|
0ko - (Bx 0kk))
-
Σ
Pr. |1.8
0k 11
(33)
where the last inequality follows by Bk Ec= (32) and (33), we have
.
PaPa
and
BB|
Er 12 I E| Combining
1
E110.
Bxll2
Pe
1
0KD 12
8
1
|1
0g l12
Pe
5
0=1
N
k-1 1
2Kr² V)((0)) 1 + 4KKndT(p2 + (1 -p2)/vo)
Σ
<
8
k=kp
k-1
2k 2 v (08)) ||3 +KKdd (p² 2 + (1 - /² p21/po)
Pc
5
k=
28
< K2
2
Hp
Kndd + N(1 (2))
s
where the last inequality follows similar argument as Lemma B.3.
For scheme II, given P1 P2
PN
1 No we have 0%
5 Eces, 18 which leads to
1
1
Bl00s 1211 E
Σ
3
BK
S2 E
TeES (3t BBb)
S
OES,
2
where IA is indicator function that equals to if the event A happens.
Plugging the facts that TP(c S&) y and P(c1, E Sx) the above equation, we have
5(8-1) for any a # 02 [N] into N(N-1)
E 0
3x 2
P(c E 118E - -All + Σ P(e1, E S4)(f
Br. B -Ba)
S
cEIN]
01402
1
N
1I - -x112 + Σ
S
(891 - Bk, 802 BB)
SN
SN(N - 1)
01402
g
N Σ
N
|1 - 3x112,
S(N
- 1)
c=1
23

---

Under review as a conference paper at ICLR 2022
where the last equality holds since ECE(N) 11.g - Bull2 + Ecitos (5% Pk (kk) || Bxll2 0. 0.
Eventually, we have
N S
N Σ
E E - Ball2
E N
| - Bel12
S(N - 1)
N
$
S(N - 1) EF N
N S
Σ
11 02o12
N - S
28K2 dT. 4KndT + N(1 - -D²))
S(N - 1)
where the first inequality follows a similar argument as in Eq.(33) and the last inequality follows by Lemma B.3.
C.3 CONVERGENCE VIA PARTIAL DEVICE PARTICIPATION
Theorem C.3 (Restatement of Theorem 5. 5.10). Assume assumptions A. A.2, and A.3 hold. Con- sider Algorithm 6 with correlation coefficient p [0, 1}, fixed learning rate 1) E (0, 2r and |10g - 11 < dD² for any c E [N], we have
W2(u,, R) < (1 (1-4) ))
(V2d(D+ + VT/m))
CxdT
+ 30x rpmdd / HHl(K +)2
+x) + -//
(02 + N(1 - p2))
Sm
where CK
kmn Km
Cs = for Scheme I and Cs
N-S for Scheme II. N-T
Proof. Note that
1 - 0x 0x-1112
1. 3k+1 + 8k+1 0x+1 l12
0 - Sc+1 1 + Bl1 - 0x+1 (0(0)))) 3et1, - Ok+1)
= = ( - Sx+1 1 El| - 1x+1l12
where the last equality follows by the unbiasedness of the device-sampling scheme in Lemma C.I.
k mod K + O, we always have 3x+1 Ok+1 and E|],0++ 0x+1 _ = Following the same argument as in Lemma B.5, both schemes lead to the one-step iterate as follows
( T) < (1
7pm
2 (2 T) + 400, dL HL,HH((KK 11² + +x).
(34)
If k + 1 mod K = 0, combining Lemma C.2 and Lemma B.5, we have
4K
W (1 S (1- ) W2 m) + 45. k)
(p2 + N(1 - (2))
S
(35)
where Cs 1 for Scheme and Cs N for Scheme II. N
Repeatedly applying (34) and (35) and arranging terms, we have that
m7m k W2 +) + 17m )
2 45. H. (K2 + +1))
W3 T) - (1-
24

---

Under review as a conference paper at ICLR 2022
(1 (1 mm K [L/KK
<KKdT
-1)
(1 2 )R
mm
(2 x) + 900mmdn. 10Ommdn2Ho(l Ho((K -1)2 Rk)
kmk
4Kd
mmK (0² + ([11-p)]CC
CR
=
m
W X) + 2000mdd""Ho(('"
Ho((K - 1)2 + K)
2
4C
(0² + N(1 - ()))
Sm
where the second inequality follows by (1 - r) K se- for any 2 0.
D BOUNDING CONTRACTION, DISCRETIZATION, AND DIVERGENCE
D.1 DOMINATED CONTRACTION PROPERTY
Proof of Lemma B.1 Given client index [A], applying Theorem 1.1.12 Nesterov (2004) leads
to
mL
(y x, x, V V)) 2 L + m
11 - 2112
L (V
Vx, y E IRd.
(36)
In what follows, we have
| - -
| - 012 2n (8 0, 1V/0)))) v/0)
(37)
For the second item in the right hand side, we have
N
PcLB 0, (39 - (09))
N
Σ
Pc(B 8" + 80 - 0 + - 0, -
cml
Σ Pe ((85 - V ()-(8)) (0- 0r, V/(0
Σ (B 00, -
+
N
2
De ((m + L.) ||3e - 312 + (m + L.) 110 - -012 + 2(m L.)
|V ()9
6=1
N
mL
+
|||3- - -8112 + L+ [V)(()> - v {(((0))1
L +
cml
2 (m L)
Pe (1135 -B112 + 100- 011)
mL L +
||,3 011
+
(38)
2 + m)
25

---

Under review as a conference paper at ICLR 2022
where the first incquality follows by the AM-GM inequality and Eq.(36), respectively; the last inequality follows by Jensen's inequality such that
N
N
2 = =113 01
Pell - 09112 Σ
2
Σ
P 89)
2
N
N
2 [V v
Pe
( (0)))
Σ
- - (
V
12
Plugging (38) into Eq.(3 have
|| 0 17 (Vf(3)) ff(())
< (1 2mmL,
1
1113 011 + nl
m
m
sDif -
N Σ
+ 2m + L)
Pe
(118- - B12 + 110- - -0133
c=1
N
<1 - mm) -13 4nL Σ
Pe (1 -B112 + 110- - -012).
c=1
where the last inequality follows by m-L 2 1, m 1 L, 1 2a S (1 - a)2² for any a, and
2L
17 E (0, _1.-
D.2 DISCRETIZATION ERROR
Proof of Lemma B.2. For any s E [0, oo), there exists certain & E N+ such that s E [kT, (k + 1)m).
By the continuous dynamics of Eq. (19), we have
Ø = 8 nill
+ (s km) V (B + 22 dw.,
kn
which suggests that
|/ V2rdW,
|00
0% a 112 S (s (o-kk)ll\
sup (E+ 1)m)
sup [E| 13))
We first square the terms on both sides and take Young's inequality and expectation
E
|0 8i 112 < 2
(s
kn)
f
2
sup
(R (R41)))
kn
Then, by Burkholder-Daviss Davis-Gundyy inequality (50) and Itô isometry, we have
sup
000
0
111
- 2T (s
8
> E 2r
(k+k))
km
11/ + 167d7.
(39)
By Young's inequality and the smoothness assumption A. we have
E V
= f
-f
26

---

Under review as a conference paper at ICLR 2022
s 2B
V (()L (,1 ff 2
s 22
0 L
0.11 + 27
7
2dF )) + 27
S 211
m
m
S 4dk
4LT)
(40)
where the third inequality follows by Lemma E.2, the fourth step holds since K 2 1. Combining Eq. (39) and Eq. (40), we have
k27 + LT) + 16ndT.
100- - ge 11 11 I 817²
E sup (++1)))
D.3 BOUNDED DIVERGENCE
Proof of Lemma B.3. For any k 2 0, consider ko K such that k: S ko and 0g = 0ke for any
2 O. It is clear that k - kn S K 1 for all k 2 Consider the non-increasing learning rate such
that T/aa I 21I for all k ho S K 1.
By the iterate (18) we have
B 0xll2
Pe
0% (0x- 0kn)112
N
Pe
11o112
N
k-1
Σ
P.
Σ
2(K 1)
V (0)
+ +4(K 1)ndd(( (2 +11 (11-1²)/pp) /²)/pe)
k=ku
NV
k-1
Σ
Σ
2(K 1)m v (02)
1 2
+ 4(K 1)7 dl dT + (1 -2))/p))
Pe
k==o
< 112(K 1) nh. H, + 8(K dT (p² 2- NN11 p²))
where the first inequality holds by B| B0l < 11 for stochastic variable 0; the second inequality follows by (EK-1 a.)² < (K 1) af: the last inequality follows by Lemma E.3 and n2 4mh Hp is defined in Definition A.4.
K-1
D.4 BOUNDED VARIANCE
Proof of Lemma B.4. By assumption A.3, we have
v v '= E
Σ Pe
Vfr(8))
-
Σ
2
(
27

---

Under review as a conference paper at ICLR 2022
N
S do² 2
Σ P - do²
Σ Pc
= do
E UNIFORM UPPER BOUND
E.1 DISCRETE DYNAMICS
Lemma E.1 (Discrete dynamics). Assume assumptions A.I, A.2, and A.3 hold. We consider the generalized formulation in Algorithm 5 with the temperature
Te.p= T(p + + (1)pe)/pe)
given a correlation coefficient p. For any learning rate 1 E (0, 20)) and 108 E [N], we have the & norm upper bound as follows
0..112 S dD² for any
6d
72
sup O 0-12 dDD2+ +
max Te,p cE[N]
+
+
md)
m
where y == [NX / and 0, denotes the global minimum for the function f.
Proof. First, we consider the k-th iteration, where k E (1,2, ......... K 2, (K -))- and (K 1)- denotes the (K- 1) Jstep before synchronization, Following the iterate of Eq.(13) in local client of E [N], we have
E 1 0_112 11
Elgg - f V. 0. (B))) ElI
Eli0E O. "> + 2ndTT..
(41)
where the last equality follows from By = and the conditional independence of 8% - 0. F((EE) and &x. Note that
05 - - (BE)
0
0- -/ I (0 (" (0) (
+ 27TEBA 0. V (f) vF(()))
0 0. / / (1)) (
05
0. -7 -
(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 g >> we can upper bound the first term of (42) as follows
E 6, - v (
= =E 0. / (0E) (0.)) - (
< (1 + E - 0, - -onn((fff (V( - V/ +1² (1+22) [VS
1pm 2 2
E 0.113 + n" 1 _ x2
$ (1 ()
(43)
where the first inequality follows by the AM-GM inequality; the second inequality is a special case of Lemma B.I based on Assumption A.2, where no local steps is involved before the syn- chronization step. Similar results have been achieved in Theorem Dalalyan (2017a). In addition, y = m (
28

---

Under review as a conference paper at ICLR 2022
Choose q = ( (220) )² so that (1 q)1p² (1++)²
Moreover, since &b 2. we get
1 1- - nm.. In addition, we have + 1
Ltg
(1+#)²
S 2 It follows that
q
=
(I-WX1436)
2n
n
S
(44)
Combining Eq. (41), Eq. (42), Eq. (43), and Eq. (44), we have the following iterate
E Bl - 0.112 v (1 mm 2
E 0.1 2ndToo + +o² + 2ny²
(n)
Note that
4/) S nm given 17 (0, m) Recursively applying the above equation
1-9(n) 4 (1
times, where k E {1, 2. K 1, K- and K_ denotes the K-step without synchronization,
follows that
E| O6 0.112 s 111) -0.112 +
1 - -(1)**
(2ndTT n2=
207 (45)
1n)
m
3
217??
$108
0,112
2nydTe.p o,p 1/2
1pm.
m
6d
a² +
72
S dD2
max cE[N]
Te,p
+
m
m
md
where the second inequality holds by g(n) I 1, the third inequality holds because ||e| 0.112 S dD2 for any e E [N] and
2 In particular, the K-th -tth step before synchronization yields that
B 0.1 z dD² U.
(46)
Having all the results ready, for the K-local step after synchronization, applying Jensen's 's inequality
N
El
0.12
E
Σ
P - 02 12
N
<
P... 1
10112
<
dD +U,
(47)
Now starting from iteration K, we adapt the recursion of Eq.) for the k-th step, where k e (K + 1,... 2K- 1, (2K) and (2K)- denotes the 2 -teep without synchronization, we have
El00 - 02.12
1 - - (nd max Te,p + 17= + 2072 cE[N]
I Ig(n) k-K K 0 - 0.112
1- 9())
< (dDD2 +U) - g(1))* k mel
mn/3 3 U
<dD2 + 9(1)
K-K + (1 g(n) g(n)*~KK k-
KU
<dD2 +U,
(48)
where the second inequality follows by (47), the fact that 1 - g(n) 2 mmm3 and S definition of U. The third one holds since (()) S 1.
and the m
2
By repeating (47) and (48), we have that for all & 2 0, we can obtain the desired uniform upper bound.
Discussions: Since the above result is independent of the learning rate 17, it can be naturally applied to the setting with decreasing learning rates. The details are omitted.
29

---

Under review as a conference paper at ICLR 2022
E.2 CONTINUOUS DIFFUSION
Lemma E.2 (Continuous time). Assume assumption A.2 holds. We have the &g norm upper bound as follows
(1122
72 2d7)
sup 8: 0a112
m
m
where 7 = maX (V))()))]lz and 0., denotes the global minimum, for the function f.
Since the synchronization conducted at every time t, the essential temperature applied to each client is r. Let ((()) = | 11 For any time 2 0, applying Itô's lemma leads to
Proof.
(g)))) = -0g - 0. v (t + 2drdt + (0, - 0. ...
2drdt + V8T 0., dw S - 2m 10; 0- 1| dl 2(0 - 0., V/ (0.) 2drdl + v (0p 0,, dW,)
20g 0.. ( (f)
s 2m 0G 0- 112 dt + m 0 0.112ar dt [v dt + 2drdt + V8r 0., dW.)
("²² + 2dl dt + (8 0,, dw 4)
<-mq(0;)dt
where the first incquality follows by Assumption A.2; the second inequality follows by the AM-GM inequality; the third inequality follows by the definition that 7² max,e[N] ((.))l)
In other words, we have
d(e mt (01)) = me (0,))tt + Th de(())
S me ((8))dt + me (-mq(BF)dt + (22 + 2dr dt + VTT 6., aw,)
(I- 2dF dt + "r mh (0S dW,.
S en
m
The solution is upper bounded by
mi ((0)) < e m-0 ((()) + ("(sme) e
(72 (I+ 2dF ds vereme 02, (lW))
By the martingale property of Itô integral, taking expectations yields
Eq(0,) - (0) mn (?2
+ 2dT
m
e -mt
??
= = S Se-mEEq((f)
2d7)
m
m
22
+ 2dr
(49)
m
m
where the last inequality follows since the synchronization is conducted any time step t. Since 0g simulated from the stationary distribution by Lemma 12 Durmus & Moulines (2016) or Theorem 17 Cheng et al. (2018), we have
< 1,2 + 2dT) = m m m
V
(g = El 0.113 <
which completes the proof.
30

---

Under review as a conference paper at ICLR 2022
E.3 BOUNDED GRADIENT
Lemma E.3 (Bounded gradient in E2 norm). Given assumptions A.I, A.2, and A.3 hold, for any
client and any learning rate 1) E (0, 2/m) and 100 0-11 - dD² for any c E [N], we have the E norm upper bound as follows
E (0 & 14d1 Hp:
where Hp =D= maX Topp + m m
Proof. Decompose the lg of the gradient as follows
[ - (05)
2 (()) (
=EIVS
+ 2E () V (f)) ((F
S S (0)) odd
=E V (0)) (0.) (0))) + o²d
s 2 / (0g) (0,)))) 2 V / (0.) odd
< 2212B|0 1 + 2y² + +
+ 12dL2
0²
72
< 2d
max Te.p
+
+
+ 22. + o
CEIN]
m
md
1
T 72
o²
S 14dL?.
D2
14dT. H
m cE[N]
m2d
m
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 % max.e[N]] f (2. the fourth inequality follows by Lemma E.1; the last inequality follows
by L 2 1.
F INITIAL CONDITION
Lemma F.1 (Initial condition). Let μo denote the Dirac delta distribution at 0g. Then, we have We(no: x) I V2( - 0x l| + / dr/m).
Proof. By Cheng et al. (2018), there exists optimal coupling between #40 and v such that
W (4o, ②) S EØ 100 - 01:3]
$ 2E x ||10 0.113] + 2E |||0 9 0x1131
= 00 0.13 1/2 + +2E.- |110 0-
< 2 0_ 1| + 2d7/m,
where the second step follows from triangle incquality, the last step follows from Lemma 12 Durmus & Moulines (2016) and the temperature 7 is included to adapt to the time scaling.
Burkholder-Davis-Gume inequality Let o [0, oo) + R for some positive integers and d. addition, we assume Efe (ds < oo and let ((t) = fo (s)dW,,, where W, a d-dimensional Brownian motion. Then for all 0, have
sup [Z(s) < 4E { (dss) 05.5t
(50)
31


Reviews:
[{"review_id": "IXgZeuhTcrh", "reviewer": "Reviewer_sxs5", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper has shown extensive theoretical analysis for the proposed FA-LD method, such as convergence for independent noise and varying learning rates, with or without full device participation. However, given the existing theoretical analysis for SGLD based on 2-Wasserstein distance and theoretical analysis for FedAvg with or without full device participation, the results are straightforward. Could the authors show the difficulty of the theoretical analysis in this paper?\n\nEmpirical studies are missing in the current version. The experiment section should be added to the paper to validate the theoretical findings and investigate the performance of FA-LD. Particularly, the reviewer would like to see the performance of FA-LD in terms of different local updates, learning rates, participating devices, etc, and its comparison with other distributed LD methods. \n\n Typo:\nsection 3.2, \"a energy function\"", "summary_of_the_paper": "The paper proposes Federated Averaging Langevin Dynamic algorithm (FA-LD). The theoretical guarantees for the proposed algorithm are developed and their relationships with noise type, heterogeneity, and learning rate are studied. ", "summary_of_the_review": "The paper has given extensive theoretical convergence analysis for FA-LD algorithm via 2-Wasserstein. However, empirical studies to verify theoretical findings are missing in the main paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "G0DhCsjZmhl", "reviewer": "Reviewer_dS2k", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nAssuming strong convexity (Assumption 5.2) seems completely ridiculous in this setting. One of the major motivations for considering stochastic gradient Langevin dynamics is to improve the performance of non-convex stochastic optimization algorithms so as to escape spurious stationarity points and arrive at global extrema. In fact, this can be rigorously shown to tend towards the distribution of global extrema without convexity with non-asymptotic rate given in Raginsky et al 2017\n\nThere are no numerical experiments that demonstrate the utility of the approach. How restrictive is it to consider log-concave distributions? This should preclude most neural network training in practice, which is the primary emphasis of federated learning, as compared with more classical methodologies based upon consensus/multi-agent optimization.\n\nIs there a practical phenomenon where standard federated learning or its variants that use Lagrange multipliers such as FedDyan and SCAFFOLD get stuck, whereas this version with randomized perturbations does not? I feel that this scenario is exactly in non-convex settings, which is excluded from consideration in this work.\n\nIt is difficult to make sense of what are the specific technical innovations in the analysis here, and how they are a departure from earlier results. ", "summary_of_the_paper": "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.\n", "summary_of_the_review": "\nThis work is mainly of theoretical interest, as it considers randomized perturbations in a federated averaging setting. However, its technical scope precludes the most useful instances of it in practice, and it is missing experiments. Therefore, it seems of limited usefulness in my view.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Emx8L8HvBGz", "reviewer": "Reviewer_JGao", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main contribution is the theoretical guarantees for a federated averaging Langevin algorithm for strongly log-concave distributions with non-i.i.d data. \n\nThe proof sketch is well written and clear. It seems that the proof is sound.\nHowever, most of the analysis technique can be founded in lituerature, the proof technique is thus standard.\n\nI think the writing has room to improve, but overall it is clear and easy to follow.\n\nI have the following concerns:\n1. In Section 3.1, I think it is better to stress the $\\beta_k$ is an auxiliary sequence. Besides, how $\\beta_k^c$ involves the Local Updates step in one iterate of the FedAvg algorithm seems to be incorrect. I think it should add $\\theta_k c = \\beta_k^c$ for $2 \\le k \\le K$.\n2. Algorithm 1 is easy to understand since it is an approximation of FedAvg. However, Algorithm 2 is not so intuitive. The author didn’t explain why we should add an additional \\textbf{shared} noise $\\dot{\\xi}_{k}$ and make the final noise correlated with the original noise $\\xi_{k}^{c}$ in Algorithm 2. How does such a shared noise ensure privacy and what kind of definition of privacy is used ? The last paragraph of Section 4 is quite vague and amphibolous. To substantiate the point, I think more argument and discussion should be added.\n3. The current presentation of main result is troublesome. The author could first introduce the main result and then provide a proof sketch. After all, not everyone is interested in the proof technique. What’s more, the analysis for algorithm 2 is deferred to Section 5.3.3, but its introduction is Section 4. I think it is somewhat too late and not helpful for readers to grasp the algorithm. \n4. The first paragraph of Section 5.2 seem to have a type. $\\nabla f\\left(\\bar{\\theta}_{t}\\right)$ should be defined as $\\sum_{c=1}^{N} p_{c} \\nabla f^{c}\\left(\\bar{\\theta}_{t}\\right)$.\n5. Though this paper is theoretical, the author could still provide some numerical experiments to show the convergence of FA-LD. Those experiments can further validate the theorems, for example, the trade-off of $K$ and the effect of sampling methods. After all, FA-LD is different from FedAvg. I don’t know whether the convergence of FA-LD has been empirically studied before.", "summary_of_the_paper": "The paper proposes a federated averaging Langevin algorithm (FA-LD) and analyzes its convergence under non-iid and partical participation settings.", "summary_of_the_review": "The delcaimed contribution include FA-LD and its theoretical guarantees.\nHowever, the analysis technique is commonly used in litureature, while no empirical expeirments to illustrate the effectiveness of FA-LD.\nBesides, there are some presentation issues and ambiguity on motivation.\nI don't think this paper is ready for publication. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "8BC3pyVQ3pK", "reviewer": "Reviewer_jFv4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n\n1. The idea of sampling using Langevin diffusion in the federated learning setting is new. \n\nWeakness:\n\n Since this is a theoretical paper I am basing my decision on its theoretical contribution.\n\n1. Novelty in the proof of LD is incremental.  Only part of the proof which is different from the traditional proof of LD is the divergence term. But showing that this term is small (Lemma B.3) follows very simply from typical strongly convex optimization techniques. All the results corresponding to partial participation and varying rates follow almost directly  from the proof of the main theorem and does not have much theoretical novelty.\n\n2. I liked the idea of allowing for correlated noise but I could not find the proof of theorem B.8 (theorem 5.9) in the paper. So it's difficult to gauge the difficulty of this part or even the correctness of this result. If I have missed the proof, and the authors could point me towards it, that would be great. \n ", "summary_of_the_paper": "In this work, the authors:\n\n1.  propose a federated learning version of Langevin diffusion sampling. \n\n2. 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.\n\n3. analyze the partial participation setting. ", "summary_of_the_review": "The paper lacks theoretical novelty but I liked the idea of introducing LD in federated setting. I may change my decision after discussion with other reviewers.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "4-ZKLNPMTml", "reviewer": "Reviewer_4zTT", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors propose the federated averaging Langevin algorithm (FA-LD) for posterior inference. The authors present non-trivial non-asymptotic convergence analysis for FA-LD for distributions with strongly smooth and strongly convex Hamiltonian and with bounded variance of noise in the stochastic gradient. The assumption of bounded gradient in l2 norm is not required in contrast to the previous work FedAvg as well as many others in literature. The analysis of FA-LD indicates the number of global steps and the choice of the learning rate. More importantly, it shows that the number of local steps should be set roughly as the order of square root of the condition number. The authors further obtain non-trivial results in the cases of varying learning rates in each step, privacy-accuracy tradeoff via correlated Gaussian noises, and computation model with partial device participation. \n\nThe highlights of this paper are novel theoretical analysis of the FA-LD algorithm. It demonstrates how the injected noise, the data heterogeneity, and the stochastic-gradient noise affect the convergence. It would be nice if the authors could also provide some lower bound or worst case analysis for the FA-LD algorithm, to show that the current convergence guarantees are tight or close to tight. For example, how does FA-LD behave for Gaussian posteriors? Is the convergence rate in the main theorem asymptotically optimal for Gaussian?\n\nDetailed comments:\nPage 3, right before Section 4: “\\tilde{f} is a unbiased estimate of f” \nShouldn’t it be “\\grad \\tilde{f} is an unbiased estimate of \\grad f”? \n", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, I think this is a nice paper with strong theoretical results. It provides insights to the theoretical study of standard sampling algorithms in federated learning. The paper is also well-written.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a federated averaging Langevin dynamics (FA-LD) for numerical mean prediction with uncertainty quantification under the setting of federated learning. Convergence analysis for the proposed method under the smoothness and strong-convex assumptions is also provided, and the results are summarized in Theorems 5.7-5.10, each of which bounds the Wasserstein-2 distance $W_2(\mu_k,\pi)$ between the model distribution $\mu_k$ and the target distribution $\pi$ under different settings.

This paper received 5 reviews in total, with scores 6, 5, 3, 5, and 3. Some reviewers evaluated positively the novelty of the idea of using the Langevin dynamics in the federated setting, which I would also like to acknowledge. Upon reading the paper by myself, however, I find that the mathematical formulations are in some places not correct. What I think problematic is the third equation in equation (3): The right-hand side is a function of $N$ variables $\\\{\theta_k^c\\\}$, and they undergo different local updates at different clients when $k\not\equiv 0\mod K$ (i.e., the synchronization does not take place). Also $\nabla\tilde{f}^c$ is in general a nonlinear function of its argument. Therefore, the right-hand side cannot be written in general as a function of a single variable $\theta_k$ which is defined as $\theta_k=\sum_{c=1}^Np_c\theta_k^c$, making this equation incorrect. This problem would affect various parts of the arguments to follow in this paper, such as the first two equations in equation (16) on page 14, the two inline equations just after equation (16), equation (18), the second equality in the inline equation in page 15, line 1, and the third line in equation (25) on page 18, to mention a few. Thus I have to question the validity of the theoretical development in this paper.

Another point I would like to mention is that I did not understand the definition of Schemes I and II in Section 5.4. It is not stated at all that $\mathcal{S}_k$ is a random quantity here. Furthermore, the conditions "with/without replacement" are not described at all.  Still another point to mention is that I did not understand the claim in page 7, lines 30-31. Does it mean: If one knows the number $T_\epsilon$ of steps to achieve the precision $\epsilon$, then one should set the number $K$ of local steps per synchronization should be set of the order of $\sqrt{T_\epsilon}$. But $T_\epsilon$ depends on $K$, so that it would be unnatural to assume that one knows $T_\epsilon$ irrespective of $K$ in the first place.

Because of these, I would judge that this paper is not yet ready for presentation in its current form. I would therefore not be able to recommend acceptance of this paper.

Minor points:
- Citation style: The authors use throughout the paper what are called the *narrative citations* even though there are occasions where what are called the *parenthetical citations* (the author name and publication date are both enclosed in parentheses) should be used.
- page 3, line 7: is (the -> an) unbiased stochastic gradient; There are several unbiased estimators for the gradient, and what is mentioned here is only one instance of them.
- page 3, lines 23-24: The aggregation should take place not on each client but on the central server.
- page 3, line 36: a(n) energy function; a(n) unbiased estimate
- page 5, lines 17-20: The contents of Assumptions 5.1 and 5.2 are not assumptions but definitions.
- page 6, line 2: to obtain (the -> a)  lower bound
- page 6, line 18: $\mathcal{D}^2$ is undefined.
- page 8, line 39: (a -> the) probability $p_c$ if it is meant to be the one defined in page 3, line 8. Otherwise, use of the same symbol to represent different quantities should be avoided.
- page 14, line 25: mod ($E$ -> $K$) =0
- page 15, line 30: $H_\rho^2$ -> $H_\rho$
Decision: Reject