On the Convergence of FedAvg on Non-IID Data

  • 2019-10-08 15:07:31
  • Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, Zhihua Zhang
  • 0

Abstract

Federated learning makes a large amount of edge computing devices jointlylearn a model without data sharing. As a leading algorithm in this setting,Federated Averaging (\texttt{FedAvg}) runs Stochastic Gradient Descent (SGD) inparallel on a small subset of the total devices and averages the sequences onlyonce in a while. Despite its simplicity, it lacks theoretical guarantees underrealistic settings. In this paper, we analyze the convergence of\texttt{FedAvg} on non-iid data and establish a convergence rate of$\mathcal{O}(\frac{1}{T})$ for strongly convex and smooth problems, where $T$is the iteration number of SGDs. Importantly, our bound demonstrates atrade-off between communication-efficiency and convergence rate. As userdevices may be disconnected from the server, we relax the assumption of fulldevice participation to partial device participation and study differentaveraging schemes; low device participation rate can be achieved withoutseverely slowing down the learning. Our results indicate that heterogeneity ofdata slows down the convergence, which matches empirical observations.Furthermore, we provide a necessary condition for \texttt{FedAvg}'s convergenceon non-iid data: the learning rate $\eta$ must decay, even if full-gradient isused; otherwise, the solution will be $\Omega (\eta)$ away from the optimal.

 

Quick Read (beta)

On the Convergence of FedAvg on Non-IID Data

Xiang Li
School of Mathematical Sciences
Peking University
Beijing, 100871, China
[email protected]
&Kaixuan Huang
School of Mathematical Sciences
Peking University
Beijing, 100871, China
[email protected]
\ANDWenhao Yang
Center for Data Science
Peking University
Beijing, 100871, China
[email protected]
&Shusen Wang
Department of Computer Science
Stevens Institute of Technology
Hoboken, NJ 07030, USA
[email protected]
\AND
Zhihua Zhang
School of Mathematical Sciences
Peking University
Beijing, 100871, China
[email protected]
Abstract

Federated learning makes a large amount of edge computing devices 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 𝒪(1T) for strongly convex and smooth problems, where T is the iteration number of SGDs. Importantly, our bound demonstrates a trade-off between communication-efficiency 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 indicate that heterogeneity of data slows down the convergence, which matches empirical observations. Furthermore, we provide a necessary condition for FedAvg’s convergence on non-iid data: the learning rate η must decay, even if full-gradient is used; otherwise, the solution will be Ω(η) away from the optimal.

On the Convergence of FedAvg on Non-IID Data


Xiang Li
School of Mathematical Sciences
Peking University
Beijing, 100871, China
[email protected]
Kaixuan Huang
School of Mathematical Sciences
Peking University
Beijing, 100871, China
[email protected]
Wenhao Yang
Center for Data Science
Peking University
Beijing, 100871, China
[email protected]
Shusen Wang
Department of Computer Science
Stevens Institute of Technology
Hoboken, NJ 07030, USA
[email protected]
Zhihua Zhang
School of Mathematical Sciences
Peking University
Beijing, 100871, China
[email protected]
\@float

noticebox[b]\[email protected]

1 Introduction

Federated Learning (FL), also known as federated optimization, allows multiple parties to collaboratively train a model without data sharing [9, 29, 20, 11, 26, 48]. Similar to the centralized parallel optimization [7, 14, 15, 28, 44, 22, 24, 25, 31, 46, 36], FL lets 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 characters that distinguish it from the standard parallel optimization [16].

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 [20, 30, 26, 27]. 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.11 1 In the original paper [20], E epochs of SGD are performed in parallel. For theoretical analyses, we denote by E the times of updates rather than epochs. 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 to response is obviously infeasible; it is thus impractical to require all the devices to be active.

Third, the training data are non-iid22 2 Throughout this paper, “non-iid” means data are independent but not identically distributed. More precisely, the data distributions in the k-th and l-th devices, denoted Dk and Dl, can be different.; 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 not only brings challenges to algorithm design but also makes theoretical analysis much harder. While FedAvg actually works when the data are non-iid [20], FedAvg on non-iid data lacks theoretical guarantee even in the convex optimization setting.

There have been much efforts on developing convergence guarantees for FL algorithms based on the assumptions that (1) the data are iid and (2) all the devices are active. Khaled et al. [8], Yu et al. [40], Wang et al. [35] made the latter assumption, while Zhou and Cong [47], Stich [33], Wang and Joshi [34], Woodworth et al. [38] made both assumptions. These two assumptions obviously violate the second and third characters of FL.

Notation.

Let N be the total number of user devices and K (N) be the maximal number of the devices that participate in every round’s communication. Let T be the total number of every device’s SGDs and E be the number of local iterations performed in a device between two communications. Thus TE 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 Theorems 1, 2, and 3 that FedAvg has 𝒪(1T) convergence rate. In particular, Theorem 3 shows that to attain a fixed precision ϵ, the number of communications is

TE=𝒪[1ϵ((1+1K)EG2+k=1Npk2σk2+Γ+G2E)]. (1)

Here, G, Γ, pk, 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 sampling33 3 Throughout this paper, “sampling” refers to how the server chooses K user devices and uses their outputs for updating the model parameters. “Sampling” does not mean how a device randomly selects training samples. and averaging schemes (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 that FedAvg with certain schemes (see Table 1) can achieve 𝒪(1T) 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’s convergence with E>1, even if full gradient descent is used.44 4 It is well know that the full gradient descent (which is equivalent to FedAvg with E=1 and full batch) does not require the decay of the learning rate. 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.

Table 1: Sampling and averaging schemes. 𝒮t𝒰(N,K) means 𝒮t is a size-K subset uniformly sampled without replacement from [N]. 𝒮t𝒲(N,K,𝐩) means 𝒮t contains K elements that are iid sampled with replacement from [N] with probabilities {pk}. In the latter scheme, 𝒮t is not a set.
Paper Sampling Averaging Convergence rate
McMahan et al. [20] 𝒮t𝒰(N,K) k𝒮tpk𝐰t+k𝒮tpk𝐰tk -
Sahu et al. [26] 𝒮t𝒲(N,K,𝐩) 1Kk𝒮t𝐰tk 𝒪(1T)55 5 The sampling scheme is proposed by Sahu et al. [26], but this convergence rate is obtained in our paper.
Ours 𝒮t𝒰(N,K) k𝒮tpkNK𝐰tk 𝒪(1T)66 6 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.

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 deferred to the appendix.

2 Federated Averaging (FedAvg)

Problem formulation.

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

min𝐰{F(𝐰)k=1NpkFk(𝐰)}, (2)

where N is the number of devices, and pk is the weight of the k-th device such that pk0 and k=1Npk=1. Suppose the k-th device holds the nk training data: xk,1,xk,2,,xk,nk. The local objective Fk() is defined by

Fk(𝐰)1nkj=1nk(𝐰;xk,j), (3)

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

Algorithm description.

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

𝐰t+i+1k𝐰t+ik-ηt+iFk(𝐰t+ik,ξt+ik),i=0,1,,E-1,

where ηt+i is the learning rate (a.k.a. the step size) and ξt+ik is a sample uniformly chosen from the local data. Last, the server aggregates the local models, 𝐰t+E1,,𝐰t+EN, to produce the new global model, 𝐰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 𝒟k. Then the overall distribution is a mixture of all local data distributions: 𝒟=k=1Npk𝒟k. The prior work [41, 47, 33, 34, 38] assumes that the data are iid generated by or partitioned among the N devices, that is, 𝒟k=𝒟 for all k[N]. However, real-world applications do not typically satisfy the iid assumption. One of our theoretical contributions is avoiding the iid assumption.

Full device participation.

The prior work [3, 47, 33, 40, 34, 35] requires the full device participation in the aggregation step of FedAvg. In this case, the aggregation step performs

𝐰t+Ek=1Npk𝐰t+Ek.

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 implies that 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 (1K<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 𝒮t (|𝒮t|=K) be the set of the indices of the first K responded devices in the t-th iteration. The aggregation step performs

𝐰t+ENKk𝒮tpk𝐰t+Ek.

It can be proved that NKk𝒮tpk equals one in expectation.

Communication cost.

The FedAvg requires two-round communications—one broadcast and one aggregation—per E iterations. If T iterations are performed totally, then the number of communications is 2TE. During the broadcast, the central server sends 𝐰t to all the devices. During the aggregation, all or part of the N devices send their outputs, say 𝐰t+Ek, to the server.

3 Convergence Analysis of FedAvg in Non-iid Setting

In this section, we show that FedAvg converges to the global optimum at a rate of 𝒪(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 {𝐰t} have a larger variance, which, however, can be controlled by learning rates. These imply that 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 F1,,FN. Assumptions 1 and 2 are standard; typical examples are the 2-norm regularized linear regression, logistic regression, and softmax classifier.

Assumption 1.

F1,,FN are all L-smooth, that is, for all v and w, f(v)f(w)+(v-w)Tf(w)+L2v-w22.

Assumption 2.

F1,,FN are all μ-strongly convex, that is, for all v and w, f(v)f(w)+(v-w)Tf(w)+μ2v-w22.

Assumptions 3 and 4 have been made by the works [42, 33, 32, 40].

Assumption 3.

Let ξtk be sampled from the k-th device’s local data uniformly at random. The variance of stochastic gradients in each device is bounded: EFk(wtk,ξtk)-Fk(wtk)2σk2 for k=1,,N.

Assumption 4.

The expected squared norm of stochastic gradients is uniformly bounded, i.e., EFk(wtk,ξtk)2G2 for all k=1,,N and t=0,,T-1.

Quantifying the degree of non-iid (heterogeneity).

Let F* and Fk* be the minimum values of F and Fk, respectively. We use the term Γ=F*-k=1NpkFk* 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 𝐰T as the solution. We always require T is evenly divisible by E so that FedAvg can output 𝐰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

𝔼[F(𝐰T)]-F*2κγ+T(Bμ+2L𝐰0-𝐰*2), (4)

where

B=k=1Npk2σk2+6LΓ+8(E-1)2G2. (5)

3.3 Convergence Result: Partial Device Participation

As discussed in Section 2, partial device participation has more practical interest than full device participation. Let the set 𝒮t ([N]) index the active devices in the t-th iteration. To establish the convergence bound, we need to make assumptions on 𝒮t.

Assumption 5 assumes the K indices are selected from the distribution pk independently and with replacement. The aggregation step is simply averaging. This is first proposed in [26], but they did not provide theoretical analysis.

Assumption 5 (Scheme I).

Assume St contains a subset of K indices randomly selected with replacement according to the sampling probabilities p1,,pN. The aggregation step of FedAvg performs wt1KkStwtk.

Theorem 2.

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 C=4KE2G2. Then

𝔼[F(𝐰T)]-F*2κγ+T(B+Cμ+2L𝐰0-𝐰*2). (6)

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 that 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 St contains a subset of K indices uniformly sampled from [N] without replacement. Assume the data is balanced in the sense that p1==pN=1N. The aggregation step of FedAvg performs wtNKkStpkwtk.

Theorem 3.

Replace Assumption 5 by Assumption 6 and C by C=N-KN-14KE2G2. Then the same bound in Theorem 2 holds.

Scheme II requires p1==pN=1N which obviously violates the unbalance nature of FL. Fortunately, this can be addressed by the following transformation. Let F~k(𝐰)=pkNFk(𝐰) be a scaled local objective Fk. Then the global objective becomes a simple average of all scaled local objectives:

F(𝐰)=k=1NpkFk(𝐰)=1Nk=1NF~k(𝐰).

Theorem 3 still holds if L,μ,σk, and G are replaced by L~νL, μ~ςμ, σ~k=νσ, and G~=νG, respectively. Here, ν=Nmaxkpk and ς=Nminkpk.

3.4 Discussions

Choice of E.

Since 𝐰0-𝐰*24μ2G2 for μ-strongly convex F, the dominating term in eqn. (6) is

𝒪(k=1Npk2σk2+LΓ+(1+1K)E2G2+κG2μT). (7)

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

TϵE(1+1K)EG2+k=1Npk2σk2+LΓ+κG2E. (8)

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 [33] showed that if the data are iid, then E can be set to 𝒪(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 𝐰tk can converge to the minimizer of Fk, and thus FedAvg becomes the one-shot average [42] of the local solutions. If the data are non-iid, the one-shot averaging does not work because weighted average of the minimizers of F1,,FN can be very different from the minimizer of F.

Choice of K.

Stich [33] 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 that FedAvg is unable to achieve linear speedup. We will empirically see this phenomenon in Section 6. Thus, in practice, the participation ratio KN can be set small to alleviate the straggler’s effect without affecting the convergence rate.

Choice of sampling schemes.

We have considered two sampling and averaging schemes in Theorems 2 and 3. Scheme I selects K devices according to the probabilities p1,,pN with replacement. The non-uniform sampling results in faster convergence than uniform sampling, especially when p1,,pN are highly non-uniform. If the system can choose to active 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 p1,,pN are highly non-uniform, then ν=Nmaxkpk is big and ς=Nminkpk is small, which makes the convergence of FedAvg slow. This point of view is empirically verified in our experiments.

4 Necessity of Learning Rate Decay

In this section, we point out that diminishing learning rates is 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).

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* be the optimal solution. Then we have

𝐰~*-𝐰*2=Ω((E-1)η)𝐰*2,

where we hide some problem dependent constants.

Theorem 4 and its proof provide several implications. First, the decay of learning rate is necessary for FedAvg’s convergence. On one hand, Theorem 1 shows that 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 converge 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 [23]. However, Theorem 4 shows that FedAvg with E>1 and full batch size cannot converge to the optimum. This conclusion does not contradict with Theorem 1 in [8], which, when tranformed into our case, asserts that 𝐰~* will locate in the neighborhood of 𝐰* with a constant learning rate.

Third, Theorem 4 shows that 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.

5 Related Work

Federated learning (FL) was first proposed by McMahan et al. [20] for collaboratively learning a model without collecting users’ data. The research work on FL is focused on the communication-efficiency [10, 20, 26, 30] and data privacy [1, 2, 4, 5, 21]. Our work is focused on the communication-efficiency issue.

FedAvg, a synchronous distributed optimization algorithm, was proposed by [20] as an effective heuristic. Sattler et al. [27], Zhao et al. [45] studied the non-iid setting. However, they do not give convergence rate. A contemporaneous and independent work of Xie et al. [39] analyzed asynchronous FedAvg, but they did not require iid data and their bound does not guarantee convergence to saddle point or local minimum. Sahu et al. [26] proposed a federated optimization framework called FedProx to deal with statistical heterogeneity and provided the convergence guarantees in the 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 and all devices are active, FedAvg is referred to as LocalSGD. Thanks to the two assumptions, theoretical analysis of LocalSGD is easier than FedAvg. Stich [33] demonstrated that LocalSGD provably achieves the same linear speedup with strictly less communication for strongly-convex stochastic optimization. Coppola [3], Zhou and Cong [47], Wang and Joshi [34] studied LocalSGD in the non-convex setting and established convergence results. Yu et al. [40], Wang et al. [35] recently analyzed LocalSGD for non-convex functions in heterogeneous settings. In particular, Yu et al. [40] demonstrated that LocalSGD also achieves 𝒪(1/NT) convergence (i.e., linear speedup) for non-convex optimization. Lin et al. [18] empirically shows that variants of 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) [8]. Khaled et al. [8] used fixed learning rate η and showed 𝒪(1T) convergence to a point 𝒪(η2E2) 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.

If the data are iid, distributed optimization can be efficiently solved by the second-order algorithms [19, 24, 28, 36, 44] and the one-shot methods [13, 17, 37, 42, 43]. The primal-dual algorithms [6, 31, 30] are more generally applicable and more relevant to FL.

6 Numerical Experiments

Models and datasets

We examine our theoretical results on a logistic regression with weight decay λ=1e-4. This is a stochastic convex optimization problem. We distribute MNIST dataset [12] 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 highly unbalanced with the number of samples among devices following a power law, while the latter is balanced such that the number of samples in each device is the same. To manipulate heterogeneity more precisly, we synthesize unbalanced datasets following the setup in Sahu et al. [26] and denote it as synthetic(α,β) where α controls how much local models differ from each other and β controls how much the local data at each device differ from that of other devices. We obtain two datasets: synthetic(0,0) and synthetic(1,1). Details can be found in Appendix D.

(a) The impact of E
(b) The impact of K
(c) Different schemes
(d) Different schemes
Figure 1: (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 perform 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.

Experiment settings

For all experiments, we initialize all runnings with 𝐰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=η01+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.

Impact of E

We expect that Tϵ/E, the required communication round to achieve curtain accuracy, is a hyperbolic finction of E as Eqn. (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 the unbalanced datasets in Figure (a)a. The reason why the phenomenon does not appear in mnist balanced dataset requires future investigations.

Impact of K

Our theory suggests that a larger K may slightly accelerate convergence because Tϵ/E contains a term 𝒪(EG2K). Figure (b)b shows that K has limited influence on the convergence of FedAvg in synthetic(0,0) dataset. It reveals that the curve of an enough large K is slightly better. We observe the similar phenomenon among the other three datasets and attach additional results in Appendix D. This justifies that when the variance resulting from sampling is not too large (i.e., BC), 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 (c)c indicates that when data are balanced, Schemes I and II achieve nearly the same performance, both better than the original scheme. Figure (d)d shows that when the data are unbalanced, i.e., pk’s are uneven, Scheme I performs the best. Scheme II suffers from some instability in this case. This is not contradictory with our theory because we do not guarantee the convergence of Scheme II when data are unbalanced. As expected, transformed Scheme II performs stably at the price of a lower convergence rate. Compared with 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.

References

  • [1] E. Bagdasaryan, A. Veit, Y. Hua, D. Estrin, and V. Shmatikov (2018) How to backdoor federated learning. arXiv preprint arXiv:1807.00459. Cited by: §5.
  • [2] K. Bonawitz, V. Ivanov, B. Kreuter, A. Marcedone, H. B. McMahan, S. Patel, D. Ramage, A. Segal, and K. Seth (2017) Practical secure aggregation for privacy-preserving machine learning. In Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security, Cited by: §5.
  • [3] G. F. Coppola (2015) Iterative parameter mixing for distributed large-margin training of structured predictors for natural language processing. PhD thesis. Cited by: §2, §5.
  • [4] R. C. Geyer, T. Klein, M. Nabi, and S. SE (2017) Differentially private federated learning: a client level perspective. arXiv preprint arXiv:1712.07557. Cited by: §5.
  • [5] B. Hitaj, G. Ateniese, and F. Pérez-Cruz (2017) Deep models under the GAN: information leakage from collaborative deep learning. In ACM SIGSAC Conference on Computer and Communications Security, Cited by: §5.
  • [6] M. Hong, M. Razaviyayn, and J. Lee (2018) Gradient primal-dual algorithm converges to second-order stationary solution for nonconvex distributed optimization over networks. In International Conference on Machine Learning (ICML), Cited by: §5.
  • [7] D. Jakovetic (2013) Distributed optimization: algorithms and convergence rates. PhD, Carnegie Mellon University, Pittsburgh PA, USA. Cited by: §1.
  • [8] A. Khaled, K. Mishchenko, and P. Richtárik (2019) First analysis of local gd on heterogeneous data. arXiv preprint arXiv:1909.04715. Cited by: §1, §4, §5.
  • [9] J. Konečnỳ, B. McMahan, and D. Ramage (2015) Federated optimization: distributed optimization beyond the datacenter. arXiv preprint arXiv:1511.03575. Cited by: §1.
  • [10] J. Konečnỳ, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon (2016) Federated learning: strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492. Cited by: §5.
  • [11] J. Konečnỳ (2017) Stochastic, distributed and federated optimization for machine learning. arXiv preprint arXiv:1707.01155. Cited by: §1.
  • [12] Y. LeCun, L. Bottou, Y. Bengio, P. Haffner, et al. (1998) Gradient-based learning applied to document recognition. Proceedings of the IEEE 86 (11), pp. 2278–2324. Cited by: §D.1, §6.
  • [13] J. D. Lee, Q. Liu, Y. Sun, and J. E. Taylor (2017) Communication-efficient sparse regression. The Journal of Machine Learning Research 18 (1), pp. 115–144. Cited by: §5.
  • [14] M. Li, D. G. Andersen, J. W. Park, A. J. Smola, A. Ahmed, V. Josifovski, J. Long, E. J. Shekita, and B. Su (2014) Scaling distributed machine learning with the parameter server. In 11th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 14), pp. 583–598. Cited by: §1.
  • [15] M. Li, D. G. Andersen, A. J. Smola, and K. Yu (2014) Communication efficient distributed machine learning with the parameter server. In Advances in Neural Information Processing Systems (NIPS), Cited by: §1.
  • [16] T. Li, A. K. Sahu, A. Talwalkar, and V. Smith (2019) Federated learning: challenges, methods, and future directions. arXiv preprint arXiv:1908.07873. Cited by: §1.
  • [17] S. Lin, X. Guo, and D. Zhou (2017) Distributed learning with regularized least squares. Journal of Machine Learning Research 18 (1), pp. 3202–3232. Cited by: §5.
  • [18] T. Lin, S. U. Stich, and M. Jaggi (2018) Don’t use large mini-batches, use local sgd. arXiv preprint arXiv:1808.07217. Cited by: §5.
  • [19] D. Mahajan, N. Agrawal, S. S. Keerthi, S. Sellamanickam, and L. Bottou (2018) An efficient distributed learning algorithm based on effective local functional approximations. Journal of Machine Learning Research 19 (1), pp. 2942–2978. Cited by: §5.
  • [20] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas (2017) Communication-Efficient Learning of Deep Networks from Decentralized Data. In International Conference on Artificial Intelligence and Statistics (AISTATS), Cited by: §B.2, §D.2, Table 1, §1, §1, §1, §5, §5, footnote 1.
  • [21] L. Melis, C. Song, E. De Cristofaro, and V. Shmatikov (2019) Exploiting unintended feature leakage in collaborative learning. In IEEE Symposium on Security & Privacy (S&P), Cited by: §5.
  • [22] X. Meng, J. Bradley, B. Yavuz, E. Sparks, S. Venkataraman, D. Liu, J. Freeman, D. Tsai, M. Amde, and S. Owen (2016) MLlib: machine learning in Apache Spark. Journal of Machine Learning Research 17 (34), pp. 1–7. Cited by: §1.
  • [23] Y. Nesterov (2013) Introductory lectures on convex optimization: a basic course. Vol. 87, Springer Science & Business Media. Cited by: §4.
  • [24] S. J. Reddi, J. Konecnỳ, P. Richtárik, B. Póczós, and A. Smola (2016) AIDE: fast and communication efficient distributed optimization. arXiv preprint arXiv:1608.06879. Cited by: §1, §5.
  • [25] P. Richtárik and M. Takác (2016) Distributed coordinate descent method for learning with big data. Journal of Machine Learning Research 17 (1), pp. 2657–2681. Cited by: §1.
  • [26] A. K. Sahu, T. Li, M. Sanjabi, M. Zaheer, A. Talwalkar, and V. Smith (2018) Federated optimization for heterogeneous networks. arXiv preprint arXiv:1812.06127. Cited by: item I, §D.1, §D.2, §D.2, Table 1, §1, §1, §3.3, §5, §5, §6, footnote 5.
  • [27] F. Sattler, S. Wiedemann, K. Müller, and W. Samek (2019) Robust and communication-efficient federated learning from non-iid data. arXiv preprint arXiv:1903.02891. Cited by: §1, §5.
  • [28] O. Shamir, N. Srebro, and T. Zhang (2014) Communication-efficient distributed optimization using an approximate Newton-type method. In International conference on machine learning (ICML), Cited by: §1, §5.
  • [29] R. Shokri and V. Shmatikov (2015) Privacy-preserving deep learning. In Proceedings of the 22nd ACM SIGSAC Conference on Computer and Communications Security, Cited by: §1.
  • [30] V. Smith, C. Chiang, M. Sanjabi, and A. S. Talwalkar (2017) Federated multi-task learning. In Advances in Neural Information Processing Systems (NIPS), Cited by: §1, §5, §5.
  • [31] V. Smith, S. Forte, C. Ma, M. Takac, M. I. Jordan, and M. Jaggi (2016) CoCoA: a general framework for communication-efficient distributed optimization. arXiv preprint arXiv:1611.02189. Cited by: §1, §5.
  • [32] S. U. Stich, J. Cordonnier, and M. Jaggi (2018) Sparsified SGD with memory. In Advances in Neural Information Processing Systems (NIPS), pp. 4447–4458. Cited by: §3.1.
  • [33] S. U. Stich (2018) Local SGD converges fast and communicates little. arXiv preprint arXiv:1805.09767. Cited by: §A.1, §1, §2, §2, §3.1, §3.4, §3.4, §5.
  • [34] J. Wang and G. Joshi (2018) Cooperative SGD: a unified framework for the design and analysis of communication-efficient SGD algorithms. arXiv preprint arXiv:1808.07576. Cited by: §1, §2, §2, §5.
  • [35] S. Wang, T. Tuor, T. Salonidis, K. K. Leung, C. Makaya, T. He, and K. Chan (2019) Adaptive federated learning in resource constrained edge computing systems. IEEE Journal on Selected Areas in Communications. Cited by: §1, §2, §5.
  • [36] S. Wang, F. Roosta-Khorasani, P. Xu, and M. W. Mahoney (2018) GIANT: globally improved approximate newton method for distributed optimization. In Conference on Neural Information Processing Systems (NeurIPS), Cited by: §1, §5.
  • [37] S. Wang (2019) A sharper generalization bound for divide-and-conquer ridge regression. In The Thirty-Third AAAI Conference on Artificial Intelligence (AAAI), Cited by: §5.
  • [38] B. E. Woodworth, J. Wang, A. Smith, B. McMahan, and N. Srebro (2018) Graph oracle models, lower bounds, and gaps for parallel stochastic optimization. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §1, §2.
  • [39] C. Xie, S. Koyejo, and I. Gupta (2019) Asynchronous federated optimization. arXiv preprint arXiv:1903.03934. Cited by: §5.
  • [40] H. Yu, S. Yang, and S. Zhu (2019) Parallel restarted sgd with faster convergence and less communication: demystifying why model averaging works for deep learning. In AAAI Conference on Artificial Intelligence, Cited by: §1, §2, §3.1, §5.
  • [41] S. Zhang, A. E. Choromanska, and Y. LeCun (2015) Deep learning with elastic averaging SGD. In Advances in Neural Information Processing Systems (NIPS), Cited by: §2.
  • [42] Y. Zhang, J. C. Duchi, and M. J. Wainwright (2013) Communication-efficient algorithms for statistical optimization. Journal of Machine Learning Research 14, pp. 3321–3363. Cited by: §3.1, §3.4, §5.
  • [43] Y. Zhang, J. Duchi, and M. Wainwright (2015) Divide and conquer kernel ridge regression: a distributed algorithm with minimax optimal rates. Journal of Machine Learning Research 16, pp. 3299–3340. Cited by: §5.
  • [44] Y. Zhang and X. Lin (2015) DiSCO: distributed optimization for self-concordant empirical loss. In International Conference on Machine Learning (ICML), Cited by: §1, §5.
  • [45] Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra (2018) Federated learning with non-iid data. arXiv preprint arXiv:1806.00582. Cited by: §5.
  • [46] S. Zheng, F. Xia, W. Xu, and T. Zhang (2016) A general distributed dual coordinate optimization framework for regularized loss minimization. arXiv preprint arXiv:1604.03763. Cited by: §1.
  • [47] F. Zhou and G. Cong (2017) On the convergence properties of a k-step averaging stochastic gradient descent algorithm for nonconvex optimization. arXiv preprint arXiv:1708.01012. Cited by: §1, §2, §2, §5.
  • [48] H. H. Zhuo, W. Feng, Q. Xu, Q. Yang, and Y. Lin (2019) Federated reinforcement learning. arXiv preprint arXiv:1901.08277. Cited by: §1.

Appendix A Proof of Theorem 1

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

A.1 Additional Notation

Let 𝐰tk be the model parameter maintained in the k-th device at the t-th step. Let E be the set of global synchronization steps, i.e., E={nE|n=1,2,}. If t+1E, i.e., the time step to communication, FedAvg activates all devices. Then the update of FedAvg with partial devices active can be described as

𝐯t+1k=𝐰tk-ηtFk(𝐰tk,ξtk), (9)
𝐰t+1k={𝐯t+1k if t+1E,k=1Npk𝐯t+1k if t+1E. (10)

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

In our analysis, we define two virtual sequences 𝐯¯t=k=1Npk𝐯tk and 𝐰¯t=k=1Npk𝐰tk. This is motivated by [33]. 𝐯¯t+1 results from an single step of SGD from 𝐰¯t. When t+1E, both are inaccessible. When t+1E, we can only fetch 𝐰¯t+1. For convenience, we define 𝐠¯t=k=1NpkFk(𝐰tk) and 𝐠t=k=1NpkFk(𝐰tk,ξtk). Therefore, 𝐯¯t+1=𝐰¯t-ηt𝐠t and 𝔼𝐠t=𝐠¯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 ηt14L, we have

𝔼𝐯¯t+1-𝐰2(1-ηtμ)𝔼𝐰¯t-𝐰2+ηt2𝔼𝐠t-𝐠¯t2+6Lηt2Γ+2𝔼k=1Npk𝐰¯t-𝐰kt2

where Γ=F*-k=1NpkFk0.

Lemma 2 (Bounding the variance).

Assume Assumption 3 holds. It follows that

𝔼𝐠t-𝐠¯t2k=1Npk2σk2.
Lemma 3 (Bounding the divergence of {𝐰tk}).

Assume Assumption 4, that ηt is non-increasing and ηt2ηt+E for all t0. It follows that

𝔼[k=1Npk𝐰¯t-𝐰kt2] 4ηt2(E-1)2G2.

A.3 Completing the Proof of Theorem 1

Proof.

It is clear that no matter whether t+1E or t+1E, we always have 𝐰¯t+1=𝐯¯t+1. Let Δt=𝔼𝐰¯t+1-𝐰2. From Lemma 1, Lemma 2 and Lemma 3, it follows that

Δt+1(1-ηtμ)Δt+ηt2B (11)

where

B=k=1Npk2σk2+6LΓ+8(E-1)2G2.

For a diminishing stepsize, ηt=βt+γ for some β>1μ and γ>0 such that η1min{1μ,14L}=14L and ηt2ηt+E. We will prove Δtvγ+t where v=max{β2Bβμ-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+ηt2B
=(1-βμt+γ)vt+γ+β2B(t+γ)2
=t+γ-1(t+γ)2v+[β2B(t+γ)2-βμ-1(t+γ)2v]
vt+γ+1.

Then by the strong convexity of F(),

𝔼[F(𝐰¯t)]-F*L2ΔtL2vγ+t.

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

𝔼[F(𝐰¯t)]-F*2κγ+t(Bμ+2LΔ1).

A.4 Deferred proofs of key lemmas

Proof of Lemma 1..

Notice that 𝐯¯t+1=𝐰¯t-ηt𝒈t, then

𝐯¯t+1-𝐰2 =𝐰¯t-ηt𝐠t-𝐰-ηt𝐠¯t+ηt𝐠¯t2
=𝐰¯t-𝐰-ηt𝐠¯t2A1+2ηt𝐰¯t-𝐰-ηt𝐠¯t,𝐠¯t-𝐠tA2+ηt2𝐠t-𝐠¯t2 (12)

Note that 𝔼A2=0. We next focus on bounding A1. Again we split A1 into three terms:

𝐰¯t-𝐰-ηt𝐠¯t2=𝐰¯t-𝐰2-2ηt𝐰¯t-𝐰,𝐠¯tB1+ηt2𝐠¯t2B2 (13)

From the the L-smoothness of Fk(), it follows that

Fk(𝐰tk)22L(Fk(𝐰tk)-Fk). (14)

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

B2=ηt2𝐠¯t2ηt2k=1NpkFk(𝐰tk)22Lηt2k=1Npk(Fk(𝐰tk)-Fk*).

Note that

B1 =-2ηt𝐰¯t-𝐰,𝐠¯t=-2ηtk=1Npk𝐰¯t-𝐰,Fk(𝐰tk)
=-2ηtk=1Npk𝐰¯t-𝐰tk,Fk(𝐰tk)-2ηtk=1Npk𝐰tk-𝐰,Fk(𝐰tk). (15)

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

-2𝐰¯t-𝐰tk,Fk(𝐰tk)1ηt𝐰¯t-𝐰tk2+ηtFk(𝐰tk)2. (16)

By the μ-strong convexity of Fk(), we have

-𝐰tk-𝐰,Fk(𝐰tk)-(Fk(𝐰tk)-Fk(𝐰*))-μ2𝐰tk-𝐰2. (17)

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

A1=𝐰¯t-𝐰-ηt𝐠¯t2 𝐰¯t-𝐰2+2Lηt2k=1Npk(Fk(𝐰tk)-Fk*)
+ηtk=1Npk(1ηt𝐰¯t-𝐰kt2+ηtFk(𝐰tk)2)
-2ηtk=1Npk(Fk(𝐰tk)-Fk(𝐰*)+μ2𝐰tk-𝐰2)
(1-μηt)𝐰¯t-𝐰2+k=1Npk𝐰¯t-𝐰kt2
+4Lηt2k=1Npk(Fk(𝐰tk)-Fk*)-2ηtk=1Npk(Fk(𝐰tk)-Fk(𝐰*))C

where we use eqn. (14) again.

We next aim to bound C. We define γt=2ηt(1-2Lηt). Since ηt14L, ηtγt2ηt. Then we split C into two terms:

C =-2ηt(1-2Lηt)k=1Npk(Fk(𝐰tk)-Fk*)+2ηtk=1Npk(Fk(𝐰*)-Fk*)
=-γtk=1Npk(Fk(𝐰tk)-F*)+(2ηt-γt)k=1Npk(F*-Fk*)
=-γtk=1Npk(Fk(𝐰tk)-F*)D+4Lηt2Γ

where in the last equation, we use the notation Γ=k=1Npk(F*-Fk*)=F*-k=1NpkFk*.

To bound D, we have

k=1Npk(Fk(𝐰tk)-F*) =k=1Npk(Fk(𝐰tk)-Fk(𝐰¯t))+k=1Npk(Fk(𝐰¯t)-F*)
k=1NpkFk(𝐰¯t),𝐰¯tk-𝐰¯t+(F(𝐰¯t)-F*)
-12k=1Npk[ηtFk(𝐰¯t)2+1ηt𝐰tk-𝐰¯t2]+(F(𝐰¯t)-F*)
-k=1Npk[ηtL(Fk(𝐰¯t)-Fk*)+12ηt𝐰tk-𝐰¯t2]+(F(𝐰¯t)-F*)

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

Therefore

C =γtk=1Npk[ηtL(Fk(𝐰¯t)-Fk*)+12ηt𝐰tk-𝐰¯t2]-γt(F(𝐰¯t)-F*)+4Lηt2Γ
=γt(ηtL-1)k=1Npk(Fk(𝐰¯t)-F*)+(4Lηt2+γtηtL)Γ+γt2ηtk=1Npk𝐰tk-𝐰¯t2
6Lηt2Γ+k=1Npk𝐰tk-𝐰¯t2

where in the last inequality, we use the following facts: (1) ηtL-1-340 and k=1Npk(Fk(𝐰¯t)-F*)=F(𝐰¯t)-F*0 (2) Γ0 and 4Lηt2+γtηtL6ηt2L and (3) γt2ηt1.

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

A1 =𝐰¯t-𝐰-ηt𝐠¯t2
(1-μηt)𝐰¯t-𝐰2+2k=1Npk𝐰¯t-𝐰kt2+6ηt2LΓ (18)

Using eqn. (A.4) and taking expectation on both sides of eqn. (A.4), 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 σk2, then

𝔼𝐠t-𝐠¯t2 =𝔼k=1Npk(Fk(𝐰tk,ξtk)-Fk(𝐰tk))2,
=k=1Npk2𝔼Fk(𝐰tk,ξtk)-Fk(𝐰tk)2,
k=1Npk2σk2.

Proof of Lemma 3.

Since FedAvg requires a communication each E steps. Therefore, for any t0, there exists a t0t, such that t-t0E-1 and 𝐰t0k=𝐰¯t0 for all k=1,2,,N. Also, we use the fact that ηt is non-increasing and ηt02ηt for all t-t0E-1, then

𝔼k=1Npk𝐰¯t-𝐰tk2 =𝔼k=1Npk(𝐰tk-𝐰¯t0)-(𝐰¯t-𝐰¯t0)2
𝔼k=1Npk𝐰tk-𝐰¯t02
k=1Npk𝔼t=t0t-1(E-1)ηt2Fk(𝐰tk,ξtk)2
k=1Npkt=t0t-1(E-1)ηt02G2
k=1Npkηt02(E-1)2G2
4ηt2(E-1)2G2.

Appendix 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 𝐰tk is the model parameter maintained in the k-th device at the t-th step. 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+1E, the scenario considered here is that FedAvg randomly activates a subset of devices according to some sampling schemes. Again, 𝐠¯t=k=1NpkFk(𝐰tk) and 𝐠t=k=1NpkFk(𝐰tk,ξtk). Therefore, 𝐯¯t+1=𝐰¯t-ηt𝐠t and 𝔼𝐠t=𝐠¯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 t the multiset selected which allows any element to appear more than once. Note that t is only well defined for tE. For convenience, we denote by 𝒮t=N(t,E) the most recent set of chosen devices where N(t,E)=max{n|nt,nE}.

Updating scheme.

Limited to realistic scenarios (for communication efficiency and low straggler effect), FedAvg first samples a random multiset 𝒮t of devices and then only perform updates on them. This make the analysis a little bit intricate, since 𝒮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],

𝐯t+1k=𝐰tk-ηtFk(𝐰tk,ξtk), (19)
𝐰t+1k={𝐯t+1k if t+1E,samples𝒮t+1 and average {𝐯t+1k}k𝒮t+1 if t+1E. (20)

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 𝔼𝒮t(), when we take expectation to erase the latter type of randomness.

B.2 Key Lemmas

Two schemes.

For full device participation, we always have 𝐰¯t+1=𝐯¯t+1. This is true when t+1E for partial device participation. When t+1E, 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

𝔼𝒮t+1𝐰¯t+1=𝐯¯t+1.

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

  1. (I)

    The server establishes 𝒮t+1 by i.i.d. with replacement sampling an index k{1,,N} with probabilities p1,,pN for K times. Hence 𝒮t+1 is a multiset which allows a element to occur more than once. Then the server averages the parameters by 𝐰t+1k=1Kk𝒮t+1𝐯t+1k. This is first proposed in [26] but lacks theoretical analysis.

  2. (II)

    The server samples 𝒮t+1 uniformly in a without replacement fashion. Hence each element in 𝒮t+1 only occurs once.Then server averages the parameters by 𝐰t+1k=k𝒮t+1pkNK𝐯t+1k. Note that when the pk’s are not all the same, one cannot ensure k𝒮t+1pkNK=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., 𝐰¯t+1) is equal to the weighted average of parameters in all devices after SGD updates (i.e., 𝐯¯t+1). However, the original scheme in [20] (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 𝐯¯t+1 and 𝐰¯t+1 is bounded. 𝔼𝒮t𝐯¯t+1-𝐰¯t+12 is actually the variance of 𝐰¯t+1.

Lemma 4 (Unbiased sampling scheme).

If t+1IE, for Scheme I and Scheme II, we have

𝔼𝒮t(𝐰¯t+1)=𝐯¯t+1.
Lemma 5 (Bounding the variance of 𝐰¯t).

For t+1I, assume that ηt is non-increasing and ηt2ηt+E for all t0. We have the following results.

  1. (1)

    For Scheme I, the expected difference between 𝐯¯t+1 and 𝐰¯t+1 is bounded by

    𝔼𝒮t𝐯¯t+1-𝐰¯t+124Kηt2E2G2.
  2. (2)

    For Scheme II, assuming p1=p2==pN=1N, the expected difference between 𝐯¯t+1 and 𝐰¯t+1 is bounded by

    𝔼𝒮t𝐯¯t+1-𝐰¯t+12N-KN-14Kηt2E2G2.

B.3 Completing the Proof of Theorem 2 and 3

Proof.

Note that

𝐰¯t+1-𝐰*2 =𝐰¯t+1-𝐯¯t+1+𝐯¯t+1-𝐰*2
=𝐰¯t+1-𝐯¯t+12A1+𝐯¯t+1-𝐰*2A2+2𝐰¯t+1-𝐯¯t+1,𝐯¯t+1-𝐰*A3.

When expectation is taken over 𝒮t+1, the last term (A3) vanishes due to the unbiasedness of 𝐰¯t+1.

If t+1E, A1 vanishes since 𝐰¯t+1=𝐯¯t+1. We use eqn. (11) to bound A2. Then it follows that

𝔼𝐰¯t+1-𝐰*2(1-ηtμ)𝔼𝐰¯t-𝐰2+ηt2B.

If t+1E, we additionally use Lemma 5 to bound A1. Then

𝔼𝐰¯t+1-𝐰*2 =𝔼𝐰¯t+1-𝐯¯t+12+𝔼𝐯¯t+1-𝐰*2
(1-ηtμ)𝔼𝐰¯t-𝐰2+ηt2(B+C), (21)

where C is the upper bound of 1ηt2𝔼𝒮t𝐯¯t+1-𝐰¯t+12 (C is defined in Theorem 2 and 3).

The only difference between eqn. (B.3) 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 η1min{1μ,14L}=14L and ηt2ηt+E, we can prove 𝔼𝐰¯t+1-𝐰*2vγ+t where v=max{β2(B+C)βμ-1,(γ+1)𝐰1-𝐰*2}.

Then by the strong convexity of F(),

𝔼[F(𝐰¯t)]-F*L2ΔtL2vγ+t.

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

𝔼[F(𝐰¯t)]-F*2κγ+t(B+Cμ+2L𝐰1-𝐰*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 {xi}i=1N denote any fixed deterministic sequence. We sample a multiset 𝒮t (with size K) by the procedure where for each sampling time, we sample xk with probability qk for each time. Pay attention that two samples are not necessarily independent. We only require each sampling distribution is identically. Let 𝒮t={i1,,iK}[N] (some ik’s may have the same value). Then

𝔼𝒮tk𝒮txk=𝔼𝒮tk=1Kxik=K𝔼𝒮txi1=Kk=1Nqkxk.

For Scheme I, qk=pk and for Scheme II, qk=1N. 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 𝒮t+1={i1,,iK} denote the multiset of chosen indexes.

(1) For Scheme I, 𝐰¯t+1=1Kl=1K𝐯t+1il. Taking expectation over 𝒮t+1, we have

𝔼𝒮t𝐰¯t+1-𝐯¯t+12=𝔼𝒮t1K2l=1K𝐯t+1il-𝐯¯t+12=1Kk=1Npk𝐯t+1k-𝐯¯t+12 (22)

where the first equality follows from 𝐯t+1il are independent and unbiased.

To bound eqn. (22), we use the same argument in Lemma 5. Since t+1E, we know that the time t0=t-E+1E is the communication time, which implies {𝐰t0k}k=1N is identical. Then

k=1Npk𝐯t+1k-𝐯¯t+12 =k=1Npk(𝐯t+1k-𝐰¯t0)-(𝐯¯t+1-𝐰¯t0)2
k=1Npk𝐯t+1k-𝐰¯t02

where the last inequality results from k=1Npk(𝐯t+1k-𝐰¯t0)=𝐯¯t+1-𝐰¯t0 and 𝔼𝒙-𝔼𝒙2𝔼𝒙2. Similarly, we have

𝔼𝒮t𝐰¯t+1-𝐯¯t+12 1Kk=1Npk𝔼𝐯t+1k-𝐰¯t02
1Kk=1Npk𝔼𝐯t+1k-𝐰t0k2
1Kk=1NpkEi=t0t𝔼ηiFk(𝐰ik,ξik)2
1KE2ηt02G24Kηt2E2G2

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

(2) For Scheme II, when assuming p1=p2==pN=1N, we again have 𝐰¯t+1=1Kl=1K𝐯t+1il.

𝔼𝒮t 𝐰¯t+1-𝐯¯t+12=𝔼𝒮t1KiSt+1𝐯t+1i-𝐯¯t+12=1K2𝔼𝒮ti=1N𝕀{iSt}(𝐯t+1i-𝐯¯t+1)2
=1K2[i[N](iSt+1)𝐯t+1i-𝐯¯t+12+ij(i,jSt+1)𝐯t+1i-𝐯¯t+1,𝐯t+1j-𝐯¯t+1]
=1KNi=1N𝐯t+1i-𝐯¯t+12+ijK-1KN(N-1)𝐯t+1i-𝐯¯t+1,𝐯t+1j-𝐯¯t+1
=1K(N-1)(1-KN)i=1N𝐯t+1i-𝐯¯t+12

where we use the following equalities: (1) (iSt+1)=KN and (i,jSt+1)=K(K-1)N(N-1) for all ij and (2) i[N]𝐯t+1i-𝐯¯t+12+ij𝐯t+1i-𝐯¯t+1,𝐯t+1j-𝐯¯t+1=0.

Therefore,

𝔼𝐰¯t+1-𝐯¯t+12 =NK(N-1)(1-KN)𝔼[1Ni=1N𝐯t+1i-𝐯¯t+12]
NK(N-1)(1-KN)𝔼[1Ni=1N𝐯t+1i-𝐰¯t02]
NK(N-1)(1-KN)4ηt2E2G2.

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

Appendix 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

min𝐰F(𝐰)12N[𝐰𝐀𝐰-2𝐛𝐰]+μ2𝐰22, (23)

where 𝐀(Np+1)×(Np+1) , 𝐰,𝐛Np+1 and μ>0. Specifically, let 𝐛=𝐞1(1,0,,0), and 𝐀 be a symmetric and tri-diagonal matrix defined by

(𝐀)i,j={2,i=j[1,Np+1],-1,|j-i|=1andi,j[1,Np+1],0,otherwise, (24)

where i,j are row and column indices, respectively. We partition 𝐀 into a sum of N symmetric matrices (𝐀=k=1N𝐀k) and 𝐛 into 𝐛=k=1N𝐛k. Specifically, we choose 𝐛1=𝐛=𝐞1 and 𝐛2==𝐛N=0. To give the formulation of 𝐀k’s, we first introduce a series of sparse and symmetric matrices 𝐁k(1kN):

(𝐁k)i,j={1,i=j{(k-1)p+1,kp+1},2,i=j and (k-1)p+1<i,j<kp+1,-1,|j-i|=1andi,j[(k-1)p+1,kp+1],0,otherwise. (25)

Now 𝐀k’s are given by 𝐀1=𝐁1+𝐄1,1,𝐀k=𝐁k(2kN-1) and 𝐀N=𝐁N+𝐄Np+1,Np+1, where 𝐄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 (𝐀k,𝐛k) to the k-th device and construct its corresponding local objective by

Fk(𝐰)12[𝐰𝐀k𝐰-2𝐛k𝐰+μ𝐰22]. (26)

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

𝐰~*-𝐰*(E-1)η16𝐀1𝐀2𝐰*

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/55+ta where a is the best parameter chosen from the set {10-2,10-4,10-6}.

(a) Fixed learning rates
(b) Decayed learning rates
Figure 2: 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.

C.3 Some properties of the example

Recall that the symmetric matrix 𝐀(Np+1)×(Np+1) is defined in eqn. (24). Observe that 𝐀 is invertible and for all vector 𝐰Np+1,

𝐰𝐀𝐰=2i=1Np+1𝐰i2-2i=1Np𝐰i𝐰i+1=𝐰12+𝐰Np+12+i=1Np(𝐰i-𝐰i+1)24𝐰22. (27)

which implies that 0𝐀4𝐈.

The sparse and symmetric matrices 𝐁k(1kN) defined in eqn. (25) can be rewritten as

(𝐁k)=(𝟎(k-1)p×(k-1)p(1-1-12-1-12-1-11)(p+1)×(p+1)𝟎(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, Ak’s have following properties:

  1. 1.

    𝐀k is positive semidefinite with 𝐀k24;

  2. 2.

    rank(𝐀2)==rank(𝐀N-1)=p and rank(𝐀1)=rank(𝐀N)=p+1;

  3. 3.

    For each k, there exist a matrix 𝐗krk×(Np+1) such that 𝐀k=𝐗k𝐗k where rk=rank(𝐀k). Given any k, each row of 𝐗k has non-zero entries only on a block of coordinates, namely k={(k-1)p+1,(k-1)p+2,,kp+1}. As a result, 𝐀=k=1N𝐀k=𝐗𝐗, where 𝐗=(𝐗1,,𝐗N)(Np+2)×(Np+1).

  4. 4.

    𝐰*=𝐀-1𝐛 is the global minimizer of problem eqn. (23), given by (𝐰*)i=1-iNp+2(1iNp+1). Let 𝐰~(1,,1p+1,0,,0(N-1)p)Np+1, then 𝐀1𝐰~=𝐗1𝐗1𝐰~=𝐛1.

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

F1(𝐰) =12[𝐰𝐀1𝐰-2𝐛1𝐰+μ𝐰2],
=12[𝐰𝐗1𝐗1𝐰-2𝐰~𝐗1𝐗1𝐰+μ𝐰2],
=12𝐗1(𝐰-𝐰~)22+12μ𝐰2+C,

where C is some constant irrelevant with 𝐰). For 2kN,

Fk(𝐰) =12[𝐰𝐀k𝐰-2𝐛k𝐰+μ𝐰2],
=12𝐗k𝐰22+12μ𝐰2.

Similarly, the global quadratic objective eqn. (23) can be written as F(𝐰)=12N𝐗(𝐰-𝐰*)22+12μ𝐰2 .

Data in each device are sparse in the sense that non-zero features only occur in the block k of coordinates. Blocks on neighboring devices only overlap one coordinate, i.e., |kk+1|=1. These observations imply that the training data in this example is not identically distributed.

The k-th device has rk(=porp+1) non-zero feature vectors which are vertically concatenated into the feature matrix 𝐗k. Without loss of generality, we can assume all devices hold p+1 data points since we can always add additional zero vectors to expand the local dataset. Therefore n1==nN=p+1 in this case, which implies that the training data in this example is balanced distributed.

C.4 Proof of Theorem 5.

Proof of Theorem 5..

To prove the theorem, we assume that (i) all devices hold the same amount of data points, (ii) all devices perform local updates in parallel, (iii) all workers use the same learning rate η and (iv) all gradients computed by each device make use of its full local dataset (hence this case is a deterministic optimization problem). We first provide the result when μ=0.

For convenience, we slightly abuse the notation such that 𝐰t is the global parameter at round t rather than step t. Let 𝐰t(k) the updated local parameter at k-th worker at round t. Once the first worker that holds data (𝐀1,𝐛1) runs E step of SGD on F1(𝐰) from 𝐰t, it follows that

𝐰t(1)=(𝐈-η𝐀1)E𝐰t+ηl=0E-1(𝐈-η𝐀1)l𝐛1.

For the rest of workers, we have 𝐰t(k)=(𝐈-η𝐀i)E𝐰t(2kN).

Therefore, from the algorithm,

𝐰t+1=1Nk=1N𝐰t+1(k)=(1Ni=1N(𝐈-η𝐀i)E)𝐰t+ηNl=0E-1(𝐈-η𝐀1)l𝐛1.

Define ρ1Ni=1N(𝐈-η𝐀i)E2. Next we show that when η<14, we have ρ<1. From Proposition 1, 𝐀k24 and 𝐀k0 for k[N]. This means 𝐈-η𝐀k21 for all k[N]. Then for any 𝐱Np+1 and x2=1, we have 𝐱(𝐈-η𝐀k)E𝐱1 and it is monotonically decreasing when E is increasing. Then

𝐱(1Ni=1N(𝐈-η𝐀i)E)𝐱 𝐱(1Ni=1N(𝐈-η𝐀i))𝐱
=𝐱(𝐈-ηN𝐀)𝐱<1

since 0𝐀4𝐈 means 0(𝐈-ηN𝐀)𝐈.

Then 𝐰t+1-𝐰t2ρ𝐰t-𝐰t-12ρt𝐰1-𝐰02. By the triangle inequality,

𝐰t+n-𝐰t2i=0n-1𝐰t+i+1-𝐰t+i2i=0n-1ρt+i𝐰1-𝐰02ρt𝐰1-𝐰021-ρ

which implies that {𝐰t}t1 is a Cauchy sequence and thus has a limit denoted by 𝐰~*. We have

𝐰~*=(𝐈-1Ni=1N(𝐈-η𝐀i)E)-1[ηNl=0E-1(𝐈-η𝐀1)l𝐛]. (28)

Now we can discuss the impact of E.

  1. (1)

    When E=1, it follows from eqn. (28) that 𝐰~*=𝐀-1𝐛=𝐰*, i.e., FedAvg converges to the global minimizer.

  2. (2)

    When E=, limEηl=0E-1(𝐈-η𝐀1)l𝐛=𝐀1+𝐛1=𝐰~ and limE1Ni=1N(𝐈-η𝐀i)E=diag{(1-1N)𝐈p;(1-1N)𝐈-1N𝐌;(1-1N)𝐈p} where 𝐌(N-2)p+1×(N-2)p+1 is some a symmetric matrix. Actually 𝐌 is almost a diagonal matrix in the sense that there are totally N-2 completely the same matrices (i.e., 1p+1𝐞𝐞T(p+1)×(p+1)) placed on the diagonal of 𝐌 but each overlapping only the lower right corner element with the top left corner element of the next block. Therefore 𝐰~*=(1,,1p,𝐕11,,𝐕(N-2)p+1,1(N-2)p+1,0,,0p)T where 𝐕=(𝐈-𝐌)-1. From (4) of Proposition 1, 𝐰~* is different from 𝐰*

  3. (3)

    When 2E<, note that

    𝐰~*-𝐰*
    = (𝐈-1Ni=1N(𝐈-η𝐀i)E)-1[ηNl=0E-1(𝐈-η𝐀1)l𝐀-(𝐈-1Ni=1N(𝐈-η𝐀i)E)]𝐰*. (29)

    The right hand side of the last equation cannot be zero. Quantificationally speaking, we have the following lemma. We defer the proof for the next subsection.

    Lemma 6.

    If the step size η is sufficiently small, then in this example, we have

    𝐰~*-𝐰*(E-1)η16𝐀1𝐀2𝐰*. (30)

    Since 𝐀1𝐀2𝟎 and 𝐰* is dense, the lower bound in eqn. (30) is not vacuous.

Now we have proved the result when μ=0. For the case where μ>0, we replace 𝐀i with 𝐀i+μ𝐈 and assume μ<14+μ instead of the original. The discussion on different choice of E is unaffected. ∎

C.5 Proof of Lemma 6

Proof.

We will derive the conclusion mainly from the expression eqn. (29). Let f(η) be a function of η. We say a matrix 𝐓 is Θ(f(η)) if and only if there exist some positive constants namely C1 and C2 such that C1f(η)𝐓C2f(η) for all η>0. In the following analysis, we all consider the regime where η is sufficiently small.

Denote by 𝐕=i=1N𝐀i2. First we have

𝐈-1Ni=1N(𝐈-η𝐀i)E =𝐈-1Ni=1N(𝐈-Eη𝐀i+E(E-1)2η2𝐀i2+Θ(η3))
=EηN𝐀-E(E-1)2Nη2𝐕+Θ(η3). (31)

Then by plugging this equation into the right hand part of eqn. (29), we have

ηNl=0E-1(𝐈-η𝐀1)l𝐀-(𝐈-1Ni=1N(𝐈-η𝐀i)E)
= ηNl=0E-1(𝐈-lη𝐀1+Θ(η2))𝐀-(EηN𝐀-E(E-1)2Nη2𝐕+Θ(η3))
= η2N(E(E-1)2(𝐕-𝐀1𝐀)+Θ(η))

Second from eqn. (C.5), we have that

(𝐈-1Ni=1N(𝐈-η𝐀i)E)-1=(EηN𝐀+Θ(η2))-1=NEη𝐀-1+Θ(1).

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

𝐰~*-𝐰* =(NEη𝐀-1+Θ(1))η2N(E(E-1)2(𝐕-𝐀1𝐀)+Θ(η))𝐰*
=(E-12η𝐀-1(𝐕-𝐀1𝐀)+Θ(η))𝐰*
(E-1)η16(𝐕-𝐀1𝐀)𝐰*
=(E-1)η16𝐀1𝐀2𝐰*

where the last inequality holds because (i) we require η to be sufficiently small and (ii) 𝐀-1𝐱14𝐱 for any vector 𝐱 as a result of 0<𝐀4. The last equality uses the fact (i) 𝐕-𝐀1𝐀=𝐀1i=2n𝐀i and (ii) 𝐀1𝐀i=𝟎 for any i3. ∎

Appendix D Experimental Details

D.1 Experimental Setting

Model and loss.

We examine our theoretical results on a multinomial logistic regression. Specifically, let f(𝐰;xi) denote the prediction model with the parameter 𝐰=(𝐖,𝐛) and the form f(𝐰;𝐱i)=softmax(𝐖𝐱i+𝐛). The loss function is given by

F(𝐰)=1ni=1nCrossEntropy(f(𝐰;𝐱i),𝐲i)+λ𝐰22.

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 [12] because of its widely 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.

Synthetic data allow us to manipulate heterogeneity more precisely. Here we follow the same setup as described in [26]. In particular, we generate synthetic samples (𝐗k,𝐘k) according to the model y=argmax(softmax(𝐖kx+𝐛k)) with x60,𝐖k10×60 and 𝐛k10, where 𝐗knk×60 and 𝐘knk. We model each entry of 𝐖k and 𝐛k as 𝒩(μk,1) with μk𝒩(0,α), and (xk)j𝒩(vk,1j1.2) with vk𝒩(Bk,1) and Bk𝒩(0,β). 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 nk 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 summary the information of federated datasets in Table 2.

Experimental.

For all experiments, we initialize all runnings with 𝐰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=η01+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.

Table 2: Statistics of federated datasets
Dataset Details # Devices (N) #Training samples (n) Samples/device
mean std
MNIST balanced 100 54200 542 0
unbalanced 100 62864 628 800
Synthetic Data α=0,β=0 100 42522 425 1372
α=1,β=1 100 27348 273 421

D.2 Theoretical verification

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𝒪(EG2K+EG2+k=1Npk2σ2+LΓ+κG2E),

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

𝒪(Gk=1Npk2σ2+LΓ+κG2),

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 𝒪(EG2K). 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.

(a) Balanced MNIST
(b) Unbalanced MNIST
(c) Synthetic(0, 0)
(d) Synthetic(1, 1)
Figure 3: The impact of K on four datasets. To show more clearly the differences between the curves, we zoom in the last few rounds in the upper left corner of the box.

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 pk and simply averages the models, which is proposed by Sahu et al. [26]. 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 [20] 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 (c)c and (d)d.

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. [26]. They noticed that Scheme I performs slightly better than another scheme: server first uniformly samples devices and then averages local models with weight pk/lStpl. However, our theoretical framework cannot apply to it, since for t, 𝔼St𝐰¯t=𝐯¯t does not hold in general.

Our theory does not guarantee FedAvg with Scheme II could converge when the training data are unbalanced distributed. Actually, if the number of training samples varies too much among devices, Scheme II may even diverge. To illustrate this point, we have shown the terrible performance on mnist unbalanced dataset in Figure (b)b. In Figure 4, we show additional results of Scheme II on the two synthetic datasets, which are the most unbalanced. We choose b=24,K=10,E=10 and η0=0.1 for these experiments. However, transformed Scheme II performs well except that it has a lower convergence rate than Scheme I.

(a) Synthetic(0,0)
(b) Synthetic(1,1)
Figure 4: The performance of four schemes on two synthetic datasets. The Scheme I performs stably and the best. The original performs the second. The curve of the Scheme II fluctuates and has no sign of convergence. Transformed Scheme II has a lower convergence rate than Scheme I.