Year: 2021

Abstract:
Federated learning (FL) is a distributed machine learning architecture that leverages a large number of workers to jointly learn a model with decentralized data. FL has received increasing attention in recent years thanks to its data privacy protection, communication efficiency and a linear speedup for convergence in training (i.e., convergence performance increases linearly with respect to the number of workers). However, existing studies on linear speedup for convergence are only limited to the assumptions of i.i.d. datasets across workers and/or full worker participation, both of which rarely hold in practice. So far, it remains an open question whether or not the linear speedup for convergence is achievable under non-i.i.d. datasets with partial worker participation in FL. In this paper, we show that the answer is affirmative. Specifically, we show that the federated averaging (FedAvg) algorithm (with two-sided learning rates) on non-i.i.d. datasets in non-convex settings achieves a convergence rate O( 1 √ mKT + 1 T ) for full worker participation and a convergence rate O( √ K √ nT + 1 T ) for partial worker participation, where K is the number of local steps, T is the number of total communication rounds, m is the total worker number and n is the worker number in one communication round if for partial worker participation. Our results also reveal that the local steps in FL could help the convergence and show that the maximum number of local steps can be improved to T/m in full worker participation. We conduct extensive experiments on MNIST and CIFAR-10 to verify our theoretical results.

text: mKT + 1T ) for full worker participation and a
convergence rate O( √ K√ nT
+ 1T ) for partial worker participation, where K is the number of local steps, T is the number of total communication rounds, m is the total worker number and n is the worker number in one communication round if for partial worker participation. Our results also reveal that the local steps in FL could help the convergence and show that the maximum number of local steps can be improved to T/m in full worker participation. We conduct extensive experiments on MNIST and CIFAR-10 to verify our theoretical results.

heading: 1 INTRODUCTION
text: Federated Learning (FL) is a distributed machine learning paradigm that leverages a large number of workers to collaboratively learn a model with decentralized data under the coordination of a centralized server. Formally, the goal of FL is to solve an optimization problem, which can be decomposed as:
min x∈Rd
f(x) := 1
m m∑ i=1 Fi(x),
where Fi(x) , Eξi∼Di [Fi(x, ξi)] is the local (non-convex) loss function associated with a local data distribution Di and m is the number of workers. FL allows a large number of workers (such as edge devices) to participate flexibly without sharing data, which helps protect data privacy. However, it also introduces two unique challenges unseen in traditional distributed learning algorithms that are used typically for large data centers:
• Non-independent-identically-distributed (non-i.i.d.) datasets across workers (data heterogeneity): In conventional distributed learning in data centers, the distribution for each worker’s local dataset can usually be assumed to be i.i.d., i.e., Di = D,∀i ∈ {1, ...,m}. Unfortunately, this assumption rarely holds for FL since data are generated locally at the workers based on their circumstances, i.e., Di 6= Dj , for i 6= j. It will be seen later that the non-i.i.d assumption imposes significant challenges in algorithm design for FL and their performance analysis.
• Time-varying partial worker participation (systems non-stationarity): With the flexibility for workers’ participation in many scenarios (particularly in mobile edge computing), workers may randomly join or leave the FL system at will, thus rendering the active worker set stochastic and time-varying across communication rounds. Hence, it is often infeasible to wait for all workers’ responses as in traditional distributed learning, since inactive workers or stragglers will significantly slow down the whole training process. As a result, only a subset of the workers may be chosen by the server in each communication round, i.e., partial worker participation.
In recent years, the Federated Averaging method (FedAvg) and its variants (McMahan et al., 2016; Li et al., 2018; Hsu et al., 2019; Karimireddy et al., 2019; Wang et al., 2019a) have emerged as a prevailing approach for FL. Similar to the traditional distributed learning, FedAvg leverages local computation at each worker and employs a centralized parameter server to aggregate and update the model parameters. The unique feature of FedAvg is that each worker runs multiple local stochastic gradient descent (SGD) steps rather than just one step as in traditional distributed learning between two consecutive communication rounds. For i.i.d. datasets and the full worker participation setting, Stich (2018) and Yu et al. (2019b) proposed two variants of FedAvg that achieve a convergence rate of O(mKT + 1√ mKT
) with a bounded gradient assumption for both strongly convex and nonconvex problems, where m is the number of workers, K is the local update steps, and T is the total communication rounds. Wang & Joshi (2018) and Stich & Karimireddy (2019) further proposed improved FedAvg algorithms to achieve an O(mT + 1√ mKT
) convergence rate without bounded gradient assumption. Notably, for a sufficiently large T , the above rates become O( 1√
mKT )1, which
implies a linear speedup with respect to the number of workers.2 This linear speedup is highly desirable for an FL algorithm because the algorithm is able to effectively leverage the massive parallelism in a large FL system. However, with non-i.i.d. datasets and partial worker participation in FL, a fundamental open question arises: Can we still achieve the same linear speedup for convergence, i.e., O( 1√
mKT ), with non-i.i.d. datasets and under either full or partial worker participation?
In this paper, we show the answer to the above question is affirmative. Specifically, we show that a generalized FedAvg with two-sided learning rates achieves linear convergence speedup with non-i.i.d. datasets and under full/partial worker participation. We highlight our contributions as follows:
• For non-convex problems, we show that the convergence rate of the FedAvg algorithm on non-i.i.d. dataset areO( 1√
mKT + 1T ) andO( √ K√ nT
+ 1T ) for full and partial worker participation, respectively, where n is the size of the partially participating worker set. This indicates that our proposed algorithm achieves a linear speedup for convergence rate for a sufficiently large T . When reduced to the i.i.d. case, our convergence rate is O( 1TK + 1√ mKT
), which is also better than previous works. We summarize the convergence rate comparisons for both i.i.d. and non-i.i.d. cases in Table 1. It is worth noting that our proof does not require the bounded gradient assumption. We note that the SCAFFOLD algorithm (Karimireddy et al., 2019) also achieves the linear speedup but extra variance reduction operations are required, which lead to higher communication costs and implementation complexity. By contrast, we do not have such extra requirements in this paper.
• In order to achieve a linear speedup, i.e., a convergence rate O( 1√ mKT ), we show that the number
of local updatesK can be as large as T/m, which improves the T 1/3/m result previously shown in Yu et al. (2019a) and Karimireddy et al. (2019). As shown later in the communication complexity comparison in Table 1, a larger number of local steps implies relatively fewer communication rounds, thus less communication overhead. Interestingly, our results also indicate that the number of local updates K does not hurt but rather help the convergence with a proper learning rates choice in full worker participation. This overcomes the limitation as suggested in Li et al. (2019b) that local SGD steps might slow down the convergence (O(KT ) for strongly convex case). This result also reveals new insights on the relationship between the number of local steps and learning rate.
1This rate also matches the convergence rate order of parallel SGD in conventional distributed learning. 2To attain accuracy for an algorithm, it needs to take O( 1 2 ) steps with a convergence rate O( 1√
T ), while
needing O( 1 m 2 ) steps if the convergence rate is O( 1√ mT ) (the hidden constant in Big-O is the same). In this sense, one achieves a linear speedup with respect to the number of workers.
Notation. In this paper, we let m be the total number of workers and St be the set of active workers for the t-th communication round with size |St| = n for some n ∈ (0,m]. 3 We use K to denote the number of local steps per communication round at each worker. We let T be the number of total communication rounds. In addition, we use boldface to denote matrices/vectors. We let [·]it,k represent the parameter of k-th local step in the i-th worker after the t-th communication. We use ‖·‖2 to denote the `2-norm. For a natural number m, we use [m] to represent the set {1, · · · ,m}. The rest of the paper is organized as follows. In Section 2, we review the literature to put our work in comparative perspectives. Section 3 presents the convergence analysis for our proposed algorithm. Section 4 discusses the implication of the convergence rate analysis. Section 5 presents numerical results and Section 6 concludes this paper. Due to space limitation, the details of all proofs and some experiments are provided in the supplementary material.

heading: 2 RELATED WORK
text: The federated averaging (FedAvg) algorithm was first proposed by McMahan et al. (2016) for FL as a heuristic to improve communication efficiency and data privacy. Since then, this work has sparked many follow-ups that focus on FL with i.i.d. datasets and full worker participation (also known as LocalSGD (Stich, 2018; Yu et al., 2019b; Wang & Joshi, 2018; Stich & Karimireddy, 2019; Lin et al., 2018; Khaled et al., 2019a; Zhou & Cong, 2017)). Under these two assumptions, most of the theoretical works can achieve a linear speedup for convergence, i.e., O( 1√
mKT ) for a sufficiently
large T , matching the rate of the parallel SGD. In addition, LocalSGD is empirically shown to be communication-efficient and enjoys better generalization performance (Lin et al., 2018). For a comprehensive introduction to FL, we refer readers to Li et al. (2019a) and Kairouz et al. (2019).
3For simplicity and ease of presentation in this paper, we let |St| = n. We note that this is not a restrictive condition and our proofs and results still hold for |St| ≥ n, which can be easily satisfied in practice.
Algorithm 1 A Generalized FedAvg Algorithm with Two-Sided Learning Rates. Initialize x0 for t = 0, · · · , T − 1 do
The server samples a subset St of workers with |St| = n. for each worker i ∈ St in parallel do xit,0 = xt for k = 0, · · · ,K − 1 do
Compute an unbiased estimate git,k = ∇Fi(xit,k, ξit,k) of∇Fi(xit,k). Local worker update: xit,k+1 = x i t,k − ηLgit,k.
end for Let ∆it = x i t,K − xit,0 = −ηL ∑K−1 k=0 g i t,k. Send ∆ i t to the server.
end for At Server:
Receive ∆it, i ∈ S. Let ∆t = 1|S| ∑ i∈S ∆ i t. Server Update: xt+1 = xt + η∆t. Broadcasting xt+1 to workers.
end for
For non-i.i.d. datasets, many works (Sattler et al., 2019; Zhao et al., 2018; Li et al., 2018; Wang et al., 2019a; Karimireddy et al., 2019; Huang et al., 2018; Jeong et al., 2018) heuristically demonstrated the performance of FedAvg and its variants. On convergence rate with full worker participation, many works (Stich et al., 2018; Yu et al., 2019a; Wang & Joshi, 2018; Karimireddy et al., 2019; Reddi et al., 2020) can achieve linear speedup, but their convergence rate bounds could be improved as shown in this paper. On convergence rate with partial worker participation, Li et al. (2019b) showed that the original FedAvg can achieve O(K/T ) for strongly convex functions, which suggests that local SGD steps slow down the convergence in the original FedAvg. Karimireddy et al. (2019) analyzed a generalized FedAvg with two-sided learning rates under strongly convex, convex and non-convex cases. However, as shown in Table 1, none of them indicates that linear speedup is achievable with non-i.i.d. datasets under partial worker participation. Note that the SCAFFOLD algorithm (Karimireddy et al., 2019) can achieve linear speedup but extra variance reduction operations are required, which lead to higher communication costs and implementation complexity. In this paper, we show that this linear speedup can be achieved without any extra requirements. For more detailed comparisons and other algorithmic variants in FL and decentralized settings, we refer readers to Kairouz et al. (2019).

heading: 3 LINEAR SPEEDUP OF THE GENERALIZED FEDAVG WITH TWO-SIDED LEARNING RATES FOR NON-IID DATASETS
text: In this paper, we consider a FedAvg algorithm with two-sided learning rates as shown in Algorithm 1, which is generalized from previous works (Karimireddy et al., 2019; Reddi et al., 2020). Here, workers perform multiple SGD steps using a worker optimizer to minimize the local loss on its own dataset, while the server aggregates and updates the global model using another gradient-based server optimizer based on the returned parameters. Specifically, between two consecutive communication rounds, each worker performs K SGD steps with the worker’s local learning rate ηL. We assume an unbiased estimator in each step, which is denoted by git,k = ∇Fi(xit,k, ξit,k), where ξit,k is a random local data sample for k-th steps after t-th communication round at worker i. Then, each worker sends the accumulative parameter difference ∆it to the server. On the server side, the server aggregates all available ∆it-values and updates the model parameters with a global learning rate η. The FedAvg algorithm with two-sided learning rates provides a natural way to decouple the learning of workers and server, thus utilizing different learning rate schedules for workers and the server. The original FedAvg can be viewed as a special case of this framework with server-side learning rate being one.
In what follows, we show that a linear speedup for convergence is achievable by the generalized FedAvg for non-convex functions on non-i.i.d. datasets. We first state our assumptions as follows.
Assumption 1. (L-Lipschitz Continuous Gradient) There exists a constant L > 0, such that ‖∇Fi(x)−∇Fi(y)‖ ≤ L‖x− y‖,∀x,y ∈ Rd, and i ∈ [m]. Assumption 2. (Unbiased Local Gradient Estimator) Let ξit be a random local data sample in the t-th step at the i-th worker. The local gradient estimator is unbiased, i.e., E[∇Fi(xt, ξit)] = ∇Fi(xt), ∀i ∈ [m], where the expectation is over all local datasets samples. Assumption 3. (Bounded Local and Global Variance) There exist two constants σL > 0 and σG > 0, such that the variance of each local gradient estimator is bounded by E[‖∇Fi(xt, ξit)−∇Fi(xt)‖2] ≤ σ2L, ∀i ∈ [m], and the global variability of the local gradient of the cost function is bounded by ‖∇Fi(xt)−∇f(xt)‖2 ≤ σ2G, ∀i ∈ [m],∀t.
The first two assumptions are standard in non-convex optimization (Ghadimi & Lan, 2013; Bottou et al., 2018). For Assumption 3, the bounded local variance is also a standard assumption. We use a universal bound σG to quantify the heterogeneity of the non-i.i.d. datasets among different workers. In particular, σG = 0 corresponds to i.i.d. datasets. This assumption is also used in other works for FL under non-i.i.d. datasets (Reddi et al., 2020; Yu et al., 2019b; Wang et al., 2019b) as well as in decentralized optimization (Kairouz et al., 2019). It is worth noting that we do not require a bounded gradient assumption, which is often assumed in FL optimization analysis.

heading: 3.1 CONVERGENCE ANALYSIS FOR FULL WORKER PARTICIPATION
text: In this subsection, we first analyze the convergence rate of the generalized FedAvg with two-sided learning rates under full worker participation, for which we have the following result: Theorem 1. Let constant local and global learning rates ηL and η be chosen as such that ηL ≤ 18LK and ηηL ≤ 1KL . Under Assumptions 1–3 and with full worker participation, the sequence of outputs {xk} generated by Algorithm 1 satisfies:
min t∈[T ] E[‖∇f(xt)‖22] ≤ f0 − f∗ cηηLKT + Φ,
where Φ , 1c [ LηηL 2m σ 2 L +
5Kη2LL 2
2 (σ 2 L + 6Kσ 2 G)], c is a constant, f0 , f(x0), f∗ , f(x∗) and the
expectation is over the local dataset samples among workers.
Remark 1. The convergence bound contains two parts: a vanishing term f0−f∗cηηLKT as T increases and a constant term Φ whose size depends on the problem instance parameters and is independent of T . The vanishing term’s decay rate matches that of the typical SGD methods.
Remark 2. The first part of Φ (i.e., LηηL2m σ 2 L) is due to the local stochastic gradients at each worker, which shrinks at rate 1m as m increases. The cumulative variance of the K local steps contributes to the second term in Φ (i.e., 5Kη 2 LL 2
2 (σ 2 L + 6Kσ 2 G)), which is independent of m and largely affected
by the data heterogeneity. To make the second part small, an inverse relationship between the local learning rate and local steps should be satisfied, i.e., ηL = O( 1K ). Specifically, note that the global and local variances are quadratically and linearly amplified by K. This requires a sufficiently small ηL to offset the variance between two successive communication rounds to make the second term in Φ small. This is consistent with the observation in strongly convex FL that a decaying learning rate is needed for FL to converge under non-i.i.d. datasets even if full gradients used in each worker (Li et al., 2019b). However, we note that our explicit inverse relationship between ηL and K in the above is new. Intuitively, the K local steps with a sufficiently small ηL can be viewed as one SGD step with a large learning rate.
With Theorem 1, we immediately have the following convergence rate for the generalized FedAvg algorithm with a proper choice of two-sided learning rates: Corollary 1. Let ηL = 1√TKL and η = √ Km. The convergence rate of the generalized FedAvg
algorithm under full worker participation is mint∈[T ] E[‖∇f(xt)‖22] = O ( 1√ mKT + 1T ) .
Remark 3. The generalized FedAvg algorithm with two-sided learning rates can achieve a linear speedup for non-i.i.d. datasets, i.e., a O( 1√
mKT ) convergence rate as long as T ≥ mK. Although
many works have achieved this convergence rate asymptotically, we improve the maximum number
of local steps K to T/m, which is significantly better than the state-of-art bounds such as T 1/3/m shown in (Karimireddy et al., 2019; Yu et al., 2019a; Kairouz et al., 2019). Note that a larger number of local steps implies relatively fewer communication rounds, thus less communication overhead. See also the communication complexity comparison in Table 1. For example, when T = 106 and m = 100 (as used in (Kairouz et al., 2019)), the local steps in our algorithm is K ≤ T/m = 104. However, K ≤ T 1/3
m = 1 means that no extra local steps can be taken to reduce communication costs.
Remark 4. When degenerated to the i.i.d. case (σG = 0), the convergence rate becomes O( 1TK + 1√ mKT ), which has a better first term in the bound compared with previous work as shown in Table 1.

heading: 3.2 CONVERGENCE ANALYSIS FOR PARTIAL WORKER PARTICIPATION
text: Partial worker participation in each communication round may be more practical than full worker participation due to many physical limitations of FL in practice (e.g., excessive delays because of too many devices to poll, malfunctioning devices, etc.). Partial worker participation can also accelerate the training by neglecting stragglers. We consider two sampling strategies proposed by Li et al. (2018) and Li et al. (2019b). Let St be the participating worker index set at communication round t with |St| = n, ∀t, for some n ∈ (0,m]. St is randomly and independently selected either with replacement (Strategy 1) or without replacement (Strategy 2) sequentially according to the sampling probabilities pi,∀i ∈ [m]. For each member in St, we pick a worker from the entire set [m] uniformly at random with probability pi = 1m ,∀i ∈ [m]. That is, selection likelihood for anyone worker i ∈ St is p = nm . Then we have the following results: Theorem 2. Under Assumptions 1–3 with partial worker participation, the sequence of outputs {xk} generated by Algorithm 1 with constant learning rates η and ηL satisfies:
min t∈[T ] E[‖∇f(xt)‖22] ≤ f0 − f∗ cηηLKT + Φ,
where f0 = f(x0), f∗ = f(x∗), and the expectation is over the local dataset samples among workers.
For sampling Strategy 1, let η and ηL be chosen as such that ηL ≤ 18LK , ηηLKL < n−1 n and 30K2η2LL 2 − LηηLn (90K 3L2η2L + 3K) < 1. It then holds that:
Φ , 1
c
[ LηηL
2n σ2L + 3LKηηL 2n σ2G + ( 5Kη2LL 2 2 + 15K2ηη3LL 3 2n )(σ2L + 6Kσ 2 G)
] .
For sampling Strategy 2, let η and ηL be chosen as such that ηL ≤ 18LK , ηηLKL ≤ n(m−1) m(n−1) and 10K2η2LL 2 − LηηL m−nn(m−1) (90K 3η2LL 2 + 3K) < 1. It then holds that:
Φ , 1
c
[ LηηL
2n σ2L+3LKηηL m− n 2n(m− 1) σ2G+
( 5Kη2LL 2
2 +15K2ηη3LL 3 m− n 2n(m− 1)
) (σ2L+6Kσ 2 G) ] .
From Theorem 2, we immediately have the following convergence rate for the generalized FedAvg algorithm with a proper choice of two-sided learning rates: Corollary 2. Let ηL = 1√TKL and η = √ Kn. The convergence rate of the generalized FedAvg algorithm under partial worker participation and both sampling strategies are:
min t∈[T ]
E‖∇f(xt)‖22 ≤ O ( √
K√ nT + 1 T
) .
Remark 5. The convergence rate bound for partial worker participation has the same structure but with a larger variance term. This implies that the partial worker participation through the uniform sampling does not result in fundamental changes in convergence (in order sense) except for an amplified variance due to fewer workers participating and random sampling. The intuition is that the uniform sampling (with/without replacement) for worker selection yields a good approximation of the entire worker distribution in expectation, which reduces the risk of distribution deviation due to the partial worker participation. As shown in Section 5, the distribution deviation due to fewer worker participation could render the training unstable, especially in highly non-i.i.d. cases.
Remark 6. The generalized FedAvg with partial worker participation under non-i.i.d. datasets can still achieve a linear speedup O(
√ K√ nT
) with proper learning rate settings as shown in Corollary 2. In addition, when degenerated to i.i.d. case (σG = 0), the convergence rate becomes O( 1TK + 1√ nKT ).
Remark 7. Here, we let |St| = n only for ease of presentation and better readability. We note that this is not a restrictive condition. We can show that |St| = n can be relaxed to |St| ≥ n, ∀t ∈ [T ] and the same convergence rate still holds. In fact, our full proof in Appendix A.2 is for |St| ≥ n.

heading: 4 DISCUSSION
text: In light of above results, in what follows, we discuss several insights from the convergence analysis:
Convergence Rate: We show that the generalized FedAvg algorithm with two-sided learning rates can achieve a linear speedup, i.e., an O( 1√
mKT ) convergence rate with a proper choice of hyper-
parameters. Thus, it works well in large FL systems, where massive parallelism can be leveraged to accelerate training. The key challenge in convergence analysis stems from the different local loss functions (also called “model drift” in the literature) among workers due to the non-i.i.d. datasets and local steps. As shown above, we obtain a convergence bound for the generalized FedAvg method containing a vanishing term and a constant term (the constant term is similar to that of SGD). In contrast, the constant term in SGD is only due to the local variance. Note that, similar to SGD, the iterations do not diminish the constant term. The local variance σ2L (randomness of stochastic gradients), global variability σ2G (non-i.i.d. datasets), and the number of local steps K (amplification factor) all contribute to the constant term, but the total global variability in K local steps dominates the term. When the local learning rate ηL is set to an inverse relationship with respect to the number of local steps K, the constant term is controllable. An intuitive explanation is that the K small local steps can be approximately viewed as one large step in conventional SGD. So this speedup and the more allowed local steps can be largely attributed to the two-sided learning rates setting.
Number of Local Steps: Besides the result that the maximum number of local steps is improved to K ≤ T/m, we also show that the local steps could help the convergence with the proper hyperparameter choices, which supports previous numerical results (McMahan et al., 2016; Stich, 2018; Lin et al., 2018) and is verified in different models with different non-i.i.d. degree datasets in Section 5. However, there are other results showing the local steps slow down the convergence (Li et al., 2019b). We believe that whether local steps help or hurt the convergence in FL worths further investigations.
Number of Workers: We show that the convergence rate improves substantially as the the number of workers in each communication round increases. This is consistent with the results for i.i.d. cases in Stich (2018). For i.i.d. datasets, more workers means more data samples and thus less variance and better performance. For non-i.i.d. datasets, having more workers implies that the distribution of the sampled workers is a better approximation for the distribution of all workers. This is also empirically observed in Section 5. On the other hand, the sampling strategy plays an important role in non-i.i.d. case as well. Here, we adopt the uniform sampling (with/without replacement) to enlist workers to participate in FL. Intuitively, the distribution of the sampled workers’ collective datasets under uniform sampling yields a good approximation of the overall data distribution in expectation.
Note that, in this paper, we assume that every worker is available to participate once being enlisted. However, this may not always be feasible. In practice, the workers need to be in certain states in order to be able to participate in FL (e.g., in charging or idle states, etc. (Eichner et al., 2019)). Therefore, care must be taken in sampling and enlisting workers in practice. We believe that the joint design of sampling schemes and the generalized FedAvg algorithm will have a significant impact on the convergence, which needs further investigations.

heading: 5 NUMERICAL RESULTS
text: We perform extensive experiments to verify our theoretical results. We use three models: logistic regression (LR), a fully-connected neural network with two hidden layers (2NN) and a convolution neural network (CNN) with the non-i.i.d. version of MNIST (LeCun et al., 1998) and one ResNet model with CIFAR-10 (Krizhevsky et al., 2009). Due to space limitation, we relegate some experimental results in the supplementary material.
0 20 40 60 80 100 Communication Round
0.0
0.5
1.5
2.0
2.5
Tr ai
ni ng
L os
s
digits_1 digits_2 digits_5 digits_10
0 20 40 60 80 100 Communication Round
0.0
0.5
1.5
2.0
2.5
Tr ai
ni ng
L os
s
digits_10, worker number = 100 digits_10, worker number = 50 digits_10, worker number = 10 digits_2, worker number = 100 digits_2, worker number = 50 digits_2, worker number = 10
0 20 40 60 80 100 Communication Round
0.0
0.5
1.5
2.0
2.5
Tr ai
ni ng
L os
s
worker number = 100, local steps = 1 epochs worker number = 100, local steps = 5 epochs worker number = 100, local steps = 10 epochs worker number = 10, local steps = 1 epochs worker number = 10, local steps = 5 epochs worker number = 10, local steps = 10 epochs
0 20 40 60 80 100 Communication Round
0.2
0.4
0.6
0.8
1.0
Te st
A cc
ur ac
y
digits_1 digits_2 digits_5 digits_10
(a) Impact of non-i.i.d. datasets.
0 20 40 60 80 100 Communication Round
0.2
0.4
0.6
0.8
1.0
Te st
A cc
ur ac
y
digits_10, worker number = 100 digits_10, worker number = 50 digits_10, worker number = 10 digits_2, worker number = 100 digits_2, worker number = 50 digits_2, worker number = 10
(b) Impact of worker number.
0 20 40 60 80 100 Communication Round
0.2
0.4
0.6
0.8
1.0
Te st
A cc
ur ac
y
worker number = 100, local steps = 1 epochs worker number = 100, local steps = 5 epochs worker number = 100, local steps = 10 epochs worker number = 10, local steps = 1 epochs worker number = 10, local steps = 5 epochs worker number = 10, local steps = 10 epochs
(c) Impact of local steps
Figure 1: Training loss (top) and test accuracy (bottom) for the 2NN model with hyper-parameters setting: local learning rate 0.1, global learning rate 1.0: (a) worker number 100, local steps 5 epochs; (b) local steps 5 epochs; (c) 5 digits in each worker’s dataset.
In this section, we elaborate the results under non-i.i.d. MNIST datasets for the 2NN. We distribute the MNIST dataset among m = 100 workers randomly and evenly in a digit-based manner such that the local dataset for each worker contains only a certain class of digits. The number of digits in each worker’s dataset represents the non-i.i.d. degree. For digits_10, each worker has training/testing samples with ten digits from 0 to 9, which is essentially an i.i.d. case. For digits_1, each worker has samples only associated with one digit, which leads to highly non-i.i.d. datasets among workers. For partial worker participation, we set the number of workers n = 10 in each communication round.
Impact of non-i.i.d. datasets: As shown in Figure 1(a), for the 2NN model with full worker participation, the top-row figures are for training loss versus communication round and the bottomrow are for test accuracy versus communication round. We can see that the generalized FedAvg algorithm converges under non-i.i.d. datasets with a proper learning rate choice in both cases. For five digits (digits_5) in each worker’s dataset with full (partial) worker participation in Figure 1(a), the generalized FedAvg algorithm achieves a convergence speed comparable to that of the i.i.d. case (digits_10). Another key observation is that non-i.i.d. datasets slow down the convergence under the same learning rate settings for both cases. The higher the non-i.i.d. degree, the slower the convergence speed. As the non-i.i.d. degree increases (from case digits_10 to case digits_1), it is obvious that the training loss is increasing and test accuracy is decreasing. This trend is more obvious from the zigzagging curves for partial worker participation. These two observations can also be verified for other models as shown in the supplementary material, which confirms our theoretical analysis.
Impact of worker number: As shown in Figure 1(b), we compare the training loss and test accuracy between full worker participation n = 100 and partial worker participation n = 10 with the same hyper-parameters. Compared with full worker participation, partial worker participation introduces another source of randomness, which leads to zigzagging convergence curves and slower convergence. This problem is more prominent for highly non-i.i.d. datasets. For full worker participation, it can neutralize the the system heterogeneity in each communication round. However, it might not be able to neutralize the gaps among different workers for partial worker participation. That is, the datasets’ distribution does not approximate the overall distribution well. Specifically, it is not unlikely that the digits in these datasets among all active workers are only a proper subset of the total 10 digits in the original MNIST dataset, especially with highly non-i.i.d. datasets. This trend is also obvious for complex models and complicated datasets as shown in the supplementary material. The sampling strategy here is random sampling with equal probability without replacement. In practice, however, the actual sampling of the workers in FL could be more complex, which requires further investigations.
Impact of local steps: One open question of FL is that whether the local steps help the convergence or not. In Figure 1(c), we show that the local steps could help the convergence for both full and partial worker participation. These results verify our theoretical analysis. However, Li et al. (2019b) showed that the local steps may hurt the convergence, which was demonstrated under unbalanced non-i.i.d. MNIST datasets. We believe that this may be due to the combined effect of unbalanced datasets and local steps rather than just the use of local steps only.
Comparison with SCAFFOLD: Lastly, we compare with the SCAFFOLD algorithm (Karimireddy et al., 2019) since it also achieves the same linear speedup effect under non-i.i.d. datasets. We compare communication rounds, total communication load, and estimated wall-clock time under the same settings to achieve certain test accuracy, and the results are reported in Table 2. The non-i.i.d. dataset is digits_2 and the i.i.d. dataset is digits_10. The learning rates are ηL = 0.1, η = 1.0, and number of local steps K is 5 epochs. We set the target accuracy = 95% for MNIST and = 75% for CIFAR-10. Note that the total training time contains two parts: i) the computation time for training the local model at each worker and ii) the communication time for information exchanges between the workers and the server. We assume the bandwidth 20 MB/s for both uplink and downlink connections. For MNIST datasets, we can see that our algorithm is similar to or outperforms SCAFFOLD. This is because the numbers of communication rounds of both algorithms are relatively small for such simple tasks. For non-i.i.d. CIFAR-10, the SCAFFOLD algorithm takes slightly fewer number of communication rounds than our FedAvg algorithm to achieve = 75% thanks to its variance reduction. However, it takes more than 1.5 times of communication cost and wall-clock time compared to those of our FedAvg algorithm. Due to space limitation, we relegate the results of time proportions for computation and communication to Appendix B (see Figure 7).

heading: 6 CONCLUSIONS AND FUTURE WORK
text: In this paper, we analyzed the convergence of a generlized FedAvg algorithm with two-sided learning rates on non-i.i.d. datasets for general non-convex optimization. We proved that the generalized FedAvg algorithm achieves a linear speedup for convergence under full and partial worker participation. We showed that the local steps in FL could help the convergence and we improve the maximum number of local steps to T/m. While our work sheds light on theoretical understanding of FL, it also opens the doors to many new interesting questions in FL, such as how to sample optimally in partial worker participation, and how to deal with active participant sets that are both time-varying and size-varying across communication rounds. We hope that the insights and proof techniques in this paper can pave the way for many new research directions in the aforementioned areas.

heading: ACKNOWLEDGEMENTS
text: This work is supported in part by NSF grants CAREER CNS-1943226, CIF-2110252, ECCS-1818791, CCF-1934884, ONR grant ONR N00014-17-1-2417, and a Google Faculty Research Award.

heading: A APPENDIX I: PROOFS
text: In this section, we give the proofs in detail for full and partial worker participation in Section A.1 and Section A.2, respectively.

heading: A.1 PROOF OF THEOREM 1
text: Theorem 1. Let constant local and global learning rates ηL and η be chosen as such that ηL ≤ 18LK and ηηL ≤ 1KL . Under Assumptions 1–3 and with full worker participation, the sequence of outputs {xk} generated by Algorithm 1 satisfies:
min t∈[T ] E[‖∇f(xt)‖22] ≤ f0 − f∗ cηηLKT + Φ,
where Φ , 1c [ LηηL 2m σ 2 L +
5Kη2LL 2
2 (σ 2 L + 6Kσ 2 G)], c is a constant, f0 , f(x0), f∗ , f(x∗) and the
expectation is over the local dataset samples among workers. Proof. For convenience, we define ∆̄t , 1m ∑m i=1 ∆ i t. Under full device participation (i.e., St =
[m]), it is clear that ∆t = 1m ∑m i=1 ∆ i t = ∆̄t.
Due to the smoothness in Assumption 1, taking expectation of f(xt+1) over the randomness at communication round t, we have:
Et[f(xt+1)] ≤ f(xt) + 〈 ∇f(xt),Et[xt+1 − xt] 〉 + L
2 Et[‖xt+1 − xt‖2]
= f(xt)+ 〈 ∇f(xt),Et[η∆̄t + ηηLK∇f(xt)− ηηLK∇f(xt)] 〉 + L
2 η2Et[‖∆̄t‖2]
= f(xt)−ηηLK‖∇f(xt)‖2+η 〈 ∇f(xt),Et[∆̄t+ηLK∇f(xt)] 〉︸ ︷︷ ︸ A1 + L 2 η2 Et[‖∆̄t‖2]︸ ︷︷ ︸ A2 .
(1)
Note that the term A1 in (1) can be bounded as follows: A1 = 〈 ∇f(xt),Et[∆̄t + ηLK∇f(xt)]
〉 = 〈 ∇f(xt),Et [ − 1 m m∑ i=1 K−1∑ k=0 ηLg i t,k + ηLK∇f(xt) ]〉
= 〈 ∇f(xt),Et [ − 1 m m∑ i=1 K−1∑ k=0 ηL∇Fi(xit,k) + ηLK 1 m m∑ i=1 ∇Fi(xt) ]〉
= 〈√ ηLK∇f(xt),− √ ηL
m √ K Et m∑ i=1 K−1∑ k=0 (∇Fi(xit,k)−∇Fi(xt)) 〉
(a1) = ηLK
2 ‖∇f(xt)‖2+ ηL 2Km2 Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k)−∇Fi(xt) ∥∥∥∥2− ηL2Km2Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2
(a2)
≤ ηLK 2 ‖∇f(xt)‖2+ ηL 2m m∑ i=1 K−1∑ k=0 Et‖∇Fi(xit,k)−∇Fi(xt)‖2− ηL 2Km2 Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2
(a3)
≤ ηLK 2 ‖∇f(xt)‖2 +
ηLL 2
2m m∑ i=1 K−1∑ k=0 Et‖xit,k − xt‖2 − ηL 2Km2 Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2
(a4)
≤ηLK( 1
2 +15K2η2LL
2)‖∇f(xt)‖2+ 5K2η3LL 2
2 (σ2L+6Kσ 2 G)− ηL 2Km2 Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2,
(2)
where (a1) follows from that 〈 x,y 〉 = 12 [‖x‖ 2 + ‖y‖2 − ‖x − y‖2] for x = √ ηLK∇f(xt) and y = − √ ηL
m √ K ∑m i=1 ∑K−1 k=0 (∇Fi(xit,k) − ∇Fi(xt)), (a2) is due to that E[‖x1 + · · · + xn‖2] ≤
nE[‖x1‖2 + · · ·+ ‖xn‖2] , (a3) is due to Assumption 1 and (a4) follows from Lemma 2. The term A2 in (1) can be bounded as:
A2 = Et[‖∆̄t‖2] = Et [∥∥∥∥ 1m m∑ i=1 ∆it ∥∥∥∥2]
≤ 1 m2 Et [∥∥∥∥ m∑
i=1
∆it ∥∥∥∥2]
= η2L m2 Et [∥∥∥∥ m∑
i=1 K−1∑ k=0 git,k ∥∥∥∥2] (a5) =
η2L m2 Et [∥∥∥∥ m∑
i=1 K−1∑ k=0 (git,k −∇Fi(xit,k)) ∥∥∥∥2]+ η2Lm2Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2
(a6) ≤ Kη 2 L
m σ2L + η2L m2 Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2, (3)
where (a5) follows from the fact that E[‖x‖2] = E[‖x − E[x]‖2] + ‖E[x]‖2] and (a6) is due to the bounded variance assumption in Assumption 3 and the fact that E[‖x1 + · · · + xn‖2] = E[‖x1‖2 + · · ·+ ‖xn‖2] if x ′
is are independent with zero mean and E[git,j ] = ∇Fi(xit,j). Substituting the inequalities in (2) of A1 and (3) of A2 into inequality (1), we have:
Et[f(xt+1)] ≤ f(xt)−ηηLK‖∇f(xt)‖2+η < ∇f(xt),Et[∆̄t+ηLK∇f(xt)] >︸ ︷︷ ︸ A1 + L 2 η2 Et[‖∆̄t‖2]︸ ︷︷ ︸ A2
≤ f(xt)− ηηLK( 1
2 − 15K2η2LL2)‖∇f(xt)‖2 + LKη2η2L 2m σ2L
+ 5ηK2η3LL 2
2 (σ2L + 6Kσ 2 G)− ( ηηL 2Km2 − Lη 2η2L 2m2 )Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2
(a7)
≤ f(xt)−ηηLK( 1
2 −5K2η2LL2)‖∇f(xt)‖2+ LKη2η2L 2m σ2L+ 5ηK2η3LL 2 2 (σ2L+6Kσ 2 G)
(a8) ≤ f(xt)− cηηLK‖∇f(xt)‖2 + LKη2η2L
2m σ2L +
5ηK2η3LL 2
2 (σ2L + 6Kσ 2 G),
where (a7) follows from ( ηηL2Km2 − Lη2η2L 2m2 ) ≥ 0 if ηηL ≤ 1 KL , (a8) holds because there exists a constant c > 0 satisfying ( 12 − 15K 2η2LL 2) > c > 0 if ηL < 1√30KL .
Rearranging and summing from t = 0, · · · , T − 1, we have: T−1∑ t=0 cηηLKE[∇f(xt)] ≤ f(x0)− f(xT ) + T (ηηLK) [ LηηL 2m σ2L + 5Kη2LL 2 2 (σ2L + 6Kσ 2 G) ] which implies,
min t∈[T ] E‖∇f(xt)‖22 ≤ f0 − f∗ cηηLKT + Φ,
where Φ = 1c [ LηηL 2m σ 2 L +
5Kη2LL 2
2 (σ 2 L + 6Kσ 2 G)]. This completes the proof.

heading: A.2 PROOF OF THEOREM 2
text: Theorem 2. Under Assumptions 1–3 with partial worker participation, the sequence of outputs {xk} generated by Algorithm 1 with constant learning rates η and ηL satisfies:
min t∈[T ] E[‖∇f(xt)‖22] ≤ f0 − f∗ cηηLKT + Φ,
where f0 = f(x0), f∗ = f(x∗), and the expectation is over the local dataset samples among workers.
For sampling Strategy 1, let η and ηL be chosen as such that ηL ≤ 18LK , ηηLKL < n−1 n and 30K2η2LL 2 − LηηLn (90K 3L2η2L + 3K) < 1. It then holds that:
Φ , 1
c
[ LηηL
2n σ2L + 3LKηηL 2n σ2G + ( 5Kη2LL 2 2 + 15K2ηη3LL 3 2n )(σ2L + 6Kσ 2 G)
] .
For sampling Strategy 2, let η and ηL be chosen as such that ηL ≤ 18LK , ηηLKL ≤ n(m−1) m(n−1) and 10K2η2LL 2 − LηηL m−nn(m−1) (90K 3η2LL 2 + 3K) < 1. It then holds that:
Φ , 1
c
[ LηηL
2n σ2L+3LKηηL m− n 2n(m− 1) σ2G+
( 5Kη2LL 2
2 +15K2ηη3LL 3 m− n 2n(m− 1)
) (σ2L+6Kσ 2 G) ] .
Proof. Let ∆̄t be defined the same as in the proof of Theorem 1. Under partial device participation, note that ∆̄t 6= ∆t (recall that ∆̄t , 1m ∑m i=1 ∆ i t, ∆t = 1 n ∑ i∈St ∆ i t, and |St| = n). The randomness for partial worker participation contains two parts: the random sampling and the stochastic gradient. We still use Et[·] to represent the expectation with respect to both types of randomness. Due to the smoothness assumption in Assumption 1, taking expectation of f(xt+1) over the randomness at communication round t:
Et[f(xt+1)] ≤ f(xt) + 〈 ∇f(xt),Et[xt+1 − xt] 〉 + L
2 Et[‖xt+1 − xt‖2]
= f(xt) + 〈 ∇f(xt),Et[η∆t + ηηLK∇f(xt)− ηηLK∇f(xt)] 〉 + L
2 η2Et[‖∆t‖2]
= f(xt)−ηηLK‖∇f(xt)‖2+η 〈 ∇f(xt),Et[∆t+ηLK∇f(xt)] 〉︸ ︷︷ ︸ A ′ 1 + L 2 η2 Et[‖∆t‖2]︸ ︷︷ ︸ A ′ 2
(4)
The term A ′ 1 in (4) can be bounded as follows: Since ESt [A ′
1] = A1 due to Lemma 1 for both sampling strategies, we have the same bound as in inequality 2 for A ′
1:
A ′ 1 ≤ ηLK( 1
2 + 15K2η2LL
2)‖∇f(xt)‖2 + 5K2η3LL 2
2 (σ2L + 6Kσ 2 G)
− ηL 2Km2 Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2, (5)
For strategy 1: We can bound A ′
2 in (4) as follows.
Note St is an index set (multiset) for independent sampling (equal probability) with replacement in which some elements may have the same value. Suppose St = {l1, . . . , ln}.
A ′
2 = Et[‖∆t‖2] = Et [∥∥∥∥ 1n ∑
i∈St
∆it ∥∥∥∥2]
= 1 n2 Et [∥∥∥∥∑
i∈St
∆it ∥∥∥∥2]
= 1 n2 Et [∥∥∥∥ n∑
z=1
∆lzt ∥∥∥∥2] (b1) =
η2L n2 Et [∥∥∥∥ n∑
z=1 K−1∑ j=0 [glzt,j −∇Flz (x lz t,j)] ∥∥∥∥2]+ η2Ln2Et [∥∥∥∥ n∑ z=1 K−1∑ j=0 ∇Flz (x lz t,j) ∥∥∥∥2] (b2)
≤ Kη 2 L
n σ2L + η2L n2 Et [∥∥∥∥ n∑
z=1 K−1∑ j=0 ∇Flz (x lz t,j) ∥∥∥∥2], where (b1) follows from the fact that E[‖x‖2] = E[‖x− E[x]‖2] + ‖E[x]‖2] and (b2) is due to the bounded variance assumption 3 and E[‖x1 + · · ·+ xn‖2] ≤ nE[‖x1‖2 + · · ·+ ‖xn‖2].
By letting ti = ∑K−1 j=0 ∇Fi(xit,j), we have:
Et [∥∥∥∥ n∑
z=1 K−1∑ j=0 ∇Flz (x lz t,j) ∥∥∥∥2 = Et[∥∥∥∥ n∑ z=1 tlz ∥∥∥∥2]
= Et [ n∑ z=1 ‖tlz‖2 + ∑
i 6=j;li,lj∈St
〈 tli , tlj 〉] (b3) = Et [ n‖tl1‖2 + n(n− 1) 〈 tl1 , tl2
〉] = n
m m∑ i=1 ‖ti‖2 + n(n− 1) m2 ∑ i,j∈[m] 〈 ti, tj 〉 = n
m m∑ i=1 ‖ti‖2 + n(n− 1) m2 ‖ m∑ i=1 ti‖2,
where (b3) is due to the independent sampling with replacement.
So we can bound A ′
2 as follows.
A ′
2 = Et[‖∆t‖2]
≤ Kη 2 L
n σ2L + η2L mn m∑ i=1 Et‖ti‖2 + (n− 1)η2L m2n Et ∥∥∥∥ m∑ i=1 ti ∥∥∥∥2, (6) For ti, we have:
m∑ i=1 Et‖ti‖2 = m∑ i=1 Et ∥∥∥∥K−1∑ j=0 ∇Fi(xit,j)−∇Fi(xt) +∇Fi(xt)−∇f(xt) +∇f(xt) ∥∥∥∥2
(b4) ≤ 3KL2 m∑ i=1 K−1∑ j=0 Et‖xit,j − xt‖2 + 3mK2σ2G + 3mK2‖∇f(xt)‖2
(b5)
≤ 15mK3L2η2L(σ2L+6Kσ2G)+(90mK4L2η2L + 3mK2)‖∇f(xt)‖2+3mK2σ2G, (7)
where (b4) is due to the fact that E[‖x1 + · · ·+ xn‖2] ≤ nE[‖x1‖2 + · · ·+ ‖xn‖2] , Assumptions 3 and 1, and (b5) follows from Lemma 2.
Substituting the inequalities in ( 5) of A ′ 1 and ( 6) of A ′ 2 into inequality (4), we have:
Et[f(xt+1)] ≤ f(xt)−ηηLK‖∇f(xt)‖2+η 〈 ∇f(xt),Et[∆t+ηLK∇f(xt)] 〉︸ ︷︷ ︸ A ′ 1 + L 2 η2 Et[‖∆t‖2]︸ ︷︷ ︸ A ′ 2
≤ f(xt)− ηηLK( 1
2 − 15K2η2LL2)‖∇f(xt)‖2 +
5ηK2η3LL 2
2 (σ2L + 6Kσ 2 G)
+
[ (n− 1)Lη2η2L
2m2n − ηηL 2Km2 ] Et ∥∥∥∥ m∑ i=1 ti ∥∥∥∥2+LKη2η2L2n σ2L+Lη2η2L2mn m∑ i=1 Et‖ti‖2
(b6)
≤ f(xt)− ηηLK( 1
2 − 15K2η2LL2)‖∇f(xt)‖2 +
5ηK2η3LL 2
2 (σ2L + 6Kσ 2 G)
+ LKη2η2L
2n σ2L + Lη2η2L 2mn m∑ i=1 Et‖ti‖2
(b7)
≤ f(xt)− ηηLK( 1
2 − 15K2η2LL2 − LηηL 2n (90K3L2η2L + 3K))‖∇f(xt)‖2
+
[ 5ηK2η3LL 2
2 + 15K3L3η2η4L 2n
] (σ2L+6Kσ 2 G)+
LKη2η2L 2n σ2L+ 3K2Lη2η2L 2n σ2G
(b8) ≤ f(xt)− cηηLK‖∇f(xt)‖2 + LKη2η2L
2n σ2L + 3K2Lη2η2L 2n σ2G
+ ηηLK
[ 5Kη2LL 2
2 +
15K2η3LηL 3
2n
] (σ2L + 6Kσ 2 G), (8)
where (b6) follows from (n−1)Lη 2η2L
2m2n − ηηL 2Km2 ≤ 0 if ηηLKL ≤ n−1 n , (b7)is due to inequality (7) and
(b8) holds since there exists a constant c > 0 such that [ 12−15K 2η2LL 2− LηηL2n (90K 3L2η2L+3K)] > c > 0 if 30K2η2LL 2 − LηηLn (90K 3L2η2L + 3K) < 1.
Note that the requirement of |St| = n can be relaxed to |St| ≥ n. With pt ≥ n workers in t-th communication round, 8 is
Et[f(xt+1)] ≤ f(xt)− cηηLK‖∇f(xt)‖2 + LKη2η2L
2pt σ2L + 3KLη2η2L 2pt σ2G
+ ηηLK
[ 5Kη2LL 2
2 +
15Kη3LηL 3
2pt
] (σ2L + 6Kσ 2 G)
≤ f(xt)− cηηLK‖∇f(xt)‖2 + LKη2η2L
2n σ2L + 3K2Lη2η2L 2n σ2G
+ ηηLK
[ 5Kη2LL 2
2 +
15K2η3LηL 3
2n
] (σ2L + 6Kσ 2 G).
That is, the same convergence rate can be guaranteed if at least n workers in each communication round (no need to be exactly n).
Rearranging and summing from t = 0, · · · , T − 1, we have the convergence for partial device participation with sampling strategy 1 as follows:
min t∈[T ] E[‖∇f(xt)‖22] ≤ f0 − f∗ cηηLKT + Φ,
where Φ = 1c [ LηηL 2n σ 2 L + 3KLηηL 2n σ 2 G + ( 5Kη2LL 2 2 + 15K2ηη3LL 3 2n )(σ 2 L + 6Kσ 2 G) ] and c is a constant.
For strategy 2: Under the strategy of independent sampling with equal probability without replacement. We bound A ′
2 as follows.
A ′
2 = Et[‖∆t‖2] = Et [∥∥∥∥ 1n ∑
i∈St
∆it ∥∥∥∥2]
= 1 n2 Et [∥∥∥∥∑
i∈St
∆it ∥∥∥∥2]
= 1 n2 Et [∥∥∥∥ m∑
i=1
I{i ∈ St}∆it ∥∥∥∥2]
= η2L n2 Et [∥∥∥∥ m∑
i=1 I{i ∈ St} K−1∑ j=0 [git,j−∇Fi(xit,j)] ∥∥∥∥2]+ η2Ln2Et [∥∥∥∥ m∑ i=1 I{i ∈ St} K−1∑ j=0 ∇Fi(xit,j)] ∥∥∥∥2]
= η2L n2 Et [∥∥∥∥ m∑
i=1 P{i ∈ St} K−1∑ j=0 [git,j −∇Fi(xit,j)] ∥∥∥∥2 + η2Ln2 ∥∥∥∥ m∑ i=1 I{i ∈ St} K−1∑ j=0 ∇Fi(xit,j) ∥∥∥∥2]
(b9) = η2L nm Et [ m∑ i=1 K−1∑ j=0 ∥∥∥∥git,j −∇Fi(xit,j)∥∥∥∥2]+ η2Ln2Et [∥∥∥∥ m∑ i=1 I{i ∈ St} K−1∑ j=0 ∇Fi(xit,j) ∥∥∥∥2]
(b10) ≤ Kη 2 L
n σ2L + η2L n2 ∥∥∥∥ m∑ i=1 P{i ∈ St} K−1∑ j=0 ∇Fi(xit,j) ∥∥∥∥2, (9)
where (b9) is due to the fact that E[‖x1 + · · · + xn‖2] = E[‖x1‖2 + · · · + ‖xn‖2] if x ′
is are independent with zero mean, xi = git,j −∇Fi(xit,j) is independent random variable with mean zero, and P{i ∈ St} = nm . (b10) is due to bounded variance assumption in Assumption 3
Substituting the inequalities in (5) of A ′ 1 and (9) of A ′ 2 into inequality (4), we have:
Et[f(xt+1)]≤f(xt)−ηηLK‖∇f(xt)‖2+η 〈 ∇f(xt),Et[∆t+ηLK∇f(xt)] 〉︸ ︷︷ ︸ A ′ 1 + L 2 η2 Et[‖∆t‖2]︸ ︷︷ ︸ A ′ 2
≤ ∇f(xt)− ηηLK( 1
2 − 15K2η2LL2)‖∇f(xt)‖2 + LKη2η2L 2n σ2L + 5ηK2η3LL 2 2 (σ2L + 6Kσ 2 G)
+ Lη2η2L 2n2 Et ∥∥∥∥ m∑ i=1 P{i ∈ St} K−1∑ j=0 ∇Fi(xit,j) ∥∥∥∥2 − ηηL2Km2Et ∥∥∥∥ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k) ∥∥∥∥2︸ ︷︷ ︸
A ′ 3
.
Then we bound A ′
3 as follows. By letting ti = ∑K−1 j=0 ∇Fi(xit,j), we have:
m∑ i=1 Et‖ti‖2 ≤ 15mK3L2η2L(σ2L + 6Kσ2G) + (90mK4L2η2L + 3mK2)‖∇f(xt)‖2 + 3mK2σ2G.
It then follows that
‖ m∑ i=1 ti‖2 = ∑ i∈[m] ‖ti‖2 + ∑ i6=j < ti, tj >
(b11) = ∑ i∈[m] m‖ti‖2 − 1 2 ∑ i6=j ‖ti − tj‖2
‖ m∑ i=1 P{i ∈ St}ti‖2 = ∑ i∈[m] P{i ∈ St}‖ti‖2 + ∑ i 6=j P{i, j ∈ St} < ti, tj >
(b12) =
n
m ∑ i∈[m] ‖ti‖2 + n(n− 1) m(m− 1) ∑ i 6=j < ti, tj >
(b13) =
n2
m ∑ i∈[m] ‖ti‖2 − n(n− 1) 2m(m− 1) ∑ i 6=j ‖ti − tj‖2,
where (b11) and (b13) are due to the fact that 〈 x,y 〉 = 12 [‖x‖ 2+‖y‖2−‖x−y‖2] ≤ 12 [‖x‖ 2+‖y‖2], (b12) follows from the fact that P{i ∈ St} = nm and P{i, j ∈ St} = n(n−1) m(m−1) . Therefore, we have
A ′ 3 = Lη2η2L 2n2 ‖ m∑ i=1 P{i ∈ St} K−1∑ j=0 ∇Fi(xit,j)]‖2 − ηηL 2Km2 ‖ m∑ i=1 K−1∑ k=0 ∇Fi(xit,k)‖2
= ( Lη2η2L 2m − ηηL 2Km ) m∑ i=1 ‖ti‖2 + ( ηηL 4Km2 − Lη 2η2L(n− 1) 4mn(m− 1) ) ∑ i 6=j ‖ti − tj‖2
(b14) = ( Lη2η2L 2m − Lη 2η2L(n− 1) 2n(m− 1) ) m∑ i=1 ‖ti‖2 − ( ηηL 2Km2 − Lη 2η2L(n− 1) 2mn(m− 1) )‖ ∑ i∈[m] ti‖2
(b15) ≤ (Lη 2η2L 2m − Lη 2η2L(n− 1) 2n(m− 1) ) m∑ i=1 ‖ti‖2
= Lη2η2L m− n
2mn(m− 1) m∑ i=1 ‖ti‖2,
where (b14) follows from the fact that ‖ ∑ i∈[m] ti‖2 = ∑ i∈[m]m‖ti‖2 − 1 2 ∑ i 6=j ‖ti − tj‖2, and (b15) is due to the fact that ( ηηL2Km2 − Lη2η2L(n−1) 2mn(m−1) ) ≥ 0 if ηηLKL ≤ n(m−1) m(n−1) .
Then we have
Et[f(xt+1)] ≤ f(xt)−ηηLK( 1
2 −15K2η2LL2−LηηL m− n 2n(m− 1) (90K3η2LL 2+3K))‖∇f(xt)‖2
+ LKη2η2L
2n σ2L + 3K
2Lη2η2L m− n
2n(m− 1) σ2G
+ ηηLK( 5Kη2LL 2
2 + 15Kηη3LL 3 m− n 2n(m− 1) )(σ2L + 6Kσ 2 G)
(b16) ≤ f(xt)− cηηLK‖∇f(xt)‖2 + LKη2η2L
2n σ2L + 3KLη
2η2L m− n
2n(m− 1) σ2G
+ ηηLK( 5Kη2LL 2
2 + 15K2ηη3LL 3 m− n 2n(m− 1) )(σ2L + 6Kσ 2 G), (10)
where (b16) holds because there exists a constant c > 0 satisfying ( 12 − 5K 2η2LL 2 − LηηL m−n 2n(m−1) (90K 3η2LL 2 + 3K)) > c > 0 if 10K2η2LL 2−LηηL m−nn(m−1) (90K 3η2LL 2 + 3K) < 1.
Note that the requirement of |St| = n can be relaxed to |St| ≥ n. With pt ≥ n workers in t-th communication round, 10 is
Et[f(xt+1)] ≤ f(xt)− cηηLK‖∇f(xt)‖2 + LKη2η2L
2pt σ2L + 3KLη
2η2L m− pt
2pt(m− 1) σ2G
+ ηηLK( 5Kη2LL 2
2 + 15K2ηη3LL 3 m− pt 2pt(m− 1) )(σ2L + 6Kσ 2 G)
≤ f(xt)− cηηLK‖∇f(xt)‖2 + LKη2η2L
2n σ2L + 3KLη
2η2L m− n
2n(m− 1) σ2G
+ ηηLK( 5Kη2LL 2
2 + 15K2ηη3LL 3 m− n 2n(m− 1) )(σ2L + 6Kσ 2 G)
That is, the same convergence rate can be guaranteed if at least n workers in each communication round (no need to be exactly n).
Rearranging and summing from t = 0, · · · , T − 1, we have the convergence for partial device participation with sampling strategy 2 as follows:
min t∈[T ] E[‖∇f(xt)‖22] ≤ f0 − f∗ cηηLKT + Φ,
where Φ = 1c [ LηηL 2n σ 2 L + 3KLηηL m−n 2n(m−1)σ 2 G + ( 5Kη2LL 2 2 + 15K 2ηη3LL 3 m−n 2n(m−1) )(σ 2 L + 6Kσ 2 G) ] and c is a constant. This completes the proof.

heading: A.2.1 KEY LEMMAS
text: Lemma 1 (Unbiased Sampling). For strategies 1 and 2, the estimator ∆t is unbiased, i.e.,
ESt [∆t] = ∆̄t.

heading: Proof of Lemma 1.
text: Let St = {t1, · · · , tn} with size n. Both for sampling strategies 1 and 2, each sampling distribution is identical. Then we have:
ESt [∆t] = 1
n ESt [ ∑ ti∈St ∆tit ] = 1 n ESt [ n∑ i=1 ∆tit ] = ESt [∆ t1 t ] = 1 m m∑ i=1 ∆it = ∆̄t.

heading: A.3 AUXILIARY LEMMAS
text: Lemma 2 (Lemma 4 in Reddi et al. (2020)). For any step-size satisfying ηL ≤ 18LK , we can have the following results:
1
m m∑ i=1 E[‖xit,k − xt‖2] ≤ 5Kη2L(σ2L + 6Kσ2G) + 30K2η2L‖∇f(xt)‖2.
Proof. In order for this paper to be self-contained, we restate the proof of Lemma 4 in (Reddi et al., 2020) here.
For any worker i ∈ [m] and k ∈ [K], we have:
E[‖xit,k − xt‖2] = E[‖xit,k−1 − xt − ηLgtt,k−1‖2] ≤E[‖xit,k−1−xt−ηL(gtt,k−1−∇Fi(xit,k−1)+∇Fi(xit,k−1)−∇Fi(xt)+∇Fi(xt)−∇f(xt)+∇f(xt))‖2]
≤ (1 + 1 2K − 1 )E[‖xit,k−1 − xt‖2] + E[‖ηL(gtt,k−1 −∇Fi(xit,k−1))‖2]
+6KE[‖ηL(∇Fi(xit,k−1)−∇Fi(xt))‖2]+6KE[‖ηL(∇Fi(xt)−∇f(xt)))‖2]+6K‖ηL∇f(xt)‖2
≤(1+ 1 2K−1 )E[‖xit,k−1−xt‖2]+η2Lσ2L+6Kη2LL2E[‖xit,k−1−xt‖2]+6Kη2Lσ2G+6K‖ηL∇f(xt)‖2
= (1 + 1
2K − 1 + 6Kη2LL 2)E[‖xit,k−1 − xt‖2] + η2Lσ2L + 6Kη2Lσ2G + 6K‖ηL∇f(xt)‖2
≤ (1 + 1 K − 1 )E[‖xit,k−1 − xt‖2] + η2Lσ2L + 6Kη2Lσ2G + 6K‖ηL∇f(xt)‖2
Unrolling the recursion, we get:
1
m m∑ i=1 E[‖xit,k − xt‖2] ≤ k−1∑ p=0 (1 + 1 K − 1 )p[η2Lσ 2 L + 6Kσ 2 G + 6Kη 2 L‖ηL∇f(xt))‖2]
≤ (K − 1)[(1 + 1 K − 1 )K − 1][η2Lσ2L + 6Kσ2G + 6Kη2L‖ηL∇f(xt))‖2]
≤ 5Kη2L(σ2L + 6Kσ2G) + 30K2η2L‖∇f(xt)‖2
This completes the proof.

heading: B APPENDIX II: EXPERIMENTS
text: We provide the full detail of the experiments. We uses non-i.i.d. versions for MNIST and CIFAR-10, which are described as follows:

heading: B.1 MNIST
text: We study image classification of handwritten digits 0-9 in MNIST and modify the MNIST dataset to a non-i.i.d. version.
To impose statistical heterogeneity, we split the data based on the digits (p) they contain in their dataset. We distribute the data to m = 100 workers such that each worker contains only a certain class of digits with the same number of training/test samples. For example, for p = 1, each worker only has training/testing samples with one digit, which causes heterogeneity among different workers. For p = 10, each worker has samples with 10 digits, which is essentially i.i.d. case. In this way, we can use the digits in worker’s local dataset to represent the non-i.i.d. degree qualitatively. In each communication round, 100 workers run K epochs locally in parallel and then the server samples n workers for aggregation and update. We make a grid-search experiments for the hyper-parameters as shown in Table 3.
We run three models: multinomial logistic regression, fully-connected network with two hidden layers (2NN) (two 200 neurons hidden layers with ReLU followed by an output layer), convolutional neural network (CNN), as shown in Table 4. The results are shown in Figures 2, 3 and 4.
Unless stated otherwise, we use the following default parameter setting: the server learning rate and client learning rate are set to η = 1.0 and ηL = 0.1, respectively. The local epochs is set to K = 10. The total number of clients is set to 100, and the clients partition number is set to n = 10. We use the same strategy to distribute the data over clients as suggested in McMahan et al. (2016). For the i.i.d. setting, we evenly partition all the training data among all clients, i.e., each client observes 500 data; for the non-i.i.d. setting, we first sort the training data by label, then divide all the training data into 200 shards of size 250, and randomly assign two shards to each client. For the CIFAR-10 dataset, we train our classifier with the ResNet model. The results are shown in Figure 5 and Figure 6.

heading: B.3 DISCUSSION
text: Impact of non-i.i.d. datasets: Figure 2 shows the results of training loss (top) and test accuracy (bottom) for three models under different non-i.i.d. datasets with full and partial worker participation
on MNIST. We can see that the FedAvg algorithm converges under non-i.i.d. datasets with a proper learning rate choice in these cases. We believe that the major challenge in FL is the non-i.i.d. datasets. For these datasets with a lower degree of non-i.i.d., the FedAvg algorithm can achieve a good result compared with the i.i.d. case. For example, when the local dataset in each worker has five digits (p = 5) with full (partial) worker participation, the FedAvg algorithm achieves a convergence speed comparable with that of the i.i.d. case (p = 10). This result can be observed in Figure 2 for all three models. As the degree of non-i.i.d. datasets increases, its negative impact on the convergence is becoming more obvious. The higher the degree of non-i.i.d., the slower the convergence speed. As the non-i.i.d. degree increases (from case p = 10 to case p = 1), it is obvious that the training loss is increasing and test accuracy is decreasing. For these with high degree of non-i.i.d., the convergence curves oscillate and are highly unstable. This trend is more obvious for complex models such for CNN in Figure 2(c).
Impact of worker number: For full worker participation, the server can have an accurate estimation of the system heterogeneity after receiving the updates for all workers and neutralize this heterogeneity in each communication round. However, partial worker participation introduces another source of randomness, which leads to zigzagging convergence curves and slower convergence. In each
communication round, the server can only receive a subset of workers based on the sampling strategy. So the server could only have a coarse estimation of the system heterogeneity and might not be able to neutralize the heterogeneity among different workers for partial worker participation. This problem is more prominent for highly non-i.i.d. datasets. It is not unlikely that the digits in these datasets among all active workers are only a proper subset of the total 10 digits in the original MNIST dataset, especially with highly non-i.i.d. datasets. For example, for p = 1 with 10 workers in each
communication round, it is highly likely that the datasets formed by these ten workers only includes certain small number of digits (say, 4 or 5) rather than total 10 digits. But for p = 5, it is the opposite, that is, the digits in these datasets among these 10 workers are highly likely to be 10. So in each communication round, the server can mitigate system heterogeneity since it covers the training samples with all 10 digits. This trend is more obvious for complex models and datasets given the dramatic drop of test accuracy in the result of CIFAR-10 in Figure 5.
The sample strategy here is random sampling with equal probability without replacement. In practice, the workers need to be in certain states in order to be able to participate in FL (e.g., in charging or idle states, etc.(Eichner et al., 2019)). Therefore, care must be taken in sampling and enlisting workers in practice. We believe that the joint design of sampling schemes, number of workers and the FedAvg algorithm will have a significant impact on the convergence, which needs further investigations.
Impact of local steps: Figure 3 and Figure 4 shows the results of training loss (top) and test accuracy (bottom) for three models under different local steps with full and partial worker participation respectively. Figure 6 shows the impact of local steps in CIFAR-10. One open question of FL is that whether the local steps help the convergence or not. Li et al. (2019b) showed a convergence rate O(KT ), i.e., the local steps may hurt the convergence for full and partial worker participation. In this two figures, we can see that local steps could help the convergence for both full and partial worker participation. However, it only has a slight effect on the convergence compared to the effects of non-i.i.d. datasets and number of workers.
Comparison with SCAFFOLD: We compare SCAFFOLD (Karimireddy et al., 2019) with the generalized FedAVg algorithm in this paper in terms of communication rounds, total communication overloads and estimated wall-clock time to achieve certain test accuracy in Table 2. We run the experiments using the same GPU (NVIDIA V100) to ensure the same conditions. Here, we give a specific comparison for these two algorithms under exact condition. Note that we divide the total training time to two parts: the computation time when the worker trains the local model and the communication time when information exchanges between the worker and server. We only compare the computation time and communication time with a fixed bandwidth 20MB/s for both uploading and downloading connections. As shown in Figure 7, to achieve = 75%, SCAFFOLD performs less communication round due to the variance reduction techniques. That is, it spends less time on computation. However, it needs to communicates as twice as the FedAvg since the control variate to perform variance reduction in each worker needs to update in each round. In this way, the communication time would be largely prolonged.

Review summaries:
This paper provides a new analysis for the FedAvg algorithm, which assumes the data on different workers are non-IID and the objective functions are non-convex. The new analysis improved the existing bounds of FedAvg. Besides, the analysis is also extended to the non-stationary network, where the number of workers participating in the optimization may vary.