Title: ON THE CONVERGENCE OF FEDAVG ON NON-IID DATA

Abstract: Federated learning enables a large amount of edge computing devices to jointly learn a model without data sharing. As a leading algorithm in this setting, Federated Averaging (FedAvg) runs Stochastic Gradient Descent (SGD) in parallel on a small subset of the total devices and averages the sequences only once in a while. Despite its simplicity, it lacks theoretical guarantees under realistic settings. In this paper, we analyze the convergence of FedAvg on non-iid data and establish a convergence rate of O( 1T ) for strongly convex and smooth problems, where T is the number of SGDs. Importantly, our bound demonstrates a trade-off between communicationefficiency and convergence rate. As user devices may be disconnected from the server, we relax the assumption of full device participation to partial device participation and study different averaging schemes; low device participation rate can be achieved without severely slowing down the learning. Our results indicates that heterogeneity of data slows down the convergence, which matches empirical observations. Furthermore, we provide a necessary condition for FedAvg on non-iid data: the learning rate η must decay, even if full-gradient is used; otherwise, the solution will be Ω(η) away from the optimal. * Equal contribution.
Here, G, Γ, p k , and σ k are problem-related constants defined in Section 3.1. The most interesting insight is that E is a knob controlling the convergence rate: neither setting E over-small (E = 1 makes FedAvg equivalent to SGD) nor setting E over-large is good for the convergence. This work also makes algorithmic contributions. We summarize the existing sampling 3 and averaging schemes for FedAvg (which do not have convergence bounds before this work) and propose a new scheme (see Table 1 ). We point out that a suitable sampling and averaging scheme is crucial for the convergence of FedAvg. To the best of our knowledge, we are the first to theoretically demonstrate

1. INTRODUCTION:

Federated Learning (FL), also known as federated optimization, allows multiple parties to collaboratively train a model without data sharing (Konevcnỳ et al., 2015; Shokri and Shmatikov, 2015; McMahan et al., 2017; Konevcnỳ, 2017; Sahu et al., 2018; Zhuo et al., 2019) . Similar to the centralized parallel optimization (Jakovetic, 2013; Li et Wang et al., 2018) , FL let the user devices (aka worker nodes) perform most of the computation and a central parameter server update the model parameters using the descending directions returned by the user devices. Nevertheless, FL has three unique characters that distinguish it from the standard parallel optimization Li et al. (2019) .

First, the training data are massively distributed over an incredibly large number of devices, and the connection between the central server and a device is slow. A direct consequence is the slow communication, which motivated communication-efficient FL algorithms (McMahan et al., 2017; Smith et al., 2017; Sahu et al., 2018; Sattler et al., 2019) . Federated averaging (FedAvg) is the first and perhaps the most widely used FL algorithm. It runs E steps of SGD in parallel on a small sampled subset of devices and then averages the resulting model updates via a central server once in a while. 1 In comparison with SGD and its variants, FedAvg performs more local computation and less communication.

Second, unlike the traditional distributed learning systems, the FL system does not have control over users' devices. For example, when a mobile phone is turned off or WiFi access is unavailable, the central server will lose connection to this device. When this happens during training, such a non-responding/inactive device, which is called a straggler, appears tremendously slower than the other devices. Unfortunately, since it has no control over the devices, the system can do nothing but waiting or ignoring the stragglers. Waiting for all the devices' response is obviously infeasible; it is thus impractical to require all the devices be active.

Third, the training data are non-iid 2 , that is, a device's local data cannot be regarded as samples drawn from the overall distribution. The data available locally fail to represent the overall distribution. This does not only bring challenges to algorithm design but also make theoretical analysis much harder. While FedAvg actually works when the data are non-iid McMahan et al. (2017) , FedAvg on non-iid data lacks theoretical guarantee even in convex optimization setting.

There have been much efforts developing convergence guarantees for FL algorithm based on the assumptions that (1) the data are iid and (2) all the devices are active. Khaled 2018) doesn't require the two mentioned assumptions and incorporates FedAvg as a special case when the added proximal term vanishes. However, their theory fails to cover FedAvg.

Notation. Let N be the total number of user devices and K (≤ N ) be the maximal number of devices that participate in every round's communication. Let T be the total number of every device's SGDs, E be the number of local iterations performed in a device between two communications, and thus T E is the number of communications.

Contributions. For strongly convex and smooth problems, we establish a convergence guarantee for FedAvg without making the two impractical assumptions: (1) the data are iid, and (2) all the devices are active. To the best of our knowledge, this work is the first to show the convergence rate of FedAvg without making the two assumptions.

We show in Theorem 1, 2, and 3 that FedAvg has O( 1 T ) convergence rate. In particular, Theorem 3 shows that to attain a fixed precision , the number of communications is S t ∼ W(N, K, p)

EQUATION

1 K k∈St w k t O( 1 T ) 5 Ours S t ∼ U(N, K) k∈St p k N K w k t O( 1 T ) 6

that FedAvg with certain schemes (see Table 1 ) can achieve O( 1 T ) convergence rate in non-iid federated setting. We show that heterogeneity of training data and partial device participation slow down the convergence. We empirically verify our results through numerical experiments.

Our theoretical analysis requires the decay of learning rate (which is known to hinder the convergence rate.) Unfortunately, we show in Theorem 4 that the decay of learning rate is necessary for FedAvg with E > 1, even if full gradient descent is used. 4 If the learning rate is fixed to η throughout, FedAvg would converge to a solution at least Ω(η(E -1)) away from the optimal. To establish Theorem 4, we construct a specific 2 -norm regularized linear regression model which satisfies our strong convexity and smoothness assumptions.

Paper organization. In Section 2, we elaborate on FedAvg. In Section 3, we present our main convergence bounds for FedAvg. In Section 4, we construct a special example to show the necessity of learning rate decay. In Section 5, we discuss and compare with prior work. In Section 6, we conduct empirical study to verify our theories. All the proofs are left to the appendix.

2. FEDERATED AVERAGING (FE DAV G ):

Problem formulation. In this work, we consider the following distributed optimization model:

EQUATION

)

where N is the number of devices, and p k is the weight of the k-th device such that p k ≥ 0 and

EQUATION

where (•; •) is a user-specified loss function.

Algorithm description.:

Here, we describe one around (say the t-th) of the standard FedAvg algorithm. First, the central server broadcasts the latest model, w t , to all the devices. Second, every device (say the k-th) lets w k t = w t and then performs E (≥ 1) local updates: 

w k t+i+1 ←-w k t+i -η t+i ∇F k (w k t+i , ξ k t+i ), i = 0, 1, • • • , E -1 where η t+i

w t+E ←- N k=1 p k w k t+E .

Unfortunately, the full device participation requirement suffers from serious "straggler's effect" (which means everyone waits for the slowest) in real-world applications. For example, if there are thousands of users' devices in the FL system, there are always a small portion of devices offline. Full device participation means the central server must wait for these "stragglers", which is obviously unrealistic.

Partial device participation. This strategy is much more realistic because it does not require all the devices' output. We can set a threshold K (1 ≤ K < N ) and let the central server collect the outputs of the first K responded devices. After collecting K outputs, the server stops waiting for the rest; the K + 1-th to N -th devices are regarded stragglers in this iteration. Let S t (|S t | = K) be the set of the indices of the first K responded devices in the t-th iteration. The aggregation step performs

w t+E ←- N K k∈St p k w k t+E .

It can be proved that N K k∈St p k equals one in expectation.

Communication cost. The FedAvg requires two rounds communications-one broadcast and one aggregation-per E iterations. If T iterations are performed totally, then the number of communications is 2T E . During the broadcast, the central server sends w t to all the devices. During the aggregation, all or part of the N devices sends its output, say w k t+E , to the server.

3. CONVERGENCE ANALYSIS OF FE DAV G IN NON-IID SETTING:

In this section, we show that FedAvg converges to the global optimum at a rate of O(1/T ) for strongly convex and smooth functions and non-iid data. The main observation is that when the learning rate is sufficiently small, the effect of E steps of local updates is similar to one step update with a larger learning rate. This coupled with appropriate sampling and averaging schemes would make each global update behave like an SGD update. Partial device participation (K < N ) only makes the averaged sequence {w t } have a larger variance, which, however, can be controlled by learning rates. These imply the convergence property of FedAvg should not differ too much from SGD. Next, we will first give the convergence result with full device participation (i.e., K = N ) and then extend this result to partial device participation (i.e., K < N ).

3.1. NOTATION AND ASSUMPTIONS:

We make the following assumptions on the functions 

F 1 , • • • , F N .

(v) ≤ F k (w) + (v - w) T ∇F k (w) + L 2 v -w 2 2 . Assumption 2. F 1 , • • • , F N are all µ-strongly convex: for all v and w, F k (v) ≥ F k (w) + (v - w) T ∇F k (w) + µ 2 v -w 2 2 .

Assumptions 

E ∇F k (w k t , ξ k t ) -∇F k (w k t ) 2 ≤ σ 2 k for k = 1, • • • , N . Assumption 4.

The expected squared norm of stochastic gradients is uniformly bounded, i.e.,

E ∇F k (w k t , ξ k t ) 2 ≤ G 2 for all k = 1, • • • , N and t = 0, • • • , T -1

Quantifying the degree of non-iid (heterogeneity). Let F * and F * k be the minimum values of F and F k , respectively. We use the term Γ = F * -N k=1 p k F * k for quantifying the degree of non-iid. If the data are iid, then Γ obviously goes to zero as the number of samples grows. If the data are non-iid, then Γ is nonzero, and its magnitude reflects the heterogeneity of the data distribution.

3.2. CONVERGENCE RESULT: FULL DEVICE PARTICIPATION:

Here we analyze the case that all the devices participate in the aggregation step; see Section 2 for the algorithm description. Let the FedAvg algorithm terminate after T iterations and return w T as the solution. We always require T is evenly divisible by E so that FedAvg can output w T as expected. Theorem 1. Let Assumptions 1 to 4 hold and L, µ, σ k , G be defined therein. Choose κ = L µ , γ = max{8κ, E} and the learning rate η t = 2 µ(γ+t) . Then FedAvg with full device participation satisfies

EQUATION

where 

EQUATION

EQUATION

)

Alternatively, we can select K indices from [N ] uniformly at random without replacement. As a consequence, we need a different aggregation strategy. Assumption 6 assumes the K indices are selected uniformly without replacement and the aggregation step is the same as in Section 2. However, to guarantee convergence, we require an additional assumption of balanced data. Assumption 6 (Scheme II). Assume S t contains a subset of K indices uniformly sampled from [N ] without replacement. Assume the data is balanced in the sense that 

p 1 = • • • = p N = 1 N .

= N -K N -1 4 K E 2 G 2 .

Then the same bound in Theorem 2 holds.

Scheme II requires p:

1 = • • • = p N = 1

N which obviously violates the unbalance nature of FL. Fortunately, this can be addressed by the following transformation. Let F k (w) = p k N F k (w) be a scaled local objective F k . Then the global objective becomes a simple average of all scaled local objectives:

F (w) = N k=1 p k F k (w) = 1 N N k=1 F k (w). Theorem 3 still holds if L, µ, σ k , G are replaced by L νL, µ ςµ, σ k = √ νσ, and G = √ νG, respectively. Here, ν = N • max k p k and ς = N • min k p k . 3.4 DISCUSSIONS Choice of E. Since w 0 -w * 2 ≤ 4

µ 2 G 2 for µ-strongly convex F , the dominating term in eqn. ( 6) is

EQUATION

Let T denote the number of required steps for FedAvg to achieve an accuracy. It follows from eqn. ( 7) that the number of required communication rounds is roughly

EQUATION

Thus, T E is a function of E that first decreases and then increases, which implies that over-small or over-large E may lead to high communication cost and that the optimal E exists. Stich (2018) showed that if the data are iid, then E can be set to O( √ T ). However, this setting does not work if the data are non-iid. Theorem 1 implies that E must not exceed Ω( √ T ); otherwise, convergence is not guaranteed. Here we give an intuitive explanation. If E is set big, then w k t can converge to the minimizer of F k , and thus FedAvg becomes the one-shot average Zhang et al. (2013) of the local solutions. If the data are non-iid, the one-shot averaging does not work because weighted average of the minimizers of F 1 , • • • , F N can be very different from the minimizer of F .

Choice of K. Stich (2018) showed that if the data are iid, the convergence rate improves substantially as K increases. However, under the non-iid setting, the convergence rate has a weak dependence on K, as we show in Theorems 2 and 3. This implies FedAvg is unable to achieve linear speedup. We have empirically observed this phenomenon (see Section 6). Thus, in practice, the participation ratio K N can be set small to alleviate the straggler's effect without affecting the convergence rate.

Choice of sampling schemes. We considered two sampling and averaging schemes in Theorems 2 and 3. Scheme I selects K devices according to the probabilities p 1 , • • • , p N with replacement. The non-uniform sampling results in faster convergence than uniform sampling, especially when p 1 , • • • , p N are highly non-uniform. If the system can choose to activate any of the N devices at any time, then Scheme I should be used.

However, oftentimes the system has no control over the sampling; instead, the server simply uses the first K returned results for the update. In this case, we can assume the K devices are uniformly sampled from all the N devices and use Theorem 3 to guarantee the convergence. If p 1 , • • • , p N are highly non-uniform, then ν = N • max k p k is big and ς = N • min k p k is small, which makes the convergence of FedAvg slow. This point of view is empirically verified in our experiments.

Theorem 4. We artificially construct a strongly convex and smooth distributed optimization problem.

With full batch size, E > 1, and any fixed step size, FedAvg will converge to sub-optimal points. Specifically, let w * be the solution produced by FedAvg with a small enough and constant η, and w * the optimal solution. Then we have

w * -w * 2 = Ω((E -1)η) • w * 2 .

where we hide some problem dependent constants.

Theorem 4 and its proof provide several implications. First, the decay of learning rate is necessary of FedAvg. On the one hand, Theorem 1 shows with E > 1 and a decaying learning rate, FedAvg converges to the optimum. On the other hand, Theorem 4 shows that with E > 1 and any fixed learning rate, FedAvg does not converges to the optimum.

Second, FedAvg behaves very differently from gradient descent. Note that FedAvg with E = 1 and full batch size is exactly the Full Gradient Descent; with a proper and fixed learning rate, its global convergence to the optimum is guaranteed Nesterov (2013). However, Theorem 4 shows that FedAvg with E > 1 and full batch size cannot possibly converge to the optimum. This conclusion doesn't contradict with Theorem 1 in Khaled et al. ( 2019), which, when translated into our case, asserts that w * will locate in the neighborhood of w * with a constant learning rate.

Third, Theorem 4 shows the requirement of learning rate decay is not an artifact of our analysis; instead, it is inherently required by FedAvg. An explanation is that constant learning rates, combined with E steps of possibly-biased local updates, form a sub-optimal update scheme, but a diminishing learning rate can gradually eliminate such bias.

The efficiency of FedAvg principally results from the fact that it performs several update steps on a local model before communicating with other workers, which saves communication. Diminishing step sizes often hinders fast convergence, which may counteract the benefit of performing multiple local updates. Theorem 4 motivates more efficient alternatives to FedAvg. Experiment settings For all experiments, we initialize all runnings with w 0 = 0. In each round, all selected devices run E steps of SGD in parallel. We decay the learning rate at the end of each round by the following scheme η t = η0 1+t , where η 0 is chosen from the set {1, 0.1, 0.01}. We evaluate the averaged model after each global synchronization on the corresponding global objective. For fair comparison, we control all randomness in experiments so that the set of activated devices is the same across all different algorithms on one configuration.

5. RELATED WORK:

Impact of E We expect that T /E, the required communication round to achieve curtain accuracy, is a hyperbolic finction of E as equ (8) indicates. Intuitively, a small E means a heavy communication burden, while a large E means a low convergence rate. One needs to trade off between communication efficiency and fast convergence. We empirically observe this phenomenon on unbalanced datasets in Figure 1a . The reason why the phenomenon does not appear in mnist balanced dataset requires future investigations.

Published as a conference paper at ICLR 2020 Impact of K Our theory suggests that a larger K may slightly accelerate convergence since T /E contains a term O EG 2 K . Figure 1b shows that K has limited influence on the convergence of FedAvg in synthetic(0,0) dataset. It reveals that the curve of a large enough K is slightly better. We observe similar phenomenon among the other three datasets and attach additional results in Appendix D. This justifies that when the variance resulting sampling is not too large (i.e., B C), one can use a small number of devices without severely harming the training process, which also removes the need to sample as many devices as possible in convex federated optimization.

Effect of sampling and averaging schemes. We compare four schemes among four federated datasets. Since the original scheme involves a history term and may be conservative, we carefully set the initial learning rate for it. Figure 1c indicates that when data are balanced, Schemes I and II achieve nearly the same performance, both better than the original scheme. Figure 1d shows that when the data are unbalanced, i.e., p k 's are uneven, Scheme I performs the best. Scheme II suffers from some instability in this case. This is not contradictory with our theory since we don't guarantee the convergence of Scheme II when data is unbalanced. As expected, transformed Scheme II performs stably at the price of a lower convergence rate. Compared to Scheme I, the original scheme converges at a slower speed even if its learning rate is fine tuned. All the results show the crucial position of appropriate sampling and averaging schemes for FedAvg.

7. CONCLUSION:

Federated learning becomes increasingly popular in machine learning and optimization communities. In this paper we have studied the convergence of FedAvg, a heuristic algorithm suitable for federated setting. We have investigated the influence of sampling and averaging schemes. We have provided theoretical guarantees for two schemes and empirically demonstrated their performances. Our work sheds light on theoretical understanding of FedAvg and provides insights for algorithm design in realistic applications. Though our analyses are constrained in convex problems, we hope our insights and proof techniques can inspire future work.

Yue Zhao, Meng 

= {nE | n = 1, 2, • • • }. If t + 1 ∈ I E , i.e.

, the time step to communication, FedAvg activates all devices. Then the update of FedAvg with partial devices active can be described as

EQUATION

EQUATION

Here, an additional variable v k t+1 is introduced to represent the immediate result of one step SGD update from w k t . We interpret w k t+1 as the parameter obtained after communication steps (if possible).

In our analysis, we define two virtual sequences v t = N k=1 p k v k t and w t = N k=1 p k w k t . This is motivated by (Stich, 2018). v t+1 results from an single step of SGD from w t . When t + 1 / ∈ I E , both are inaccessible. When t + 1 ∈ I E , we can only fetch w t+1 . For convenience, we define

g t = N k=1 p k ∇F k (w k t ) and g t = N k=1 p k ∇F k (w k t , ξ k t )

. Therefore, v t+1 = w t -η t g t and Eg t = g t .

A.2 KEY LEMMAS:

To convey our proof clearly, it would be necessary to prove certain useful lemmas. We defer the proof of these lemmas to latter section and focus on proving the main theorem. Lemma 1 (Results of one step SGD). Assume Assumption 1 and 2. If η t ≤ 1 4L , we have

E v t+1 -w 2 ≤ (1 -η t µ)E w t -w 2 + η 2 t E g t -g t 2 + 6Lη 2 t Γ + 2E N k=1 p k w t -w t k 2 where Γ = F * - N k=1 p k F k ≥ 0.

Lemma 2 (Bounding the variance). Assume Assumption 3 holds. It follows that

E g t -g t 2 ≤ N k=1 p 2 k σ 2 k .

Lemma 3 (Bounding the divergence of {w k t }). Assume Assumption 4, that η t is non-increasing and η t ≤ 2η t+E for all t ≥ 0. It follows that

E N k=1 p k w t -w t k 2 ≤ 4η 2 t (E -1) 2 G 2 .

A.3 COMPLETING THE PROOF OF THEOREM 1:

Proof. It is clear that no matter whether t + 1 ∈ I E or t + 1 / ∈ I E , we always have w t+1 = v t+1 . Let ∆ t = E w t+1 -w 2 . From Lemma 1, Lemma 2 and Lemma 3, it follows that

EQUATION

where

B = N k=1 p 2 k σ 2 k + 6LΓ + 8(E -1) 2 G 2 .

For a diminishing stepsize, η t = β t+γ for some β > 1 µ and γ > 0 such that

η 1 ≤ min{ 1 µ , 1 4L } = 1 4L

and η t ≤ 2η t+E . We will prove

∆ t ≤ v γ+t where v = max β 2 B βµ-1 , (γ + 1)∆ 1 .

We prove it by induction. Firstly, the definition of v ensures that it holds for t = 1. Assume the conclusion holds for some t, it follows that

∆ t+1 ≤ (1 -η t µ)∆ t + η 2 t B = 1 - βµ t + γ v t + γ + β 2 B (t + γ) 2 = t + γ -1 (t + γ) 2 v + β 2 B (t + γ) 2 - βµ -1 (t + γ) 2 v ≤ v t + γ + 1 .

Then by the strong convexity of F (•),

E[F (w t )] -F * ≤ L 2 ∆ t ≤ L 2 v γ + t .

Specifically, if we choose β = 2 µ , γ = max{8 L µ -1, E} and denote κ = L µ , then η t = 2 µ 1 γ+t and

E[F (w t )] -F * ≤ 2κ γ + t B µ + 2L∆ 1 .

A.4 DEFERRED PROOFS OF KEY LEMMAS:

Proof of Lemma 1. Notice that v t+1 = w t -η t g t , then

EQUATION

Note that EA 2 = 0. We next focus on bounding A 1 . Again we split A 1 into three terms:

EQUATION

From the the L-smoothness of F k (•), it follows that

EQUATION

)

By the convexity of • 2 and eqn. ( 14), we have

B 2 = η 2 t g t 2 ≤ η 2 t N k=1 p k ∇F k w k t 2 ≤ 2Lη 2 t N k=1 p k F k (w k t ) -F * k .

Published as a conference paper at ICLR 2020 Note that

EQUATION

By Cauchy-Schwarz inequality and AM-GM inequality, we have

EQUATION

)

By the µ-strong convexity of F k (•), we have

EQUATION

)

By combining eqn. (13), eqn. (15), eqn. ( 16) and eqn. ( 17), it follows that

A 1 = w t -w -η t g t 2 ≤ w t -w 2 + 2Lη 2 t N k=1 p k F k (w k t ) -F * k + η t N k=1 p k 1 η t w t -w t k 2 + η t ∇F k w k t 2 -2η t N k=1 p k F k w k t -F k (w * ) + µ 2 w k t -w 2 = (1 -µη t ) w t -w 2 + N k=1 p k w t -w t k 2 + 4Lη 2 t N k=1 p k F k (w k t ) -F * k -2η t N k=1 p k F k w k t -F k (w * ) C

where we use eqn. ( 14) again.

We next aim to bound C. We define

γ t = 2η t (1 -2Lη t ). Since η t ≤ 1 4L , η t ≤ γ t ≤ 2η t .

Then we split C into two terms:

C = -2η t (1 -2Lη t ) N k=1 p k F k (w k t ) -F * k + 2η t N k=1 p k (F k (w * ) -F * k ) = -γ t N k=1 p k F k (w k t ) -F * + (2η t -γ t ) N k=1 p k (F * -F * k ) = -γ t N k=1 p k F k (w k t ) -F * D +4Lη 2 t Γ

where in the last equation, we use the notation

Γ = N k=1 p k (F * -F * k ) = F * - N k=1 p k F * k .

To bound D, we have

N k=1 p k F k (w k t ) -F * = N k=1 p k F k (w k t ) -F k (w t ) + N k=1 p k (F k (w t ) -F * ) ≥ N k=1 p k ∇F k (w t ), w k t -w t + (F (w t ) -F * ) ≥ - 1 2 N k=1 p k η t ∇F k (w t ) 2 + 1 η t w k t -w t 2 + (F (w t ) -F * ) ≥ - N k=1 p k η t L (F k (w t ) -F * k ) + 1 2η t w k t -w t 2 + (F (w t ) -F * )

where the first inequality results from the convexity of F k (•), the second inequality from AM-GM inequality and the third inequality from eqn. ( 14).

Therefore

C = γ t N k=1 p k η t L (F k (w t ) -F * k ) + 1 2η t w k t -w t 2 -γ t (F (w t ) -F * ) + 4Lη 2 t Γ = γ t (η t L -1) N k=1 p k (F k (w t ) -F * ) + 4Lη 2 t + γ t η t L Γ + γ t 2η t N k=1 p k w k t -w t 2 ≤ 6Lη 2 t Γ + N k=1 p k w k t -w t 2

where in the last inequality, we use the following facts: (1)

η t L -1 ≤ -3 4 ≤ 0 and N k=1 p k (F k (w t ) -F * ) = F (w t ) -F * ≥ 0 (2) Γ ≥ 0 and 4Lη 2 t + γ t η t L ≤ 6η 2 t L and (3) γt 2ηt ≤ 1.

Recalling the expression of A 1 and plugging C into it, we have

EQUATION

Using eqn. (18) and taking expectation on both sides of eqn. (12), we erase the randomness from stochastic gradients, we complete the proof.

Proof of Lemma 2. From Assumption 3, the variance of the stochastic gradients in device k is bounded by σ 2 k , then

E g t -g t 2 = E N k=1 p k (∇F k (w k t , ξ k t ) -∇F k (w k t )) 2 , = N k=1 p 2 k E ∇F k (w k t , ξ k t ) -∇F k (w k t ) 2 , ≤ N k=1 p 2 k σ 2 k .

Proof of Lemma 3. Since FedAvg requires a communication each E steps. Therefore, for any t ≥ 0, there exists a t 0 ≤ t, such that t -t 0 ≤ E -1 and w k t0 = w t0 for all k = 1, 2, • • • , N . Also, we use the fact that η t is non-increasing and η t0 ≤ 2η t for all t -t 0 ≤ E -1, then

E N k=1 p k w t -w k t 2 = E N k=1 p k (w k t -w t0 ) -(w t -w t0 ) 2 ≤ E N k=1 p k w k t -w t0 2 ≤ N k=1 p k E t-1 t=t0 (E -1)η 2 t ∇F k (w k t , ξ k t ) 2 ≤ N k=1 p k t-1 t=t0 (E -1)η 2 t0 G 2 ≤ N k=1 p k η 2 t0 (E -1) 2 G 2 ≤ 4η 2 t (E -1) 2 G 2 .

B PROOFS OF THEOREMS 2 AND 3:

We analyze FedAvg in the setting of partial device participation in this section.

B.1 ADDITIONAL NOTATION:

Recall that w k t is the model parameter maintained in the k-th device at the t-th step.

I E = {nE | n = 1, 2, • • • }

is the set of global synchronization steps. Unlike the setting in Appendix A, when it is the time to communicate, i.e., t + 1 ∈ I E , the scenario considered here is that FedAvg randomly activates a subset of devices according to some sampling schemes. Again,

g t = N k=1 p k ∇F k (w k t ) and g t = N k=1 p k F k (w k t , ξ k t )

. Therefore, v t+1 = w t -η t g t and Eg t = g t .

Multiset selected. All sampling schemes can be divided into two groups, one with replacement and the other without replacement. For those with replacement, it is possible for a device to be activated several times in a round of communication, even though each activation is independent with the rest. We denote by H t the multiset selected which allows any element to appear more than once. Note that H t is only well defined for t ∈ I E . For convenience, we denote by S t = H N (t,E) the most recent set of chosen devices where N (t, E) = max{n|n ≤ t, n ∈ I E }.

Updating scheme. Limited to realistic scenarios (for communication efficiency and low straggler effect), FedAvg first samples a random multiset S t of devices and then only perform updates on them. This make the analysis a little bit intricate, since S t varies each E steps. However, we can use a thought trick to circumvent this difficulty. We assume that FedAvg always activates all devices at the beginning of each round and then uses the parameters maintained in only a few sampled devices to produce the next-round parameter. It is clear that this updating scheme is equivalent to the original.

Then the update of FedAvg with partial devices active can be described as: for all k ∈ [N ],

EQUATION

EQUATION

Sources of randomness. In our analysis, there are two sources of randomness. One results from the stochastic gradients and the other is from the random sampling of devices. All the analysis in Appendix A only involve the former. To distinguish them, we use the notation E St (•), when we take expectation to erase the latter type of randomness.

B.2 KEY LEMMAS:

Two schemes. For full device participation, we always have w t+1 = v t+1 . This is true when t + 1 / ∈ I E for partial device participation. When t+1 ∈ I E , we hope this relation establish in the sense of expectation. To that end, we require the sampling and averaging scheme to be unbiased in the sense that

E St+1 w t+1 = v t+1 .

We find two sampling and averaging schemes satisfying the requirement and provide convergence guarantees.

(I) The server establishes S t+1 by i.i.d. with replacement

sampling an index k ∈ {1, • • • , N } with probabilities p 1 , • • • , p N for K times.

Hence S t+1 is a multiset which allows a element to occur more than once. Then the server averages the parameters by

w k t+1 = 1 K k∈St+1 v k t+1

. This is first proposed in (Sahu et al., 2018) but lacks theoretical analysis.

(II) The server samples S t+1 uniformly in a without replacement fashion. Hence each element in S t+1 only occurs once.Then server averages the parameters by w k t+1 = k∈St+1 p k N K v k t+1 . Note that when the p k 's are not all the same, one cannot ensure k∈St+1 p k N K = 1.

Unbiasedness and bounded variance. Lemma 4 shows the mentioned two sampling and averaging schemes are unbiased. In expectation, the next-round parameter (i.e., w t+1 ) is equal to the weighted average of parameters in all devices after SGD updates (i.e., v t+1 ). However, the original scheme in (McMahan et al., 2017) (see Table 1 ) does not enjoy this property. But it is very similar to Scheme II except the averaging scheme. Hence our analysis cannot cover the original scheme.

Lemma 5 shows the expected difference between v t+1 and w t+1 is bounded.

E St v t+1 -w t+1 2

is actually the variance of w t+1 .

Lemma 4 (Unbiased sampling scheme). If t + 1 ∈ I E , for Scheme I and Scheme II, we have

E St (w t+1 ) = v t+1 .

Lemma 5 (Bounding the variance of w t ). For t + 1 ∈ I, assume that η t is non-increasing and η t ≤ 2η t+E for all t ≥ 0. We have the following results.

(1) For Scheme I, the expected difference between v t+1 and w t+1 is bounded by

E St v t+1 -w t+1 2 ≤ 4 K η 2 t E 2 G 2 .

(2) For Scheme II, assuming

p 1 = p 2 = • • • = p N = 1

N , the expected difference between v t+1 and w t+1 is bounded by

E St v t+1 -w t+1 2 ≤ N -K N -1 4 K η 2 t E 2 G 2 .

B.3 COMPLETING THE PROOF OF THEOREM 2 AND 3:

Proof. Note that

w t+1 -w * 2 = w t+1 -v t+1 + v t+1 -w * 2 = w t+1 -v t+1 2 A1 + v t+1 -w * 2 A2 + 2 w t+1 -v t+1 , v t+1 -w * A3 .

When expectation is taken over S t+1 , the last term (A 3 ) vanishes due to the unbiasedness of w t+1 .

If t + 1 / ∈ I E , A 1 vanishes since w t+1 = v t+1 . We use Lemma 5 to bound A 2 . Then it follows that

E w t+1 -w * 2 ≤ (1 -η t µ)E w t -w 2 + η 2 t B.

If t + 1 ∈ I E , we additionally use Lemma 5 to bound A 1 . Then

EQUATION

where C is the upper bound of 1

η 2 t E St v t+1 -w t+1 2 (C

is defined in Theorem 2 and 3).

The only difference between eqn. (21) and eqn. ( 11) is the additional C. Thus we can use the same argument there to prove the theorems here. Specifically, for a diminishing stepsize, η t = β t+γ for some β > 1 µ and γ > 0 such that

η 1 ≤ min{ 1 µ , 1 4L } = 1 4L and η t ≤ 2η t+E , we can prove E w t+1 -w * 2 ≤ v γ+t where v = max β 2 (B+C) βµ-1 , (γ + 1) w 1 -w * 2 .

Then by the strong convexity of F (•),

E[F (w t )] -F * ≤ L 2 ∆ t ≤ L 2 v γ + t . Specifically, if we choose β = 2 µ , γ = max{8 L µ -1, E} and denote κ = L µ , then η t = 2 µ 1 γ+t and E[F (w t )] -F * ≤ 2κ γ + t B + C µ + 2L w 1 -w * 2 .

B.4 DEFERRED PROOFS OF KEY LEMMAS:

Proof of Lemma 4. We first give a key observation which is useful to prove the followings. Let {x i } N i=1 denote any fixed deterministic sequence. We sample a multiset S t (with size K) by the procedure where for each sampling time, we sample x k with probability q k for each time. Pay attention that two samples are not necessarily independent. We only require each sampling distribution is identically. Let S t = {i 1 , • • • , i K } ⊂ [N ] (some i k 's may have the same value). Then

E St k∈St x k = E St K k=1 x i k = KE St x i1 = K N k=1 q k x k .

For Scheme I, q k = p k and for Scheme II, q k = 1 N . It is easy to prove this lemma when equipped with this observation.

Proof of Lemma 5. We separately prove the bounded variance for two schemes. Let S t+1 = {i 1 , • • • , i K } denote the multiset of chosen indexes.

(1) For Scheme I,

w t+1 = 1 K K l=1 v i l t+1 .

Taking expectation over S t+1 , we have

EQUATION

)

where the first equality follows from v i l t+1 are independent and unbiased. To bound eqn. (22), we use the same argument in Lemma 5. Since t + 1 ∈ I E , we know that the time t 0 = t -E + 1 ∈ I E is the communication time, which implies {w k t0 } N k=1 is identical. Then

N k=1 p k v k t+1 -v t+1 2 = N k=1 p k (v k t+1 -w t0 ) -(v t+1 -w t0 ) 2 ≤ N k=1 p k v k t+1 -w t0

where the last inequality results from

N k=1 p k (v k t+1 -w t0 ) = v t+1 -w t0 and E x -Ex 2 ≤ E x 2 . Similarly, we have E St w t+1 -v t+1 2 ≤ 1 K N k=1 p k E v k t+1 -w t0 2 ≤ 1 K N k=1 p k E v k t+1 -w k t0 2 ≤ 1 K N k=1 p k E t i=t0 E η i ∇F k (w k i , ξ k i ) 2 ≤ 1 K E 2 η 2 t0 G 2 ≤ 4 K η 2 t E 2 G 2

where in the last inequality we use the fact that η t is non-increasing and η t0 ≤ 2η t .

(2) For Scheme II, when assuming

p 1 = p 2 = • • • = p N = 1 N , we again have w t+1 = 1 K K l=1 v i l t+1 . E St w t+1 -v t+1 2 = E St 1 K i∈St+1 v i t+1 -v t+1 2 = 1 K 2 E St N i=1 I {i ∈ S t } (v i t+1 -v t+1 ) 2 = 1 K 2   i∈[N ] P (i ∈ S t+1 ) v i t+1 -v t+1 2 + i =j P (i, j ∈ S t+1 ) v i t+1 -v t+1 , v j t+1 -v t+1   = 1 KN N i=1 v i t+1 -v t+1 2 + i =j K -1 KN (N -1) v i t+1 -v t+1 , v j t+1 -v t+1 = 1 K(N -1) 1 - K N N i=1 v i t+1 -v t+1 2

where we use the following equalities: (1) P (i ∈ S t+1 ) = K N and P (i, j ∈ S t+1 ) = K(K-1) N (N -1) for all i = j and (2)

i∈[N ] v i t+1 -v t+1 2 + i =j v i t+1 -v t+1 , v j t+1 -v t+1 = 0. Therefore, E w t+1 -v t+1 2 = N K(N -1) 1 - K N E 1 N N i=1 v i t+1 -v t+1 2 ≤ N K(N -1) 1 - K N E 1 N N i=1 v i t+1 -w t0 2 ≤ N K(N -1) 1 - K N 4η 2 t E 2 G 2 .

where in the last inequality we use the same argument in (1) .

C THE EMPIRICAL RISK MINIMIZATION EXAMPLE IN SECTION 4 C.1 DETAIL OF THE EXAMPLE:

Let p > 1 be a positive integer. To avoid the trivial case, we assume N > 1. Consider the following quadratic optimization

EQUATION

where A ∈ R (N p+1)×(N p+1) , w, b ∈ R N p+1 and µ > 0. Specifically, let b = e 1 (1, 0, • • • , 0) , and A be a symmetric and tri-diagonal matrix defined by

EQUATION

where i, j are row and column indices, respectively. We partition A into a sum of N symmetric matrices (A = 

EQUATION

Now A k 's are given by

A 1 = B 1 +E 1,1 , A k = B k (2 ≤ k ≤ N -1) and A N = B N +E N p+1,N p+1

, where E i,j is the matrix where only the (i, j)th entry is one and the rest are zero.

Back to the federated setting, we distribute the k-th partition (A k , b k ) to the k-th device and construct its corresponding local objective by

EQUATION

In the next subsection (Appendix C.3), we show that the quadratic minimization with the global objective (23) and the local objectives ( 26) is actually a distributed linear regression. In this example, training data are not identically but balanced distributed. Moreover, data in each device are sparse in the sense that non-zero features only occur in one block. The following theorem (Theorem 5) shows that FedAvg might converge to sub-optimal points even if the learning rate is small enough. We provide a numerical illustration in Appendix C.2 and a mathematical proof in Appendix C.4.

Theorem 5. In the above problem of the distributed linear regression, assume that each device computes exact gradients (which are not stochastic). With a constant and small enough learning rate η and E > 1, FedAvg converges to a sub-optimal solution, whereas FedAvg with E = 1 (i.e., gradient descent) converges to the optimum. Specifically, in a quantitative way, we have

w * -w * ≥ (E -1)η 16 A 1 A 2 w *

where w * is the solution produced by FedAvg and w * is the optimal solution.

C.2 NUMERICAL ILLUSTRATION ON THE EXAMPLE:

We conduct a few numerical experiments to illustrate the poor performance of FedAvg on the example introduced in Section 4. Here we set N = 5, p = 4, µ = 2 × 10 -4 . The annealing scheme of learning rates is given by η t = 1/5 5+t•a where a is the best parameter chosen from the set {10 -2 , 10 -4 , 10 -6 }.

C.3 SOME PROPERTIES OF THE EXAMPLE:

Recall that the symmetric matrix A ∈ R (N p+1)×(N p+1) is defined in eqn. (24). Observe that A is invertible and for all vector w ∈ R N p+1 ,

EQUATION

)

which implies that 0 ≺ A 4I. The sparse and symmetric matrices B k (1 ≤ k ≤ N ) defined in eqn. (25) can be rewritten as

(B k ) =            0 (k-1)p×(k-1)p        1 -1 -1 2 -1 -1 . . . . . . . . . 2 -1 -1 1        (p+1)×(p+1) 0 (N -k)p×(N -k)p           

.

From theory of linear algebra, it is easy to follow this proposition. Proposition 1. By the way of construction, A k 's have following properties:

1. A k is positive semidefinite with A k 2 ≤ 4; 2. rank(A 2 ) = • • • = rank(A N -1 ) = p and rank(A 1 ) = rank(A N ) = p + 1; 3. For each k, there exist a matrix X k ∈ R r k ×(N p+1) such that A k = X k X k where r k = rank(A k ).

Given any k, each row of X k has non-zero entries only on a block of coordinates, namely

I k = {(k -1)p + 1, (k -1)p + 2, • • • , kp + 1}. As a result, A = N k=1 A k = X X, where X = (X 1 , • • • , X N ) ∈ R (N p+2)×(N p+1) . 4. w * = A -1 b is the global minimizer of problem eqn. (23), given by (w * ) i = 1-i N p+2 (1 ≤ i ≤ N p + 1). Let w (1, • • • , 1 p+1 , 0, • • • , 0 (N -1)p ) ∈ R N p+1 , then A 1 w = X 1 X 1 w = b 1 .

From Proposition 1, we can rewrite these local quadratic objectives in form of a ridge linear regression. Specifically, for k = 1,

F 1 (w) = 1 2 w A 1 w -2b 1 w + µ w 2 , = 1 2 w X 1 X 1 w -2 w X 1 X 1 w + µ w 2 , = 1 2 X 1 (w -w) 2 2 + 1 2 µ w 2 + C,

Second from eqn. (31), we have that

I - 1 N N i=1 (I -ηA i ) E -1 = Eη N A + Θ(η 2 ) -1 = N Eη A -1 + Θ(1).

Plugging the last two equations into eqn. (29), we have

w * -w * = N Eη A -1 + Θ(1) η 2 N E(E -1) 2 (V -A 1 A) + Θ(η) w * = E -1 2 ηA -1 (V -A 1 A) + Θ(η) w * ≥ (E -1)η 16 (V -A 1 A)w * = (E -1)η 16 A 1 A 2 w *

where the last inequality holds because (i) we require η to be sufficiently small and (ii) A -1 x ≥ 1 4 x for any vector x as a result of 0 < A ≤ 4. The last equality uses the fact (i)

V -A 1 A = A 1 n i=2 A i and (ii) A 1 A i = 0 for any i ≥ 3.

D EXPERIMENTAL DETAILS D.1 EXPERIMENTAL SETTING:

Model and loss. We examine our theoretical results on a multinomial logistic regression. Specifically, let f (w; x i ) denote the prediction model with the parameter w = (W, b) and the form f (w; x i ) = softmax(Wx i + b). The loss function is given by

F (w) = 1 n n i=1 CrossEntropy (f (w; x i ), y i ) + λ w 2 2 .

This is a convex optimization problem. The regularization parameter is set to λ = 10 -4 .

Datasets. We evaluate our theoretical results on both real data and synthetic data. For real data, we choose MNIST dataset (LeCun et al., 1998) because of its wide academic use. To impose statistical heterogeneity, we distribute the data among N = 100 devices such that each device contains samples of only two digits. To explore the effect of data unbalance, we further vary the number of samples among devices. Specifically, for unbalanced cases, the number of samples among devices follows a power law, while for balanced cases, we force all devices to have the same amount of samples. Here α and β allow for more precise manipulation of data heterogeneity: α controls how much local models differ from each other and β controls how much the local data at each device differs from that of other devices. There are N = 100 devices in total. The number of samples n k in each device follows a power law, i.e., data are distributed in an unbalanced way. We denote by synthetic(α, β) the synthetic dataset with parameter α and β.

We summarize the information of federated datasets in Table 2 .

Experiments. For all experiments, we initialize all runnings with w 0 = 0. In each round, all selected devices run E steps of SGD in parallel. We decay the learning rate at the end of each round by the following scheme η t = η0 1+t , where η 0 is chosen from the set {1, 0.1, 0.01}. We evaluate the averaged model after each global synchronization on the corresponding global objective. For fair comparison, we control all randomness in experiments so that the set of activated devices is the same across all different algorithms on one configuration. The impact of E. From our theory, when the total steps T is sufficiently large, the required number of communication rounds to achieve a certain precision is

T /E ≈ O EG 2 K + EG 2 + N k=1 p 2 k σ 2 + LΓ + κG 2 E ,

which is s a function of E that first decreases and then increases. This implies that the optimal local step E * exists. What's more, the T /E evaluated at

E * is O   G N k=1 p 2 k σ 2 + LΓ + κG 2   ,

which implies that FedAvg needs more communication rounds to tackle with severer heterogeneity.

To validate these observations, we test FedAvg with Scheme I on our four datasets as listed in Table 2 . In each round, we activate K = 30 devices and set η 0 = 0.1 for all experiments in this part.

For unbalanced MNIST, we use batch size b = 64. The target loss value is 0.29 and the minimum loss value found is 0.2591. For balanced MNIST, we also use batch size b = 64. The target loss value is 0.50 and the minimum loss value found is 0.3429. For two synthetic datasets, we choose b = 24. The target loss value for synthetic(0,0) is 0.95 and the minimum loss value is 0.7999. Those for synthetic(1,1) are 1.15 and 1.075.

The impact of K. Our theory suggests that a larger K may accelerate convergence since T /E contains a term O EG 2 K . We fix E = 5 and η 0 = 0.1 for all experiments in this part. We set the batch size to 64 for two MNIST datasets and 24 for two synthetic datasets. We test Scheme I for illustration. Our results show that, no matter what value K is, FedAvg converges. From Figure 3 , all the curves in each subfigure overlap a lot. To show more clearly the differences between the curves, we zoom in the last few rounds in the upper left corner of the figure. It reveals that the curve of a large enough K is slightly better. This result also shows that there is no need to sample as many devices as possible in convex federated optimization.

Sampling and averaging schemes. We analyze the influence of sampling and averaging schemes. As stated in Section 3.3, Scheme I iid samples (with replacement) K indices with weights p k and simply averages the models, which is proposed by Sahu et al. (2018) . Scheme II uniformly samples (without replacement) K devices and weightedly averages the models with scaling factor N/K. Transformed Scheme II scales each local objective and uses uniform sampling and simple averaging. We compare Scheme I, Scheme II and transformed Scheme II, as well as the original scheme (McMahan et al., 2017) on four datasets. We carefully tuned the learning rate for the original scheme. In particular, we choose the best step size from the set {0.1, 0.5, 0.9, 1.1}. We did not fine tune the rest schemes and set η 0 = 0.1 by default. The hyperparameters are the same for all schemes: E = 20, K = 10 and b = 64. The results are shown in Figure 1c and 1d .

Our theory renders Scheme I the guarantee of convergence in common federated setting. As expected, Scheme I performs well and stably across most experiments. This also coincides with the findings of Sahu et al. (2018) . They noticed that Scheme I performs slightly better than another scheme: server

:

In original paper (McMahan et al.,

2017), E epochs of SGD are performed in parallel. For theoretical analyses, we denote by E the times of updates rather than epochs.2 Throughout this paper, "non-iid" means data are not identically distributed. More precisely, the data distributions in the k-th and l-th devices, denote D k and D l , can be different.3 Throughout this paper, "sampling" refers to how the server chooses K user devices and use their outputs for updating the model parameters. "Sampling" does not mean how a device randomly selects training samples.

It is well know that the full gradient descent (which is equivalent to FedAvg with E = 1 and full batch) do not require the decay of learning rate.

The sampling scheme is proposed by Sahu et al.(2018) for FedAvg as a baseline, but this convergence rate is our contribution.

The convergence relies on the assumption that data are balanced, i.e., n1 = n2 = • • • = nN . However, we can use a rescaling trick to get rid of this assumption. We will discuss this point later in Section 3.

NECESSITY OF LEARNING RATE DECAYIn this section, we point out that diminishing learning rates are crucial for the convergence of FedAvg in the non-iid setting. Specifically, we establish the following theorem by constructing a ridge regression model (which is strongly convex and smooth).

Figures and Tables:

et al. (2019); Yu et al. (2019); Wang et al. (2019) made the latter assumption, while Zhou and Cong (2017); Stich (2018); Wang and Joshi (2018); Woodworth et al. (2018) made both assumptions. The two assumptions violates the second and third characters of FL. Previous algorithm Fedprox Sahu et al. (

is the learning rate (a.k.a. step size) and ξ k t+i is a sample uniformly chosen from the local data. Last, the server aggregates the local models, w 1 t+E , • • • , w N t+E , to produce the new global model, w t+E . Because of the non-iid and partial device participation issues, the aggregation step can vary. IID versus non-iid. Suppose the data in the k-th device are i.i.d. sampled from the distribution D k . Then the overall distribution is a mixture of all local data distributions: D = N k=1 p k D k . The prior work Zhang et al. (2015a); Zhou and Cong (2017); Stich (2018); Wang and Joshi (2018); Woodworth et al. (2018) assumes the data are iid generated by or partitioned among the N devices, that is, D k = D for all k ∈ [N ]. However, real-world applications do not typically satisfy the iid assumption. One of our theoretical contributions is avoiding making the iid assumption. Full device participation. The prior work Coppola (2015); Zhou and Cong (2017); Stich (2018); Yu et al. (2019); Wang and Joshi (2018); Wang et al. (2019) requires the full device participation in the aggregation step of FedAvg. In this case, the aggregation step performs

Figure1: (a) To obtain an accuracy, the required rounds first decrease and then increase when we increase the local steps E. (b) In Synthetic(0,0) dataset, decreasing the numbers of active devices each round has little effect on the convergence process. (c) In mnist balanced dataset, Scheme I slightly outperforms Scheme II. They both performs better than the original scheme.Here transformed Scheme II coincides with Scheme II due to the balanced data. (d) In mnist unbalanced dataset, Scheme I performs better than Scheme II and the original scheme. Scheme II suffers from instability while transformed Scheme II has a lower convergence rate.

k=1 A k ) and b into b = N k=1 b k . Specifically, we choose b 1 = b = e 1 and b 2 = • • • = b N = 0. To give the formulation of A k 's, we first introduce a series of sparse and symmetric matrices B k (1 ≤ k ≤ N ):

Figure2: The left figure shows that the global objective value that FedAvg converges to is not optimal unless E = 1. Once we decay the learning rate, FedAvg can converge to the optimal even if E > 1.

Synthetic data allow us to manipulate heterogeneity more precisely. Here we follow the same setup as described in(Sahu et al., 2018). In particular, we generate synthetic samples(X k , Y k ) according to the model y = argmax(softmax(W k x + b k )) with x ∈ R 60 , W k ∈ R 10×60 and b k ∈ R 10 , where X k ∈ R n k ×60 and Y k ∈ R n k . We model each entry of W k and b k as N (µ k , 1) with µ k ∼ N (0, α), and (x k ) j ∼ N (v k , 1 j 1.2 ) with v k ∼ N (B k , 1) and B k ∼ N (0, β).

al., 2014a;b; Shamir et al., 2014; Zhang and Lin, 2015; Meng et al., 2016; Reddi et al., 2016; Richtárik and Takác, 2016; Smith et al., 2016; Zheng et al., 2016; Shusen

Table content: <table/>

Sampling and averaging schemes for FedAvg. S t ∼ U(N, K) means S t is a size-K subset uniformly sampled without replacement from [N ]. S t ∼ W(N, K, p) means S t contains K elements that are iid sampled with replacement from [N ] with probabilities {p k }. In the latter scheme, S t is not a set.

Table content: <table><tr><td>Paper</td><td>Sampling</td><td>Averaging</td><td>Convergence rate</td></tr><tr><td>McMahan et al. (2017)</td><td>S t ∼ U(N, K)</td><td>k / ∈St p k w t + k∈St p k w k t</td><td>-</td></tr><tr><td>Sahu et al. (2018)</td><td/><td/><td/></tr></table>

Let ξ k t be sampled from the k-th device's local data uniformly at random. The variance of stochastic gradients in each device is bounded:

Table content: <table><tr><td>3 and 4 have been made by the works Zhang et al. (2013); Stich (2018); Stich et al.</td></tr><tr><td>(2018); Yu et al. (2019).</td></tr><tr><td>Assumption 3.</td></tr></table>

Let Assumptions 1 to 4 hold and L, µ, σ k , G be defined therein. Let κ, γ, η t , and B be defined in Theorem 1. Let Assumption 5 hold and define

Table content: <table><tr><td>3.3 CONVERGENCE RESULT: PARTIAL DEVICE PARTICIPATION</td></tr><tr><td>As discussed in Section 2, partial device participation has more practical interest than full device</td></tr><tr><td>participation. Let the set S t (⊂ [N ]) index the active devices in the t-th iteration. To establish the</td></tr><tr><td>convergence bound, we need to make assumptions on S t .</td></tr><tr><td>Assumption 5 assumes the K indices are selected from the distribution p k independently and with</td></tr><tr><td>replacement. The aggregation step is simply averaging. This is first proposed in (Sahu et al., 2018),</td></tr><tr><td>but they did not provide theoretical analysis.</td></tr><tr><td>Assumption 5 (Scheme I). Assume S t contains a subset of K indices randomly selected with</td></tr><tr><td>replacement according to the sampling probabilities p 1 , • • • , p N . The aggregation step of FedAvg performs w t ←-1 K k∈St w k t .</td></tr><tr><td>Theorem 2.</td></tr></table>

FedAvg; while they did not require iid data, their bound do not guarantee convergence to saddle point or local minimum. Sahu et al.(2018)  proposed a federated optimization framework called FedProx to deal with statistical heterogeneity and provided the convergence guarantees in non-iid setting. FedProx adds a proximal term to each local objective. When these proximal terms vanish, FedProx is reduced to FedAvg. However, their convergence theory requires the proximal terms always exist and hence fails to cover FedAvg. When data are iid distributed and all devices are active, FedAvg is referred to as LocalSGD. Due to the two assumptions, theoretical analysis of LocalSGD is easier than FedAvg. Stich (2018) demonstrated LocalSGD provably achieves the same linear speedup with strictly less communication for strongly-convex stochastic optimization. Coppola (2015); Zhou and Cong (2017); Wang and Joshi (2018) studied LocalSGD in the non-convex setting and established convergence results. LocalSGD increase training efficiency and improve the generalization performance of large batch sizes while reducing communication. For LocalGD on non-iid data (as opposed to LocalSGD), the best result is by the contemporaneous work (but slightly later than our first version) (Khaled et al., 2019). Khaled et al. (2019) used fixed learning rate η and showed O( 1 T )convergence to a point O(η 2 E 2 ) away from the optimal. In fact, the suboptimality is due to their fixed learning rate. As we show in Theorem 4, using a fixed learning rate η throughout, the solution by LocalGD is at least Ω((E -1)η) away from the optimal.Models and datasetsWe examine our theoretical results on a logistic regression with weight decay λ = 1e -4. This is a stochastic convex optimization problem. We distribute MNIST dataset(LeCun et al., 1998) among N = 100 workers in a non-iid fashion such that each device contains samples of only two digits. We further obtain two datasets: mnist balanced and mnist unbalanced.The former is balanced such that the number of samples in each device is the same, while the latter is highly unbalanced with the number of samples among devices following a power law. To manipulate heterogeneity more precisly, we synthesize unbalanced datasets following the setup in Sahu et al.

Table content: <table><tr><td>If the data are iid, distributed optimization can be efficiently solved by the second-order algorithms</td></tr><tr><td>Mahajan et al. (2018); Reddi et al. (2016); Shamir et al. (2014); Shusen Wang et al. (2018); Zhang</td></tr><tr><td>and Lin (2015) and the one-shot methods Lee et al. (2017); Lin et al. (2017); Wang (2019); Zhang</td></tr><tr><td>et al. (2013; 2015b). The primal-dual algorithms Hong et al. (2018); Smith et al. (2016; 2017) are</td></tr><tr><td>more generally applicable and more relevant to FL.</td></tr><tr><td>6 NUMERICAL EXPERIMENTS</td></tr></table>

Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv:1806.00582, 2018. 7 Shun Zheng, Fen Xia, Wei Xu, and Tong Zhang. A general distributed dual coordinate optimization framework for regularized loss minimization. arXiv preprint arXiv:1604.03763, 2016. 1

Table content: <table><tr><td>Fan Zhou and Guojing Cong. On the convergence properties of a k-step averaging stochastic gradient</td></tr><tr><td>descent algorithm for nonconvex optimization. arXiv preprint arXiv:1708.01012, 2017. 2, 4, 7</td></tr><tr><td>Hankz Hankui Zhuo, Wenfeng Feng, Qian Xu, Qiang Yang, and Yufeng Lin. Federated reinforcement</td></tr><tr><td>learning. arXiv preprint arXiv:1901.08277, 2019. 1</td></tr><tr><td>A PROOF OF THEOREM 1</td></tr><tr><td>We analyze FedAvg in the setting of full device participation in this section.</td></tr><tr><td>A.1 ADDITIONAL NOTATION</td></tr><tr><td>Let w k</td></tr></table>

Statistics of federated datasets

Table content: <table><tr><td>Dataset</td><td>Details</td><td colspan="4"># Devices (N ) #Training samples (n) Samples/device</td></tr><tr><td/><td/><td/><td/><td>mean</td><td>std</td></tr><tr><td>MNIST</td><td>balanced unbalanced</td><td>100 100</td><td>54200 62864</td><td>542 628</td><td>0 800</td></tr><tr><td>Synthetic Data</td><td>α = 0, β = 0 α = 1, β = 1</td><td>100 100</td><td>42522 27348</td><td>425 273</td><td>1372 421</td></tr><tr><td colspan="3">D.2 THEORETICAL VERIFICATION</td><td/><td/><td/></tr></table>

Peer review (iclr2020):
This paper presents convergence rates for straggler-aware averaged SGD for non-identically but independent distributed data. The paper is well-written and motivated with good discussions of the algorithm and the related works. The proof techniques involve bounding how much worse can the algorithm do because of non-identical distribution and introduction of stragglers into the standard analysis of SGD-like algorithms. The presented theory is useful, and also provides new insights such as a new sampling scheme and an inherent bias for the case of non-decaying step size. The empirical evaluation is adequate and well-presented. I think this paper is a strong contribution and should spark further discussions in the community.