FEDERATED LEARNING’S BLESSING: FEDAVG
Year: 2020

Abstract:
Federated learning (FL) learns a model jointly from a set of participating devices without sharing each other’s privately held data. The characteristics of non-i.i.d. data across the network, low device participation, high communication costs, and the mandate that data remain private bring challenges in understanding the convergence of FL algorithms, particularly in regards to how convergence scales with the number of participating devices. In this paper, we focus on Federated Averaging (FedAvg)–arguably the most popular and effective FL algorithm class in use today–and provide a unified and comprehensive study of its convergence rate. Although FedAvg has recently been studied by an emerging line of literature, it remains open as to how FedAvg’s convergence scales with the number of participating devices in the fully heterogeneous FL setting–a crucial question whose answer would shed light on the performance of FedAvg in large FL systems. We fill this gap by providing a unified analysis that establishes convergence guarantees for FedAvg under three classes of problems: strongly convex smooth, convex smooth, and overparameterized strongly convex smooth problems. We show that FedAvg enjoys linear speedup in each case, although with different convergence rates and communication efficiencies. While there have been linear speedup results from distributed optimization that assumes full participation, ours are the first to establish linear speedup for FedAvg under both statistical and system heterogeneity. For strongly convex and convex problems, we also characterize the corresponding convergence rates for the Nesterov accelerated FedAvg algorithm, which are the first linear speedup guarantees for momentum variants of FedAvg in the convex setting. To provably accelerate FedAvg, we design a new momentum-based FL algorithm that further improves the convergence rate in overparameterized linear regression problems. Empirical studies of the algorithms in various settings have supported our theoretical results.

heading: 1 INTRODUCTION
text: Federated learning (FL) is a machine learning paradigm where many clients (e.g., mobile devices or organizations) collaboratively train a model under the orchestration of a central server (e.g., service provider), while keeping the training data decentralized (Smith et al. (2017); Kairouz et al. (2019)). In recent years, FL has swiftly emerged as an important learning paradigm (McMahan et al. (2017); Li et al. (2020a))–one that enjoys widespread success in applications such as personalized recommendation (Chen et al. (2018)), virtual assistant (Lam et al. (2019)), and keyboard prediction (Hard et al. (2018)), to name a few–for at least three reasons: First, the rapid proliferation of smart devices that are equipped with both computing power and data-capturing capabilities provided the infrastructure core for FL. Second, the rising awareness of privacy and the explosive growth of computational power in mobile devices have made it increasingly attractive to push the computation to the edge. Third, the empirical success of communication-efficient FL algorithms has enabled increasingly larger-scale parallel computing and learning with less communication overhead.
Despite its promise and broad applicability in our current era, the potential value FL delivers is coupled with the unique challenges it brings forth. In particular, when FL learns a single statistical model using data from across all the devices while keeping each individual device’s data isolated (Kairouz et al. (2019)), it faces two challenges that are absent in centralized optimization and distributed (stochastic) optimization (Zhou & Cong (2018); Stich (2019); Khaled et al. (2019); Liang et al. (2019); Wang & Joshi (2018); Woodworth et al. (2018); Wang et al. (2019); Jiang & Agrawal (2018); Yu et al. (2019b;a); Khaled et al. (2020); Koloskova et al. (2020); Woodworth et al. (2020b;a)):
1) Data (statistical) heterogeneity: data distributions in devices are different (and data cannot be shared);
κ1κ̃ are condition numbers defined in Section G. Since κ1 ≥ κ̃, this implies a speedup factor of√
κ1 κ̃ for accelerated FedAvg.
2) System heterogeneity: only a subset of devices may access the central server at each time both because the communications bandwidth profiles vary across devices and because there is no central server that has control over when a device is active (the presence of “stragglers”).
To address these challenges, Federated Averaging (FedAvg) McMahan et al. (2017) was proposed as a particularly effective heuristic, which has enjoyed great empirical success. This success has since motivated a growing line of research efforts into understanding its theoretical convergence guarantees in various settings. For instance, Haddadpour & Mahdavi (2019) analyzed FedAvg (for non-convex smooth problems satisfying PL conditions) under the assumption that each local device’s minimizer is the same as the minimizer of the joint problem (if all devices’ data is aggregated together), an overly restrictive assumption that restricts the extent of data heterogeneity. Very recently, Li et al. (2020b) furthered the progress and established an O( 1T ) convergence rate for FedAvg for strongly convex smooth Federated learning problems with both data and system heterogeneity. A similar result in the same setting Karimireddy et al. (2019) also established an O( 1T ) result that allows for a linear speedup when the number of participating devices is large. At the same time, Huo et al. (2020) studied the Nesterov accelerated FedAvg for non-convex smooth problems and established anO( 1√
T )
convergence rate to stationary points.
However, despite these very recent fruitful pioneering efforts into understanding the theoretical convergence properties of FedAvg, it remains open as to how the number of devices–particularly the number of devices that participate in the computation–affects the convergence speed. In particular, is linear speedup of FedAvg a universal phenomenon across different settings and for any number of devices? What about when FedAvg is accelerated with momentum updates? Does the presence of both data and system heterogeneity in FL imply different communication complexities and require technical novelties over results in distributed and decentralized optimization? These aspects are currently unexplored or underexplored in FL. We fill in the gaps here by providing affirmative answers.
Our Contributions We provide a comprehensive and unified convergence analysis of FedAvg and its accelerated variants in the presence of both data and system heterogeneity. Our contributions are threefold.
First, we establish an O(1/KT ) convergence rate under FedAvg for strongly convex and smooth problems and an O(1/ √ KT ) convergence rate for convex and smooth problems (where K is the number of participating devices), thereby establishing that FedAvg enjoys the desirable linear speedup property in the FL setup. Prior to our work here, the best and the most related convergence analysis is given by Li et al. (2020b) and Karimireddy et al. (2019), which established an O( 1T ) convergence rate for strongly convex smooth problems under FedAvg. Our rate matches the same (and optimal) dependence on T , but also completes the picture by establishing the linear dependence on K, for any K ≤ N , where N is the total number of devices, whereas Li et al. (2020b) does not have linear speedup analysis, and Karimireddy et al. (2019) only allows linear speedup close to full participation (K = O(N)). As for convex and smooth problems, there was no prior work that established the O( 1√
T ) rate under both system and data heterogeneity. Our unified analysis highlights the common
elements and distinctions between the strongly and convex settings.
Second, we establish the same convergence rates–O(1/KT ) for strongly convex and smooth problems and O(1/ √ KT ) for convex and smooth problems–for Nesterov accelerated FedAvg. We analyze the accelerated version of FedAvg here because empirically it tends to perform better; yet, its theoretical
convergence guarantee is unknown. To the best of our knowledge, these are the first results that provide a linear speedup characterization of Nesterov accelerated FedAvg in those two problem classes (that FedAvg and Nesterov accelerated FedAvg share the same convergence rate is to be expected: this is the case even for centralized stochastic optimization). Prior to our results here, the most relevant results Yu et al. (2019a); Li et al. (2020a); Huo et al. (2020) only concern the non-convex setting, where convergence is measured with respect to stationary points (vanishing of gradient norms, rather than optimality gaps). Our unified analysis of Nesterov FedAvg also illustrates the technical similarities and distinctions compared to the original FedAvg algorithm, whereas prior works (in the non-convex setting) were scattered and used different notations.
Third, we study a subclass of strongly convex smooth problems where the objective is overparameterized and establish a faster O(exp(−KTκ )) convergence rate for FedAvg, in contrast to the O(exp(−Tκ )) rate for individual solvers Ma et al. (2018). Within this class, we further consider the linear regression problem and establish an even sharper rate under FedAvg. In addition, we propose a new variant of accelerated FedAvg based on a momentum update of Liu & Belkin (2020)–MaSS accelerated FedAvg–and establish a faster convergence rate (compared to if no acceleration is used). This stands in contrast to generic (strongly) convex stochastic problems where theoretically no rate improvement is obtained when one accelerates FedAvg. The detailed convergence results are summarized in Table 1.
Connections with Distributed and Decentralized Optimization Federated learning is closely related to distributed and decentralized optimization, and as such it is important to discuss connections and distinctions between our work and related results from that literature. First, when there is neither system heterogeneity, i.e. all devices participate in parameter averaging during a communication round, nor statistical heterogeneity, i.e. all devices have access to a common set of stochastic gradients, FedAvg coincides with the “Local SGD” of Stich (2019), which showed the linear speedup rate O(1/NT ) for strongly convex and smooth functions. Woodworth et al. (2020b) and Woodworth et al. (2020a) further improved the communication complexity that guarantees the linear speedup rate. When there is only data heterogeneity, some works have continued to use the term Local SGD to refer to FedAvg, while others subsume it in more general frameworks that include decentralized model averaging based on a network topology or a mixing matrix. They have provided linear speedup analyses for strongly convex and convex problems, e.g. Khaled et al. (2020); Koloskova et al. (2020) as well as non-convex problems, e.g. Jiang & Agrawal (2018); Yu et al. (2019b); Wang & Joshi (2018). However, these results do not consider system heterogeneity, i.e. the presence of stragglers in the device network. Even with decentralized model averaging, the assumptions usually imply that model averages over all devices is the same as decentralized model averages based on network topology (e.g. Koloskova et al. (2020) Proposition 1), which precludes system heterogeneity as defined in this paper and prevalent in FL problems. For momentum accelerated FedAvg, Yu et al. (2019a) provided linear speedup analysis for non-convex problems, while results for strongly convex and convex settings are entirely lacking, even without system heterogeneity. Karimireddy et al. (2019) considers both types of heterogeneities for FedAvg, but their rate implies a linear speedup only when the number of stragglers is negligible. In contrast, our linear speedup analyses consider both types of heterogeneity present in the full federated learning setting, and are valid for any number of participating devices. We also highlight a distinction in communication efficiency when system heterogeneity is present. Moreover, our results for Nesterov accelerated FedAvg completes the picture for strongly convex and convex problems. For a detailed comparison with related works, please refer to Table 2 in Appendix Section B.

heading: 2 SETUP
text: In this paper, we study the following federated learning problem:
min w
{ F (w) , ∑N k=1 pkFk(w) } , (1)
where N is the number of local devices (users/nodes/workers) and pk is the k-th device’s weight satisfying pk ≥ 0 and ∑N k=1 pk = 1. In the k-th local device, there are nk data points: x 1 k,x 2 k, . . . ,x nk k .
The local objective Fk(·) is defined as: Fk(w) , 1nk ∑nk j=1 ` ( w;xjk ) , where ` denotes a userspecified loss function. Each device only has access to its local data, which gives rise to its own
local objective Fk. Note that we do not make any assumptions on the data distributions of each local device. The local minimum F ∗k = minw Fk(w) can be far from the global minimum of Eq (1) (data heterogeneity).

heading: 2.1 THE FEDERATED AVERAGING (FEDAVG) ALGORITHM
text: We first introduce the standard Federated Averaging (FedAvg) algorithm which was first proposed by McMahan et al. (2017). FedAvg updates the model in each device by local Stochastic Gradient Descent (SGD) and sends the latest model to the central server every E steps. The central server conducts a weighted average over the model parameters received from active devices and broadcasts the latest averaged model to all devices. Formally, the updates of FedAvg at round t is described as follows:
vkt+1 = w k t − αtgt,k, wkt+1 = { vkt+1 if t+ 1 /∈ IE ,∑ k∈St+1 qkv k t+1 if t+ 1 ∈ IE ,
(2)
where wkt is the local model parameter maintained in the k-th device at the t-th iteration and gt,k := ∇Fk(wkt , ξkt ) is the stochastic gradient based on ξkt , the data point sampled from k-th device’s local data uniformly at random. IE = {E, 2E, . . . } is the set of global communication steps, when local parameters from a set of active devices are averaged and broadcast to all devices. We use St+1 to represent the (random) set of active devices at t+ 1. qk is a set of averaging weights that are specific to the sampling procedure used to obtain the set of active devices St+1. Since federated learning usually involves an enormous amount of local devices, it is often more realistic to assume only a subset of local devices is active at each communication round (system heterogeneity). In this work, we consider both the case of full participation where the model is averaged over all devices at each communication round, in which case qk = pk for all k and wkt+1 = ∑N k=1 pkv k t+1 if t+ 1 ∈ IE , and the case of partial participation where |St+1| < N .
With partial participation, we follow Li et al. (2020a); Karimireddy et al. (2019); Li et al. (2020b) and assume that St+1 is obtained by one of two types of sampling schemes to simulate practical scenarios. One scheme establishes St+1 by i.i.d. sampling the devices with probability pk with replacement, and uses qk = 1K , where K = |St+1|, while the other scheme samples St+1 uniformly i.i.d. from all devices without replacement, and uses qk = pk NK . Both schemes guarantee that gradient updates in FedAvg are unbiased stochastic versions of updates in FedAvg with full participation, which is important in the theoretical analysis of convergence. Because the original sampling scheme and weights proposed by McMahan et al. (2017) lacks this nice property, it is not considered in this paper. For more details on the notations and setup as well as properties of the two sampling schemes, please refer to Section A in the appendix.

heading: 2.2 ASSUMPTIONS
text: We make the following standard assumptions on the objective function F1, . . . , FN . Assumptions 1 and 2 are commonly satisfied by a range of popular objective functions, such as `2-regularized logistic regression and cross-entropy loss functions. Assumption 1 (L-smooth). F1, · · · , FN are all L-smooth: for all v and w, Fk(v) ≤ Fk(w)+(v− w)T∇Fk(w) + L2 ‖v −w‖ 2 2. Assumption 2 (Strongly-convex). F1, · · · , FN are all µ -strongly convex: for all v and w, Fk(v) ≥ Fk(w) + (v −w)T∇Fk(w) + µ2 ‖v −w‖ 2 2 Assumption 3 (Bounded local variance). Let ξkt be sampled from the k-th device’s local data uniformly at random. The variance of stochastic gradients in each device is bounded: E ∥∥∇Fk (wkt , ξkt )−∇Fk (wkt )∥∥2 ≤ σ2k, for k = 1, · · · , N and any wkt . Let σ2 := ∑Nk=1 pkσ2k. Assumption 4 (Bounded local gradient). The expected squared norm of stochastic gradients is uniformly bounded. i.e., E
∥∥∇Fk (wkt , ξkt )∥∥2 ≤ G2, for all k = 1, ..., N and t = 0, . . . , T − 1. Assumptions 3 and 4 have been made in many previous works in federated learning, e.g. Yu et al. (2019b); Li et al. (2020b); Stich (2019). We provide further justification for their generality. As model average parameters become closer to w∗, the L-smoothness property implies that E‖∇Fk(wkt , ξkt )‖2 and E‖∇Fk(wkt , ξkt )−∇Fk(wkt )‖2 approach E‖∇Fk(w∗, ξkt )‖2 and
E‖∇Fk(w∗, ξkt )−∇Fk(w∗)‖2. Therefore, there is no substantial difference between these assumptions and assuming the bounds at w∗ only Koloskova et al. (2020). Furthermore, compared to assuming bounded gradient diversity as in related work Haddadpour & Mahdavi (2019); Li et al. (2020a), Assumption 4 is much less restrictive. When the optimality gap converges to zero, bounded gradient diversity restricts local objectives to have the same minimizer as the global objective, contradicting the heterogeneous data setting. For detailed discussions of our assumptions, please refer to Appendix Section B.

heading: 3 LINEAR SPEEDUP ANALYSIS OF FEDAVG
text: In this section, we provide convergence analyses of FedAvg for convex objectives in the general setting with both heterogeneous data (statistical heterogeneity) and partial participation (system heterogeneity). We show that for strongly convex and smooth objectives, the convergence of the optimality gap of averaged parameters across devices is O(1/KT ), while for convex and smooth objectives, the rate is O(1/ √ KT ). Our results improve upon Li et al. (2020b); Karimireddy et al. (2019) by showing linear speedup for any number of participating devices, and upon Khaled et al. (2020); Koloskova et al. (2020) by allowing system heterogeneity. The proofs also highlight similarities and distinctions between the strongly convex and convex settings. Detailed proofs are deferred to Appendix Section E.

heading: 3.1 STRONGLY CONVEX AND SMOOTH OBJECTIVES
text: We first show that FedAvg has an O(1/KT ) convergence rate for µ-strongly convex and L-smooth objectives. The result relies on a technical improvement over the analysis in Li et al. (2020b). Moreover, it implies a distinction in communication efficiency that guarantees this linear speedup for FedAvg with full and partial device participation. With full participation, E can be chosen as large as O( √ T/N) without degrading the linear speedup in the number of workers. On the other hand, with partial participation, E must be O(1) to guarantee O(1/KT ) convergence.
Theorem 1. Let wT = ∑N k=1 pkw k T in FedAvg, νmax = maxkNpk, and set decaying learning rates αt = 4µ(γ+t) with γ = max{32κ,E} and κ = L µ . Then under Assumptions 1 to 4 with full device participation,
EF (wT )− F ∗ = O ( κν2maxσ 2/µ
NT + κ2E2G2/µ T 2
) ,
and with partial device participation with at most K sampled devices at each communication round, EF (wT )− F ∗ = O ( κE2G2/µ
KT + κν2maxσ 2/µ NT + κ2E2G2/µ T 2
) .
Proof sketch. Because our unified analyses of results in the main text follow the same framework with variations in technical details, we first give an outline of proof for Theorem 1 to illustrate the main ideas. For full participation, the main ingredient is a recursive contraction bound
E‖wt+1 −w∗‖2 ≤ (1− µαt)E‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6α3tLE 2G2
where the O(α3tE2G2) term is the key improvement over the bound in Li et al. (2020b), which has O(α2tE2G2) instead. We then use induction to obtain a non-recursive bound on E‖wT − w∗‖2, which is converted to a bound on EF (wT ) − F ∗ using L-smoothness. For partial participation, an additional term O( 1Kα 2 tE
2G2) of leading order resulting from sampling variance is added to the contraction bound. To facilitate the understanding of our analysis, please refer to a high-level summary in Appendix C.
Linear speedup. We compare our bound with that in Li et al. (2020b), which isO( 1NT + E2 KT + E2G2
T ). Because the term E 2G2
T is also O(1/T ) without a dependence on N , for any choice of E their bound cannot achieve linear speedup. The improvement of our bound comes from the term κ
2E2G2/µ T 2 , which
now is O(E2/T 2) and so is not of leading order. As a result, all leading terms scale with 1/N in the full device participation setting, and with 1/K in the partial participation setting. This implies that in both settings, there is a linear speedup in the number of active workers during a communication round. We also emphasize that the reason one cannot recover the full participation bound by setting K = N in the partial participation bound is due to the variance generated by sampling. Communication Complexity. Our bound implies a distinction in the choice of E between the full and partial participation settings. With full participation, the term involving E, O(E2/T 2), is not of leading order O(1/T ), so we can increase E and reduce the number of communication rounds without degrading the linear speedup in iteration complexity O(1/NT ), as long as E = O( √ T/N), since then O(E2/T 2) = O(1/NT ) matches the leading term. This corresponds to a communication complexity of T/E = O( √ NT ). In contrast, the bound in Li et al. (2020b) does not allow E to scale with √ T to preserve O(1/T ) rate, even for full participation. On the other hand, with partial participation, κE 2G2/µ KT is also a leading term, and so E must be O(1). In this case, our bound still yields a linear speedup in K, which is also confirmed by experiments. The requirement that E = O(1) in order to achieve linear speedup in partial participation cannot be removed for our sampling schemes, as the term κE
2G2/µ KT comes from variance in the sampling process, which is
O(E2/T 2). In Proposition 1 in Section E of the appendix, we provide a problem instance where the dependence of the sampling variance on E is tight. Comparison with related works. To better understand the significance of the obtained bound, we compare our rates to the best-known results in related settings. Haddadpour & Mahdavi (2019) proves a linear speedup O(1/KT ) result for strongly convex and smooth objectives1, with O(K1/3T 2/3) communication complexity with non-i.i.d. data and partial participation. However, their results build on the bounded gradient diversity assumption, which implies the existence of w∗ that minimizes all local objectives (see discussions in Section 2.2 and Appendix B), effectively removing statistical heterogeneity. The bound in Koloskova et al. (2020) matches our bound in the full participation case, but their framework excludes partial participation (Koloskova et al., 2020, Proposition 1). The result of Karimireddy et al. (2019) applies to the full FL setting, but only has linear speedup when K = O(N), i.e. close to full participation, whereas our result has linear speedup for any number of participating devices. When there is no data heterogeneity, i.e. in the classical distributed optimization paradigm, communication complexity can be further improved, e.g. Woodworth et al. (2020b;a), but such results are not directly comparable to ours since we consider the setting where individual devices have access to different datasets.

heading: 3.2 CONVEX SMOOTH OBJECTIVES
text: Next we provide linear speedup analysis of FedAvg with convex and smooth objectives and show that the optimality gap is O(1/ √ KT ). This result complements the strongly convex case in the previous part, as well as the non-convex smooth setting in Jiang & Agrawal (2018); Yu et al. (2019b); Haddadpour & Mahdavi (2019), where O(1/ √ KT ) results are given in terms of averaged gradient norm, and it also extends the result in Khaled et al. (2020), which has linear speedup in the convex setting, but only for full participation.
Theorem 2. Under Assumptions 1,3,4 and constant learning rate αt = O( √ N T ), FedAvg satisfies
min t≤T
F (wt)− F (w∗) = O ( ν2maxσ 2
√ NT
+ NE2LG2
T ) with full participation, and with partial device participation with K sampled devices at each commu-
nication round and learning rate αt = O( √ K T ),
min t≤T
F (wt)− F (w∗) = O ( ν2maxσ 2
√ KT + E2G2√ KT + KE2LG2 T
) .
The analysis again relies on a recursive bound, but without contraction:
E‖wt+1 −w∗‖2 + αt(F (wt)− F (w∗)) ≤ E‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6α3tE 2LG2
1Their result applies to a larger class of non-convex objectives that satisfy the Polyak-Lojasiewicz condition.
which is then summed over time steps to give the desired bound, with αt = O( √ N T ). Choice of E and linear speedup. With full participation, as long as E = O(T 1/4/N3/4), the convergence rate is O(1/ √ NT ) with O(N3/4T 3/4) communication rounds. In the partial participation setting, E must be O(1) in order to achieve linear speedup of O(1/ √ KT ). This is again due to the fact that the sampling variance E‖wt − vt‖2 = O(α2tE2G2) cannot be made independent of E, as illustrated by Proposition 1. See also the proof in Section E for how the sampling variance and the term E2G2/ √ KT are related. Our result again demonstrates the difference in communication complexities between full and partial participation, and is to our knowledge the first result on linear speedup in the general federated learning setting with both heterogeneous data and partial participation for convex objectives.

heading: 4 LINEAR SPEEDUP ANALYSIS OF NESTEROV ACCELERATED FEDAVG
text: A natural extension of the FedAvg algorithm is to use momentum-based local updates instead of local SGD updates in order to accelerate FedAvg. As we know from stochastic optimization, Nesterov and other momentum updates fail to provably accelerate over SGD (Liu & Belkin (2020); Kidambi et al. (2018); Liu et al. (2018); Yuan & Ma (2020)). This is in contrast to the classical acceleration result of Nesterov-accelerated gradient descent over GD. Thus in the FL setting, the best provable convergence rate for FedAvg with Nesterov updates is the same as FedAvg with SGD updates. Nevertheless, Nesterov and other momentum updates are frequently used in practice, in both non-FL and FL settings, and are observed to perform better empirically. In fact, previous works such as Stich (2019) on FedAvg with vanilla SGD uses FedAvg with Nesterov or other momentum updates in their experiments to achieve target accuracy. Because of the popularity of Nesterov and other momentumbased methods, understanding the linear speedup behavior of FedAvg with such local updates is important. To our knowledge, the only convergence analyses of FedAvg with momentum-based stochastic updates focus on the non-convex smooth case Huo et al. (2020); Yu et al. (2019a); Li et al. (2020a), and no results existed in the convex smooth settings. In this section, we complete the picture by providing the first O(1/KT ) and O(1/ √ KT ) convergence results for Nesterov-accelerated FedAvg for convex objectives that match the rates for FedAvg with SGD updates. Detailed proofs of convergence results in this section are deferred to Appendix Section F.

heading: 4.1 STRONGLY CONVEX AND SMOOTH OBJECTIVES
text: The Nesterov Accelerated FedAvg algorithm follows the updates:
vkt+1 = w k t − αtgt,k, wkt+1 =
{ vkt+1 + βt(v k t+1 − vkt ) if t+ 1 /∈ IE ,∑
k∈St+1 qk [ vkt+1 + βt(v k t+1 − vkt ) ] if t+ 1 ∈ IE ,
where gt,k := ∇Fk(wkt , ξkt ) is the stochastic gradient sampled on the k-th device at time t, and qk again depends on participation and sampling schemes.
Theorem 3. Let vT = ∑N k=1 pkv k T in Nesterov accelerated FedAvg, and set learning rates αt = 6 µ 1 t+γ , βt−1 =
3 14(t+γ)(1− 6t+γ )max{µ,1} . Then under Assumptions 1,2,3,4 with full device participation,
EF (vT )− F ∗ = O ( κν2maxσ 2/µ
NT + κ2E2G2/µ T 2
) ,
and with partial device participation with K sampled devices at each communication round, EF (vT )− F ∗ = O ( κν2maxσ 2/µ
NT + κE2G2/µ KT + κ2E2G2/µ T 2
) .
Similar to FedAvg, the key step in the proof of this result is a recursive contraction bound, but different in that it involves three time steps, due to the update format of Nesterov SGD (see Lemma 7 in Appendix F.1). Then we can again use induction and L-smoothness to obtain the desired bound. To our knowledge, this is the first convergence result for Nesterov accelerated FedAvg in the strongly convex and smooth setting. The same discussion about linear speedup of FedAvg applies to the Nesterov accelerated variant. In particular, to achieve O(1/NT ) linear speedup, T iterations of the algorithm require only O( √ NT ) communication rounds with full participation.

heading: 4.2 CONVEX SMOOTH OBJECTIVES
text: We now show that the optimality gap of Nesterov-accelerated FedAvg has O(1/ √ KT ) rate for convex and smooth objectives. This result complements the strongly convex case in the previous part, as well as the non-convex smooth setting in Huo et al. (2020); Yu et al. (2019a); Li et al. (2020a), where a similar O(1/ √ KT ) rate is given in terms of averaged gradient norm.
Theorem 4. Set learning rates αt = βt = O( √ N T ). Then under Assumptions 1,3,4 Nesterov accelerated FedAvg with full device participation has rate
min t≤T
F (vt)− F ∗ = O ( ν2maxσ 2
√ NT
+ NE2LG2
T
) ,
and with partial device participation with K sampled devices at each communication round,
min t≤T
F (vt)− F ∗ = O ( ν2maxσ 2
√ KT + E2G2√ KT + KE2LG2 T
) .
We emphasize again that in the stochastic optimization setting, the optimal convergence rate that FedAvg wtih Nesterov udpates can achieve is the same as FedAvg with SGD updates. However, due to the popularity and superior performance of momentum methods in practice, it is still important to understand the linear speedup behavior of such FedAvg variants. Our results in this section fill exactly this gap.

heading: 5 NUMERICAL EXPERIMENTS
text: In this section, we empirically examine the linear speedup convergence of FedAvg and Nesterov accelerated FedAvg in various settings, including strongly convex function, convex smooth function, and overparameterized objectives, as analyzed in previous sections.
Setup. Following the experimental setting in Stich (2019), we conduct experiments on both synthetic datasets and real-world dataset w8a Platt (1998) (d = 300, n = 49749). We consider the distributed objectives F (w) = ∑N k=1 pkFk(w), and the objective function on the k-th local device includes three cases: 1) Strongly convex objective: the regularized binary logistic regression problem, Fk(w) =
1 Nk ∑Nk i=1 log(1 + exp(−ykiwTxki ) + λ 2 ‖w‖
2. The regularization parameter is set to λ = 1/n ≈ 2e − 5. 2) Convex smooth objective: the binary logistic regression problem without regularization. 3) Overparameterized setting: the linear regression problem without adding noise to the label, Fk(w) = 1Nk ∑Nk i=1(w Txki + b− yki )2.
Linear speedup of FedAvg and Nesterov accelerated FedAvg. To verify the linear speedup convergence as shown in Theorems 1 2 3 4, we evaluate the number of iterations needed to reach -accuracy in three objectives. We initialize all runs with w0 = 0d and measure the number of iterations to reach the target accuracy . For each configuration (E,K), we extensively search the learning rate from min(η0, nc1+t ), where η0 ∈ {0.1, 0.12, 1, 32} according to different problems and c can take the values c = 2i ∀i ∈ Z. As the results shown in Figure 1, the number of iterations decreases as the number of (active) workers increasing, which is consistent for FedAvg and Nesterov accelerated FedAvg across all scenarios. For additional experiments on the impact of E, detailed experimental setup, and hyperparameter setting, please refer to the Appendix Section I.

heading: A ADDITIONAL NOTATIONS AND BOUNDS FOR SAMPLING SCHEMES
text: In this section, we introduce additional notations that are used throughout the proofs. Following common practice, e.g. Stich (2019); Li et al. (2020b), we define two virtual sequences vt = ∑N k=1 pkv k t
and wt = ∑N k=1 pkw k t , where we recall the FedAvg updates from (2):
vkt+1 = w k t − αtgt,k, wkt+1 = { vkt+1 if t+ 1 /∈ IE ,∑ k∈St+1 qkv k t+1 if t+ 1 ∈ IE .
The following observations apply to FedAvg updates, while Nesterov accelerated FedAvg requires modifications. For full device participation or partial participation with t /∈ IE , note that vt = wt = ∑N k=1 pkv k t . For partial participation with t ∈ IE , wt 6= vt since vt = ∑N k=1 pkv k t while
wt = ∑ k∈St qkw k t . However, we can use unbiased sampling strategies such that EStwt = vt. Note that vt+1 is one-step SGD from wt.
vt+1 = wt − αtgt, (3) where gt = ∑N k=1 pkgt,k is the one-step stochastic gradient averaged over all devices.
gt,k = ∇Fk ( wkt , ξ k t ) ,
Similarly, we denote the expected one-step gradient gt = Eξt [gt] = ∑N k=1 pkEξkt gt,k, where
Eξkt gt,k = ∇Fk ( wkt ) , (4)
and ξt = {ξkt }Nk=1 denotes random samples at all devices at time step t. Since in this work we also consider the case of partial participation, the sampling strategy to approximate the system heterogeneity can also affect the convergence. Here we follow the prior works Li et al. (2020b) and Li et al. (2020a) and consider two types of sampling schemes that guarantee EStwt = vt. The sampling scheme I establishes St+1 by i.i.d. sampling the devices according to probabilities pk with replacement, and setting qk = 1K . In this case the upper bound of expected square norm of wt+1 − vt+1 is given by (Li et al., 2020b, Lemma 5):
ESt+1 ‖wt+1 − vt+1‖ 2 ≤ 4
K α2tE 2G2. (5)
The sampling scheme II establishes St+1 by uniformly sampling all devices without replacement and setting qk = pk NK , in which case we have
ESt+1 ‖wt+1 − vt+1‖ 2 ≤ 4(N −K)
K(N − 1) α2tE 2G2. (6)
We summarize these upper bounds as follows:
ESt+1 ‖wt+1 − vt+1‖ 2 ≤ 4
K α2tE 2G2. (7)
and this bound will be used in the convergence proof of the partial participation result.

heading: B COMPARISON OF CONVERGENCE RATES WITH RELATED WORKS
text: In this section, we compare our convergence rate with the best-known results in the literature (see Table 2). In Haddadpour & Mahdavi (2019), the authors provide O(1/NT ) convergence rate of nonconvex problems under Polyak-Łojasiewicz (PL) condition, which means their results can directly apply to the strongly convex problems. However, their assumption is based on bounded gradient diversity, defined as follows:
Λ(w) =
∑ k pk‖∇Fk(w)‖22
‖ ∑ k pk∇Fk(w)‖22 ≤ B
This is a more restrictive assumption comparing to assuming bounded gradient under the case of target accuracy → 0 and PL condition. To see this, consider the gradient diversity at the global optimal
w∗, i.e., Λ(w∗) = ∑ k pk‖∇Fk(w)‖ 2 2
‖ ∑ k pk∇Fk(w)‖22
. For Λ(w∗) to be bounded, it requires ‖∇Fk(w∗)‖22 = 0, ∀ k. This indicates w∗ is also the minimizer of each local objective, which contradicts to the practical setting of heterogeneous data. Therefore, their bound is not effective for arbitrary small -accuracy under general heterogeneous data while our convergence results still hold in this case.
In Karimireddy et al. (2019), the linear speedup convergence rate of FedAvg are provided for strongly convex, general convex, and non-convex problems under full participation setting. However, their rate does not enjoy linear speedup for any number of devices while our results apply to any valid K ≤ N . For example, they provides an optimality gap of O ( (1− KN )E/T ) for the strongly convex case (Karimireddy et al., 2019, Theorem V). With partial participation, and when K = O(1), their convergence rate is O(E/T ) which does not have linear speedup. Under partial participation, the FedAvg analyses in Karimireddy et al. (2019) requires E = O(1). For example, the strongly convex result O((1− KN )E/T ) in Theorem V is O(E/T ) when K = O(1) and is O(E/NT ) when K = O(N). In either case, to achieve a O(1/T ) convergence rate, it requires E = O(1) as well. Similar conclusion also holds for the general convex problem.

heading: C A HIGH-LEVEL SUMMARY OF FEDAVG ANALYSIS
text: To facilitate the understanding of our analysis and highlight the improvement of our work comparing to prior arts, we summarize the general steps used in the proofs across the various settings. In this section, we take the strongly convex case as an example to illustrate our analysis. The corresponding proof for general convex functions follows the same framework.

heading: One step progress bound
text: This step establishes the progress of distance (‖wt −w∗‖2) to optimal solution after one step SGD update (see line 9, Alg 1), as the following equation shows:
E‖wt+1 −w∗‖2 ≤ O(ηtE‖wt −w∗‖2 + α2tσ2/N + α3tE2G2).
Algorithm 1 FEDAVG: Federated Averaging 1: Server input: initial model w0, initial step size α0, local steps E. 2: Client input: 3: for each round r = 0, 1, ..., R, where r = t ∗ E do 4: Sample clients St ⊆ {1, ..., N} 5: Broadcast w to all clients k ∈ St 6: for each client k ⊆ St do 7: initialize local model wkt = w 8: for t = r ∗ E + 1, . . . , (r + 1) ∗ E do 9: wkt+1 = wkt − αtgt,k 10: end for 11: end for 12: Average the local models at server end: wt = ∑ k∈St w k t . 13: end for
The above bound consists of three main ingredients, the distance to optima in previous step (with ηt ∈ (0, 1) to obtained a contraction bound), the variance of stochastic gradients in local clients (second term), the variance across different clients (third term). Notice that the third term in this bound is the primary source of improvement in the rate. Comparing to the bound in Li et al. (2020b), we improve the third term from O(α2tE2G2) to O(α2tE2G2), which enables the linear speedup in the convergence rate.

heading: Iterative deduction
text: This step uses the one step progress bound iteratively to connect the the current distance to optimal solution with the initial distance (‖w0 −w∗‖2), as follows:
E‖wt+1 −w∗‖2 ≤ O(E‖w0 −w∗‖2 1
T ).
Then we can use the distance to optima to upper bound the optimality gap (F (wt)− F ∗ ≤ O(1/T )), as follows:
E(F (wt))− F ∗ ≤ O(E‖wt −w∗‖2).
The convergence rate of the optimality gap is equally obtained as the convergence rate of the distance to optima.
From full participation to partial participation There are three sources of variances that affect the convergence rate. The first two sources come from the variances of within local clients and across clients (second and third term in one step progress bound). The partial participation, which involves a sampling procedure, is the third source of variance. Therefore, comparing to the rate in full participation, this will add another term of variance into the convergence rate, where we follow a similar derivation as in Li et al. (2020b).

heading: D TECHNICAL LEMMAS
text: To facilitate reading, we first summarize some basic properties of L-smooth and µ-strongly convex functions, found in e.g. Rockafellar (1970), which are used in various steps of proofs in the appendix. Lemma 1. Let F be a convex L-smooth function. Then we have the following inequalities:
1. Quadratic upper bound: 0 ≤ F (w)− F (w′)− 〈∇F (w′),w −w′〉 ≤ L2 ‖w −w ′‖2.

heading: 2. Coercivity: 1L‖∇F (w)−∇F (w
text: ′)‖2 ≤ 〈∇F (w)−∇F (w′),w −w′〉.
3. Lower bound: F (w) ≥ F (w′) + 〈∇F (w′),w−w′〉+ 12L‖∇F (w)−∇F (w ′)‖2. In particular, ‖∇F (w)‖2 ≤ 2L(F (w)− F (w∗)). 4. Optimality gap: F (w)− F (w∗) ≤〈∇F (w),w −w∗〉. Lemma 2. Let F be a µ-strongly convex function. Then
F (w) ≤ F (w′) + 〈∇F (w′),w −w′〉+ 1 2µ ‖∇F (w)−∇F (w′)‖2
F (w)− F (w∗) ≤ 1 2µ ‖∇F (w)‖2

heading: E PROOF OF CONVERGENCE RESULTS FOR FEDAVG

heading: E.1 STRONGLY CONVEX SMOOTH OBJECTIVES
text: To organize our proofs more effectively and highlight the significance of our results compared to prior works, we first state the following key lemmas used in proofs of main results and defer their proofs to later.
Lemma 3 (One step progress, strongly convex). Let wt = ∑N k=1 pkw k t , and suppose our functions satisfy Assumptions 1,2,3,4, and set step size αt = 4µ(γ+t) with γ = max{32κ,E} and κ = L µ , then the updates of FedAvg with full participation satisfy
E‖wt+1 −w∗‖2 ≤ (1− µαt)E‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6E2Lα3tG 2.
We emphasize that the above lemma is the key step that allows us to obtain a bound that improves on the convergence result of Li et al. (2020b) with linear speedup. Its proof will make use of the following two results. Lemma 4 (Bounding gradient variance (Lemma 2 Li et al. (2020b)) ). Given Assumption 3, the upper bound of gradient variance is given as follows,
E‖gt − gt‖2 ≤ N∑ k=1 p2kσ 2 k.
Lemma 5 (Bounding the divergence of wkt (Lemma 3 Li et al. (2020b)) ). Given Assumption 4, and assume that αt is non-increasing and αt ≤ 2αt+E for all t ≥ 0, we have
E [ N∑ k=1 pk‖wt −wkt ‖2 ] ≤ 4E2α2tG2.
We now restate Theorem 1 from the main text and then prove it using Lemma 3. Theorem 1. Let wT = ∑N k=1 pkw k T in FedAvg, νmax = maxkNpk, and set decaying learning rates αt = 4µ(γ+t) with γ = max{32κ,E} and κ = L µ . Then under Assumptions 1,2,3,4 with full device participation,
EF (wT )− F ∗ = O ( κν2maxσ 2/µ
NT + κ2E2G2/µ T 2 ) and with partial device participation with at most K sampled devices at each communication round,
EF (wT )− F ∗ = O ( κE2G2/µ
KT + κν2maxσ 2/µ NT + κ2E2G2/µ T 2 ) Proof. The road map of the proof for full device participation contains three steps. First, we establish a recursive relationship between E‖wt+1−w∗‖2 and E‖wt−w∗‖2, upper bounding the progress of FedAvg from step t to step t+ 1. Second, we show that E‖wt −w∗‖2 = O(ν 2 maxσ 2/µ tN + E2LG2/µ2
t2 ) by induction using the recursive relationship from the previous step. Third, we use the property of L-smoothness to bound the optimality gap by E‖wt −w∗‖2. By Lemma 3, we have the following upper bound for the one step progress:
E‖wt+1 −w∗‖2 ≤ (1− µαt)E‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6E2Lα3tG 2.
We show next that E‖wt − w∗‖2 = O(ν 2 maxσ 2/µ tN + E2LG2/µ2
t2 ) using induction. To simplify the presentation, we denote C ≡ 6E2LG2 and D ≡ 1N ν 2 maxσ
2. Suppose that we have the bound E‖wt − w∗‖2 ≤ b · (αtD + α2tC) for some constant b and learning rates αt. Then the one step progress from Lemma 3 becomes:
E‖wt+1 −w∗‖2 ≤ (b(1− µαt) + αt)αtD + (b(1− µαt) + αt)α2tC To establish the result at step t+1, it remains to choose αt and b such that (b(1−µαt)+αt)αt ≤ bαt+1 and (b(1 − µαt) + αt)α2t ≤ bα2t+1. If we let αt = 4µ(t+γ) where γ = max{E, 32κ} (choice of γ required to guarantee the one step progress) and set b = 4µ , we have:
(b(1− µαt) + αt)αt = ( b(1− 4
t+ γ ) +
4
µ(t+ γ)
) 4
µ(t+ γ) ≤ b 4 µ(t+ γ + 1) = bαt+1
(b(1− µαt) + αt)α2t = b( t+ γ − 2 t+ γ ) 16 µ2(t+ γ)2 ≤ b 16 µ2(t+ γ + 1)2 = bα2t+1
where we have used the following inequalities: t+ γ − 1 (t+ γ)2 ≤ 1 (t+ γ + 1) t+ γ − 2 (t+ γ)3 ≤ 1 (t+ γ + 1)2
∀ γ ≥ 1
Thus we have established the result at step t+ 1 assuming the result is correct at step t:
E‖wt+1 −w∗‖2 ≤ b · (αt+1D + α2t+1C)
At step t = 0, we can ensure the following inequality by scaling b with c‖w0−w∗‖2 for a sufficiently large constant c:
‖w0 −w∗‖2 ≤ b · (α0D + α20C) = b · ( 4
µγ D +
16
µ2γ2 C)
It follows that
E‖wt −w∗‖2 ≤ c‖w0 −w∗‖2 4
µ (Dαt + Cα
2 t ) (8)
for all t ≥ 0. Finally, the L-smoothness of F implies
E(F (wT ))− F ∗ ≤ L
2 E‖wT −w∗‖2
≤ L 2 c‖w0 −w∗‖2 4 µ (DαT + Cα 2 T )
= 2c‖w0 −w∗‖2κ(DαT + Cα2T ) ≤ 2c‖w0 −w∗‖2κ [ 4
µ(T + γ) · 1 N ν2maxσ 2 + 6E2LG2 · ( 4 µ(T + γ)
)2 ]
= O(κ µ 1 N ν2maxσ 2 · 1 T + κ2 µ E2G2 · 1 T 2 )
where in the first line, we use the property of L-smooth function (see Lemma 1), and in the second line, we use the conclusion in Eq (8).
With partial participation, the update at each communication round is now given by weighted averages over a subset of sampled devices. When t+ 1 /∈ IE , vt+1 = wt+1, while when t+ 1 ∈ IE , we have Ewt+1 = vt+1 by design of the sampling schemes (Li et al. (2020b), Lemma 4), so that
E‖wt+1 −w∗‖2 = E‖wt+1 − vt+1 + vt+1 −w∗‖2
= E‖wt+1 − vt+1‖2 + E‖vt+1 −w∗‖2
This in particular implies E‖vt −w∗‖2 ≤ E‖wt −w∗‖2 for all t. Since vt = ∑N k=1 pkv k t always averages over all devices, the full participation one step progress result Lemma 3 applied to vt implies
E‖vt+1 −w∗‖2 ≤ E(1− µαt)‖vt −w∗‖2 + 6E2Lα3tG2 + α2t 1
N ν2maxσ 2
≤ E(1− µαt)‖wt −w∗‖2 + 6E2Lα3tG2 + α2t 1
N ν2maxσ 2
The bound for E‖wt+1 − vt+1‖2 for the two sampling schemes we consider is provided in Eq (7), and applying it we can write the one step progress for partial participation as
E‖wt+1 −w∗‖2 ≤ (1− µαt)E‖wt −w∗‖2 + α2t 1
N ν2maxσ
2 + 4
K α2tE 2G2 + 6E2Lα3tG 2, (9)
and the same arguments using induction and L-smoothness as the full participation case implies
EF (wT )− F ∗ = O( κν2maxσ 2/µ NT + κE2G2/µ KT + κ2E2G2/µ T 2 )

heading: E.1.1 DEFERRED PROOFS OF KEY LEMMAS
text: Here we first rewrite the proofs of lemmas 4 and 5 from Li et al. (2020b) with slight modifications for the consistency and completeness of this work, since later we will use modified versions of these results in the convergence proof for Nesterov accelerated FedAvg.

heading: Proof of lemma 4.
text: E‖gt − gt‖2 = E‖gt − Egt‖2 = N∑ k=1 p2k‖gt,k − Egt,k‖2 ≤ N∑ k=1 p2kσ 2 k
Proof of lemma 5. Now we bound E ∑N k=1 pk‖wt −wkt ‖2 following Li et al. (2020b). Since communication is done every E steps, for any t ≥ 0, we can find a t0 ≤ t such that t− t0 ≤ E − 1 and wkt0 = wt0for all k. Moreover, using αt is non-increasing and αt0 ≤ 2αt for any t− t0 ≤ E − 1, we have
E N∑ k=1 pk‖wt −wkt ‖2
=E N∑ k=1 pk‖wkt −wt0 − (wt −wt0)‖2
≤E N∑ k=1 pk‖wkt −wt0‖2
=E N∑ k=1 pk‖wkt −wkt0‖ 2
=E N∑ k=1 pk‖ − t−1∑ i=t0 αigi,k‖2
≤2 N∑ k=1 pkE t−1∑ i=t0 Eα2i ‖gi,k‖2
≤2 N∑ k=1 pkE 2α2t0G 2
≤4E2α2tG2
Based on the results of Lemma 4, 5, we now prove the upper bound of one step SGD progress. This proof improves on the previous work Li et al. (2020b) and is the first to reveal the linear speedup of convergence of FedAvg.
Proof of lemma 3. We have
‖wt+1 −w∗‖2 = ‖(wt − αtgt)−w∗‖2 = ‖(wt − αtgt −w∗)− αt(gt − gt)‖2
= ‖wt −w∗ − αtgt‖2︸ ︷︷ ︸ A1 + 2αt〈wt −w∗ − αtgt,gt − gt〉︸ ︷︷ ︸ A2 +α2t ‖gt − gt‖2︸ ︷︷ ︸ A3
where we denote:
A1 = ‖wt −w∗ − αtgt‖2 A2 = 2αt〈wt −w∗ − αtgt,gt − gt〉 A3 = α 2 t ‖gt − gt‖2
By definition of gt and gt (see Eq (4)), we have EA2 = 0. For A3, we have the following upper bound (see Lemma 4):
α2tE‖gt − gt‖2 ≤ α2t N∑ k=1 p2kσ 2 k
Next we bound A1:
‖wt −w∗ − αtgt‖2 = ‖wt −w∗‖2 + 2〈wt −w∗,−αtgt〉+ ‖αtgt‖2
and we will show that the third term ‖αtgt‖2 can be canceled by an upper bound of the second term, which is one of major improvement comparing to prior art Li et al. (2020b). The upper bound of second term can be derived as follows, using the strong convexity and L-smoothness of Fk:
− 2αt〈wt −w∗,gt〉
=− 2αt N∑ k=1 pk〈wt −w∗,∇Fk(wkt )〉
=− 2αt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉 − 2αt N∑ k=1 pk〈wkt −w∗,∇Fk(wkt )〉
≤ − 2αt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉+ 2αt N∑ k=1 pk(Fk(w ∗)− Fk(wkt ))− αtµ N∑ k=1 pk‖wkt −w∗‖2
≤2αt N∑ k=1 pk [ Fk(w k t )− Fk(wt) + L 2 ‖wt −wkt ‖2 + Fk(w∗)− Fk(wkt ) ] − αtµ‖ N∑ k=1 pkw k t −w∗‖2
=αtL N∑ k=1 pk‖wt −wkt ‖2 + 2αt N∑ k=1 pk [Fk(w ∗)− Fk(wt)]− αtµ‖wt −w∗‖2
We record the bound we have obtained so far, as it will also be used in the proof for convex case:
E‖wt+1 −w∗‖2 ≤E(1− µαt)‖wt −w∗‖2 + αtL N∑ k=1 pk‖wt −wkt ‖2
+ 2αt N∑ k=1 pk [Fk(w ∗)− Fk(wt)] + α2t N∑ k=1 p2kσ 2 k + α 2 t ‖gt‖2 (10)
For the term 2αt ∑N k=1 pk [Fk(w
∗)− Fk(wt)], which is negative, we can ignore it, but this yields a suboptimal bound that fails to provide the desired linear speedup. Instead, we upper bound it using the following derivation:
2αt N∑ k=1 pk [Fk(w ∗)− Fk(wt)]
≤2αt [F (wt+1)− F (wt)] ≤2αtE〈∇F (wt),wt+1 −wt〉+ αtLE‖wt+1 −wt‖2
=− 2α2tE〈∇F (wt),gt〉+ α3tLE‖gt‖2
=− 2α2tE〈∇F (wt),gt〉+ α3tLE‖gt‖2 =− α2t [ ‖∇F (wt)‖2 + ‖gt‖2 − ‖∇F (wt)− gt‖2 ] + α3tLE‖gt‖2
=− α2t [ ‖∇F (wt)‖2 + ‖gt‖2 − ‖∇F (wt)− ∑ k pk∇F (wkt )‖2 ] + α3tLE‖gt‖2
≤− α2t [ ‖∇F (wt)‖2 + ‖gt‖2 − ∑ k pk‖∇F (wt)−∇F (wkt )‖2 ] + α3tLE‖gt‖2
≤− α2t [ ‖∇F (wt)‖2 + ‖gt‖2 − L2 ∑ k pk‖wt −wkt ‖2 ] + α3tLE‖gt‖2
≤− α2t ‖gt‖2 + α2tL2 ∑ k pk‖wt −wkt ‖2 + α3tLE‖gt‖2 − α2t ‖∇F (wt)‖2
where we have used the smoothness of F twice.
Note that the term −α2t ‖gt‖2 exactly cancels the α2t ‖gt‖2 in the bound in Eq (10), so that plugging in the bound for −2αt〈wt −w∗,gt〉, we have so far proved
E‖wt+1 −w∗‖2 ≤ E(1− µαt)‖wt −w∗‖2 + αtL N∑ k=1 pk‖wt −wkt ‖2 + α2t N∑ k=1 p2kσ 2 k
+ α2tL 2 N∑ k=1 pk‖wt −wkt ‖2 + α3tLE‖gt‖2 − α2t ‖∇F (wt)‖2 (11)
Under Assumption 4, we have E‖gt‖2 ≤ G2. Furthermore, we can check that our choice of αt satisfies αt is non-increasing and αt ≤ 2αt+E , so we may plug in the bound E ∑N k=1 pk‖wt − wkt ‖2 ≤ 4E2α2tG2 to the above inequality (see Lemma 5). Therefore, we can conclude that, with νmax := N ·maxk pk and νmin := N ·mink pk,
E‖wt+1 −w∗‖2
≤E(1− µαt)‖wt −w∗‖2 + 4E2Lα3tG2 + 4E2L2α4tG2 + α2t N∑ k=1 p2kσ 2 k + α 3 tLG 2
=E(1− µαt)‖wt −w∗‖2 + 4E2Lα3tG2 + 4E2L2α4tG2 + α2t 1
N2 N∑ k=1 (pkN) 2σ2k + α 3 tLG 2
≤E(1− µαt)‖wt −w∗‖2 + 4E2Lα3tG2 + 4E2L2α4tG2 + α2t 1
N2 ν2max N∑ k=1 σ2k + α 3 tLG 2
≤E(1− µαt)‖wt −w∗‖2 + 6E2Lα3tG2 + α2t 1
N ν2maxσ 2
where in the last inequality we use σ2 = ∑N k=1 pkσ 2 k, and that by construction αt satisfies Lαt ≤ 18 .
One may ask whether the dependence on E in the term κE 2G2/µ KT can be removed, or equivalently whether ∑ k pk‖wkt −wt‖2 = O(1/T 2) can be independent of E. We provide a simple counterexample that shows that this is not possible in general. Proposition 1. There exists a dataset such that if E = O(T β) for any β > 0 then ∑ k pk‖wkt − wt‖2 = Ω( 1T 2−2β ) .
Proof. Suppose that we have an even number of devices and each Fk(w) = 1nk ∑nk j=1(x j k −w)2 contains data points xjk = w ∗,k, with nk ≡ n. Moreover, the w∗,k’s come in pairs around the origin. As a result, the global objective F is minimized at w∗ = 0. Moreover, if we start from w0 = 0, then by design of the dataset the updates in local steps exactly cancel each other at each iteration, resulting in wt = 0 for all t. On the other hand, if E = T β , then starting from any t = O(T ) with constant step size O( 1T ), after E iterations of local steps, the local parameters are updated towards w
∗,k with ‖wkt+E‖2 = Ω((T β · 1T ) 2) = Ω( 1 T 2−2β
). This implies that∑ k pk‖wkt+E −wt+E‖2 = ∑ k pk‖wkt+E‖2
= Ω( 1
T 2−2β )
which is at a slower rate than 1T 2 for any β > 0. Thus the sampling variance E‖wt+1 − vt+1‖ 2 = Ω( ∑ k pkE‖wkt+1 −wt+1‖2) decays at a slower rate than 1 T 2 , resulting in a convergence rate slower than O( 1T ) with partial participation.

heading: E.2 CONVEX SMOOTH OBJECTIVES
text: In this section we provide the proof of the convergence result for FedAvg with convex and smooth objectives. The key step is a one step progress result analogous to that in the strongly convex case, and their proofs share identical components as well.
Lemma 6 (One step progress, convex case). Let wt = ∑N k=1 pkw k t in FedAvg. Under assumptions 1,3,4, the following bound holds for all t:
E‖wt+1 −w∗‖2 + αt(F (wt)− F (w∗)) ≤ E‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6α3tE 2LG2
Proof. The first part of the proof follows directly from Eq (10) in the proof of Lemma 3. Setting µ = 0 in Eq (10) (since we are in the convex setting instead of strongly convex), we obtain
‖wt+1 −w∗‖2 ≤ ‖wt −w∗‖2 + αtL N∑ k=1 pk‖wt −wkt ‖2
+ 2αt N∑ k=1 pk [Fk(w ∗)− Fk(wt)] + α2t ‖gt‖2 + α2t N∑ k=1 p2kσ 2 k
The difference of this bound with that in the strongly convex case is that we no longer have a contraction factor of 1− µαt in front of ‖wt −w∗‖2. In the strongly convex case, we were able to cancel α2t ‖gt‖2 with 2αt ∑N k=1 pk [Fk(w
∗)− Fk(wt)] and obtain only lower order terms. In the convex case, we use a different strategy and preserve ∑N k=1 pk [Fk(w
∗)− Fk(wt)] in order to obtain the desired optimality gap.
We have ‖gt‖2 = ‖ ∑ k pk∇Fk(wkt )‖2
= ‖ ∑ k pk∇Fk(wkt )− ∑ k pk∇Fk(wt) + ∑ k pk∇Fk(wt)‖2
≤ 2‖ ∑ k pk∇Fk(wkt )− ∑ k pk∇Fk(wt)‖2 + 2‖ ∑ k pk∇Fk(wt)‖2
≤ 2L2 ∑ k pk‖wkt −wt‖2 + 2‖ ∑ k pk∇Fk(wt)‖2
= 2L2 ∑ k pk‖wkt −wt‖2 + 2‖∇F (wt)‖2
using ∇F (w∗) = 0. Now using the L smoothness of F , we have ‖∇F (wt)‖2 ≤ 2L(F (wt) − F (w∗)), so that
‖wt+1 −w∗‖2
≤‖wt −w∗‖2 + αtL N∑ k=1 pk‖wt −wkt ‖2 + 2αt N∑ k=1 pk [Fk(w ∗)− Fk(wt)]
+ 2α2tL 2 ∑ k pk‖wkt −wt‖2 + 4α2tL(F (wt)− F (w∗)) + α2t N∑ k=1 p2kσ 2 k
=‖wt −w∗‖2 + (2α2tL2 + αtL) N∑ k=1 pk‖wt −wkt ‖2 + αt N∑ k=1 pk [Fk(w ∗)− Fk(wt)]
+ α2t N∑ k=1 p2kσ 2 k + αt(1− 4αtL)(F (w∗)− F (wt))
Since F (w∗) ≤ F (wt), as long as 4αtL ≤ 1, we can ignore the last term, and rearrange the inequality to obtain
‖wt+1 −w∗‖2 + αt(F (wt)− F (w∗))
≤‖wt −w∗‖2 + (2α2tL2 + αtL) N∑ k=1 pk‖wt −wkt ‖2 + α2t N∑ k=1 p2kσ 2 k
≤‖wt −w∗‖2 + 3
2 αtL N∑ k=1 pk‖wt −wkt ‖2 + α2t N∑ k=1 p2kσ 2 k
The same argument as before yields E ∑N k=1 pk‖wt −wkt ‖2 ≤ 4E2α2tG2 which gives
‖wt+1 −w∗‖2 + αt(F (wt)− F (w∗)) ≤ ‖wt −w∗‖2 + α2t N∑ k=1 p2kσ 2 k + 6α 3 tE 2LG2
≤ ‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6α3tE 2LG2
With the one step progress result, we can now prove the convergence result in the convex setting, which we restate below.
Theorem 2. Under assumptions 1,3,4 and constant learning rate αt = O( √ N T ), FedAvg satisfies
min t≤T
F (wt)− F (w∗) = O ( νmaxσ 2
√ NT
+ NE2LG2
T ) with full participation, and with partial device participation with K sampled devices at each commu-
nication round and learning rate αt = O( √ K T ),
min t≤T
F (wt)− F (w∗) = O ( νmaxσ 2
√ KT + E2G2√ KT + KE2LG2 T ) Proof. We first prove the bound for full participation. Applying Lemma 6, we have
‖wt+1 −w∗‖2 + αt(F (wt)− F (w∗)) ≤ ‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6α3tE 2LG2
Summing the inequalities from t = 0 to t = T , we obtain T∑ t=0 αt(F (wt)− F (w∗)) ≤ ‖w0 −w∗‖2 + T∑ t=0 α2t · 1 N ν2maxσ 2 + T∑ t=0 α3t · 6E2LG2
so that
min t≤T F (wt)− F (w∗) ≤ 1∑T t=0 αt
( ‖w0 −w∗‖2 +
T∑ t=0 α2t · 1 N ν2maxσ 2 + T∑ t=0 α3t · 6E2LG2 )
By setting the constant learning rate αt ≡ √ N T , we have
min t≤T F (wt)− F (w∗) ≤ 1√ NT · ‖w0 −w∗‖2 + 1√ NT T · N T · 1 N ν2maxσ 2 + 1√ NT T (
√ N
T )36E2LG2
≤ 1√ NT · ‖w0 −w∗‖2 + 1√ NT T · N T · 1 N ν2maxσ 2 + N T 6E2LG2
= (‖w0 −w∗‖2 + ν2maxσ2) 1√ NT + N T 6E2LG2
= O(ν 2 maxσ 2
√ NT
+ NE2LG2
T )
For partial participation, the one step progress bound in Lemma 6 is updated in a similar manner as the strongly convex case in (9) to incorporate the sampling variance. More precisely, with partial participation,
E‖wt+1 −w∗‖2 = E‖wt+1 − vt+1 + vt+1 −w∗‖2
= E‖wt+1 − vt+1‖2 + E‖vt+1 −w∗‖2, where Ewt+1 = vt+1 for all t, by the unbiasedness of our sampling schemes. Since vt =∑N
k=1 pkv k t always averages over all devices, the full participation one step progress bound in
Lemma 6 applied to vt implies
E‖vt+1 −w∗‖2 + αt(F (vt)− F (w∗)) ≤ E‖vt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6α3tE 2LG2
≤ E‖wt −w∗‖2 + α2t 1
N ν2maxσ 2 + 6α3tE 2LG2
The bound for E‖wt+1 − vt+1‖2 for the two sampling schemes we consider is provided in Eq (7), and applying it to the above bound we can write the one step progress for partial participation as
E‖wt+1 −w∗‖2 + αt(F (wt)− F (w∗)) ≤ E‖wt −w∗‖2 + α2t ( 1
N ν2maxσ 2 + C) + 6E2Lα3tG 2,
where C = 4KE 2G2 or N−KN−1 4 KE 2G2 depending on the sampling scheme.
Summing up the one-step progress over t,
min t≤T F (wt)− F (w∗) ≤ 1∑T t=0 αt
( ‖w0 −w∗‖2 +
T∑ t=0 α2t · ( 1 N νmaxσ 2 + C) + T∑ t=0 α3t · 6E2LG2 ) ,
so that with αt = √ K T , we have
min t≤T
F (wt)− F (w∗) = O( νmaxσ
2
√ KT + E2G2√ KT + KE2LG2 T ).

heading: F PROOF OF CONVERGENCE RESULTS FOR NESTEROV ACCELERATED FEDAVG

heading: F.1 STRONGLY CONVEX SMOOTH OBJECTIVES
text: Recall that the Nesterov accelerated FedAvg follows the updates
vkt+1 = w k t − αtgt,k, wkt+1 =
{ vkt+1 + βt(v k t+1 − vkt ) if t+ 1 /∈ IE ,∑
k∈St+1 qk [ vkt+1 + βt(v k t+1 − vkt ) ] if t+ 1 ∈ IE .
The proofs of convergence results for Nesterov Accelerated FedAvg consists of components that are direct analogues of the FedAvg case. We first state these analogue results before proving the main theorem. Like before, the proofs of the lemmas are deferred to after the main proof.
Lemma 7 (One step progress, Nesterov). Let vt = ∑N k=1 pkv k t in Nesterov accelerated FedAvg, and suppose our functions satisfy Assumptions 1,2,3,4, and set step sizes αt = 6µ 1 t+γ , βt−1 = 3 14(t+γ)(1− 6t+γ )max{µ,1}
with γ = max{32κ,E} and κ = Lµ , the updates of Nesterov accelerated

heading: FedAvg satisfy
text: E‖vt+1 −w∗‖2 ≤ E(1− µαt)(1 + βt−1)2‖vt −w∗‖2 + 20E2Lα3tG2 + (1− αtµ)β2t−1‖(vt−1 −w∗)‖2
+ α2t 1
N νmaxσ
2 + 2βt−1(1 + βt−1)(1− αtµ)‖vt −w∗‖ · ‖vt−1 −w∗‖
The one step progress result makes use of the same bound on the gradient variance in Lemma 4, as well as a divergence bound analogous to Lemma 5, which we state below. Lemma 8 (Bounding the divergence of wkt , Nesterov). Given Assumption 4, and assume that αt is non-increasing, αt ≤ 2αt+E , and 2β2t−1 + 2α2t ≤ 1/2 for all t ≥ 0, wt = ∑N k=1 pkw k t in Nesterov accelerated FedAvg satisfies
E [ N∑ k=1 pk‖wt −wkt ‖2 ] ≤ 16(E − 1)2α2tG2.
Theorem 3. Let vT = ∑N k=1 pkv k T in Nesterov accelerated FedAvg and set learning rates αt = 6 µ 1 t+γ , βt−1 =
3 14(t+γ)(1− 6t+γ )max{µ,1} . Then under Assumptions 1,2,3,4 with full device participation,
EF (vT )− F ∗ = O ( κνmaxσ 2/µ
NT + κ2E2G2/µ T 2
) ,
and with partial device participation with K sampled devices at each communication round, EF (vT )− F ∗ = O ( κνmaxσ 2/µ
NT + κE2G2/µ KT + κ2E2G2/µ T 2
) .
Proof. We first prove the result for full participation. Applying the one step progress bound in Lemma 7, we have
E‖vt+1 −w∗‖2 ≤ E(1− µαt)(1 + βt−1)2‖vt −w∗‖2 + 20E2Lα3tG2 + (1− αtµ)β2t−1‖(vt−1 −w∗)‖2
+ α2t 1
N νmaxσ
2 + 2βt−1(1 + βt−1)(1− αtµ)‖vt −w∗‖ · ‖vt−1 −w∗‖
Recall that we require αt0 ≤ 2αt for any t− t0 ≤ E− 1, Lαt ≤ 15 , and 2β 2 t−1 + 2α 2 t ≤ 1/2 in order for Lemmas 8 and 7 to hold, which we can check by definition of αt and βt. We show next that E‖vt −w∗‖2 = O(ν 2 maxσ 2/µ tN + E2LG2/µ2
t2 ) by induction. Assume that we have shown
E‖vt −w∗‖2 ≤ b(Cα2t +Dαt)
for all iterations until t, where C = 20E2LG2, D = 1N ν 2 maxσ 2, and b is some constant to be chosen later. For step sizes recall that we choose αt = 6µ 1 t+γ and βt−1 =
3 14(t+γ)(1− 6t+γ )max{µ,1} where
γ = max{32κ,E}, so that βt−1 ≤ αt and
(1− µαt)(1 + 14βt−1) ≤ (1− 6
t+ γ )(1 +
3
(t+ γ)(1− 6t+γ ) )
= 1− 6 t+ γ + 3 t+ γ = 1− 3 t+ γ = 1− µαt 2
Moreover, E‖vt−1 −w∗‖2 ≤ b(Cα2t−1 + Dαt−1) ≤ 4b(Cα2t + Dαt) with the chosen step sizes. Therefore the bound for E‖vt+1 −w∗‖2 can be further simplified with
2βt−1(1 + βt−1)(1− αtµ)E‖vt −w∗‖ · ‖vt−1 −w∗‖ ≤ 4βt−1(1 + βt−1)(1− αtµ) · b(Cα2t +Dαt)
and
(1− αtµ)β2t−1E‖(vt−1 −w∗)‖2 ≤ 4(1− αtµ)β2t−1 · b(Cα2t +Dαt)
so that
E‖vt+1 −w∗‖2 ≤ (1− µαt)((1 + βt−1)2 + 4βt−1(1 + βt−1) + 4β2t−1) · b(Cα2t +Dαt)
+ 20E2Lα3tG 2 + α2t
1
N νmaxσ
2
≤ E(1− µαt)(1 + 14βt−1) · b(Cα2t +Dαt) + 20E2Lα3tG2 + α2t 1
N νmaxσ
2
≤ b(1− µαt 2 )(Cα2t +Dαt) + Cα 3 t +Dα 2 t = (b(1− µαt 2 ) + αt)α 2 tC + (b(1− µαt 2 ) + αt)αtD
and so it remains to choose b such that
(b(1− µαt 2 ) + αt)αt ≤ bαt+1 (b(1− µαt 2 ) + αt)α 2 t ≤ bα2t+1
from which we can conclude E‖vt+1 −w∗‖2 ≤ α2t+1C + αt+1D.
With b = 6µ , we have
(b(1− µαt 2 ) + αt)αt = (b(1− ( 3 t+ γ ) +
6
µ(t+ γ) )
6
µ(t+ γ)
= (b t+ γ − 3 t+ γ + 6 µ(t+ γ) )
6
µ(t+ γ)
≤ b( t+ γ − 1 t+ γ ) 6 µ(t+ γ)
≤ b 6 µ(t+ γ + 1) = bαt+1
where we have used t+γ−1(t+γ)2 ≤ 1 t+γ+1 .
Similarly
(b(1− µαt 2 ) + αt)α 2 t = (b(1− ( 3 t+ γ ) +
6
µ(t+ γ) )(
6
µ(t+ γ) )2
= (b t+ γ − 3 t+ γ + 6 µ(t+ γ) )(
6
µ(t+ γ) )2
= b( t+ γ − 2 t+ γ )( 6 µ(t+ γ) )2
≤ b 36 µ2(t+ γ + 1)2 = bα2t+1
where we have used t+γ−2(t+γ)3 ≤ 1 (t+γ+1)2 .
Finally, to ensure ‖v0 −w∗‖2 ≤ b(Cα20 +Dα0), we can rescale b by c‖v0 −w∗‖2 for some c. It follows that E‖vt −w∗‖2 ≤ b(Cα2t +Dαt) for all t ≥ 0. Using the L-smooothness of F ,
E(F (vT ))− F ∗ = E(F (vT )− F (w∗))
≤ L 2 E‖vT −w∗‖2 ≤ L 2 c‖v0 −w∗‖2 6 µ (DαT + Cα 2 T )
= 3c‖v0 −w∗‖2κ(DαT + Cα2T ) ≤ 3c‖v0 −w∗‖2κ [ 6
µ(T + γ) · 1 N νmaxσ 2 + 20E2LG2 · ( 6 µ(T + γ)
)2 ]
= O(κ µ 1 N νmaxσ 2 · 1 T + κ2 µ E2G2 · 1 T 2 )
With partial participation, the same argument with an added term for sampling error yields
EF (wT )− F ∗ = O( κνmaxσ
2/µ NT + κE2G2/µ KT + κ2E2G2/µ T 2 )

heading: F.1.1 DEFERRED PROOFS OF KEY LEMMAS
text: Proof of lemma 8. The proof of bound for E ∑N k=1 pk‖wt − wkt ‖2 in the Nesterov accelerated FedAvg follows a similar logic as in Lemma 5, but requires extra reasoning. Since communication is done every E steps, for any t ≥ 0, we can find a t0 ≤ t such that t− t0 ≤ E − 1 and wkt0 = wt0 for all k. Moreover, using αt is non-increasing, αt0 ≤ 2αt, and βt ≤ αt for any t− t0 ≤ E−1, we have E N∑ k=1 pk‖wt −wkt ‖2 = E N∑ k=1 pk‖wkt −wt0 − (wt −wt0)‖2
≤ E N∑ k=1 pk‖wkt −wt0‖2
= E N∑ k=1 pk‖wkt −wkt0‖ 2
= E N∑ k=1 pk‖ t−1∑ i=t0 βi(v k i+1 − vki )− t−1∑ i=t0 αigi,k‖2
≤ 2 N∑ k=1 pkE t−1∑ i=t0 (E − 1)α2i ‖gi,k‖2 + 2 N∑ k=1 pkE t−1∑ i=t0 (E − 1)β2i ‖(vki+1 − vki )‖2
≤ 2 N∑ k=1 pkE t−1∑ i=t0 (E − 1)α2i (‖gi,k‖2 + ‖(vki+1 − vki )‖2)
≤ 4 N∑ k=1 pkE t−1∑ i=t0 (E − 1)α2iG2
≤ 4(E − 1)2α2t0G 2 ≤ 16(E − 1)2α2tG2
where we have used E‖vkt − vkt−1‖2 ≤ G2. To see this identity for appropriate αt, βt, note the recursion
vkt+1 − vkt = wkt −wkt−1 − (αtgt,k − αt−1gt−1,k) wkt+1 −wkt = −αtgt,k + βt(vkt+1 − vkt )
so that
vkt+1 − vkt = −αt−1gt−1,k + βt−1(vkt − vkt−1)− (αtgt,k − αt−1gt−1,k) = βt−1(v k t − vkt−1)− αtgt,k
Since the identity vkt+1 − vkt = βt−1(vkt − vkt−1)− αtgt,k implies
E‖vkt+1 − vkt ‖2 ≤ 2β2t−1E‖vkt − vkt−1‖2 + 2α2tG2
as long as αt, βt−1 satisfy 2β2t−1 + 2α 2 t ≤ 1/2, we can guarantee that E‖vkt − vkt−1‖2 ≤ G2 for all k by induction. This together with Jensen’s inequality also gives E‖vt − vt−1‖2 ≤ G2 for all t.
Now we are ready to prove the one step progress result for Nesterov accelerated FedAvg. The first part of the proof is identical to that of the FedAvg case, while the main recursion takes a different form.
Proof of lemma 7. We again have
‖vt+1 −w∗‖2 = ‖(wt − αtgt)−w∗‖2
and using exactly the same derivation as the FedAvg case, we can obtain the following bound (same as Eq (11) in the proof of Lemma 3):
E‖wt+1 −w∗‖2 ≤ E(1− µαt)‖wt −w∗‖2 + αtL N∑ k=1 pk‖wt −wkt ‖2 + α2t N∑ k=1 p2kσ 2 k
+ α2tL 2 N∑ k=1 pk‖wt −wkt ‖2 + α3tLE‖gt‖2 − α2t ‖∇F (wt)‖2
Different from the FedAvg case, we no longer have wt = vt. Instead,
‖wt −w∗‖2 = ‖vt + βt−1(vt − vt−1)−w∗‖2
= ‖(1 + βt−1)(vt −w∗)− βt−1(vt−1 −w∗)‖2
= (1 + βt−1) 2‖vt −w∗‖2 − 2βt−1(1 + βt−1)〈vt −w∗,vt−1 −w∗〉+ β2t−1‖(vt−1 −w∗)‖2 ≤ (1 + βt−1)2‖vt −w∗‖2 + 2βt−1(1 + βt−1)‖vt −w∗‖ · ‖vt−1 −w∗‖+ β2t−1‖(vt−1 −w∗)‖2
which gives a recursion involving both vt and vt−1: ‖vt+1 −w∗‖2 ≤ (1− αtµ)(1 + βt−1)2‖vt −w∗‖2 + 2(1− αtµ)βt−1(1 + βt−1)‖vt −w∗‖ · ‖vt−1 −w∗‖+ α2t N∑ k=1 p2kσ 2 k
+ β2t−1(1− αtµ)‖(vt−1 −w∗)‖2 + αtL N∑ k=1 pk‖wt −wkt ‖2 + α2tL2 ∑ k pk‖wt −wkt ‖2 + α3tLG2
and we will using this recursive relation to obtain the desired bound.
We can check that our choice of αt and βt satisfy αt is non-increasing, αt ≤ 2αt+E , and 2β2t−1 + 2α2t ≤ 1/2 for all t ≥ 0, so that we can apply the bound from Lemma 8 on E ∑N k=1 pk‖wt −wkt ‖2 to conclude that, with νmax := N ·maxk pk,
E‖vt+1 −w∗‖2 ≤ E(1− µαt)(1 + βt−1)2‖vt −w∗‖2 + 16E2Lα3tG2 + 16E2L2α4tG2 + α3tLG2
+ (1− αtµ)β2t−1‖(vt−1 −w∗)‖2 + α2t N∑ k=1 p2kσ 2 k + 2βt−1(1 + βt−1)(1− αtµ)‖vt −w∗‖ · ‖vt−1 −w∗‖ ≤ E(1− µαt)(1 + βt−1)2‖vt −w∗‖2 + 20E2Lα3tG2 + (1− αtµ)β2t−1‖(vt−1 −w∗)‖2
+ α2t 1
N νmaxσ
2 + 2βt−1(1 + βt−1)(1− αtµ)‖vt −w∗‖ · ‖vt−1 −w∗‖
where we have used σ2 = ∑ k pkσ 2 k, and by construction our αt satisfies Lαt ≤ 15 .

heading: F.2 CONVEX SMOOTH OBJECTIVES
text: In this section we provide proof of the convergence result for Nesterov accelerated FedAvg with convex and smooth objectives. Unlike with the FedAvg algorithm, where convex and strongly convex results share identical components, the proof for the convergence result in the convex setting for Nesterov FedAvg uses a change of variables, although the general ideas are in the same vein: we have a one step progress bound for E‖wt+1 −w∗‖2 + ηt(F (wt)− F (w∗)), which is then used to form a telescoping sum that gives an upper bound on mint≤T F (wt)− F (w∗).
Lemma 9 (One step progress, convex case, Nesterov). Let wt = ∑N k=1 pkw k t in Nesterov accelerated FedAvg, and define ηt = αt1−βt . Under assumptions 1,3,4, the following bound holds for all t:
E‖wt+1 −w∗‖2 + ηt(F (wt)− F (w∗)) ≤ E‖wt −w∗‖2 + 32LE2α2t ηtG2 + η2t νmax 1
N σ2 + 2ηt β2t 1− βt G2
Theorem 4. Set learning rates αt = βt = O( √ N T ). Then under Assumptions 1,3,4 Nesterov accelerated FedAvg with full device participation has rate
min t≤T
F (wt)− F ∗ = O ( νmaxσ 2
√ NT
+ NE2LG2
T
) ,
and with partial device participation with K sampled devices at each communication round,
min t≤T
F (wt)− F ∗ = O ( νmaxσ 2
√ KT + E2G2√ KT + KE2LG2 T
) .
Proof. Applying the bound from Lemma 9, with ηt = αt1−βt we have
E‖wt+1 −w∗‖2 + ηt(F (wt)− F (w∗)) ≤ E‖wt −w∗‖2 + 32LE2α2t ηtG2 + η2t νmax 1
N σ2 + 2ηt β2t 1− βt G2
Summing the inequalities from t = 0 to t = T , we obtain T∑ t=0 ηt(F (wt)− F (w∗)) ≤ ‖w0 −w∗‖2 + T∑ t=0 η2t · 1 N νmaxσ 2 + T∑ t=0 ηtα 2 t · 32LE2G2 + T∑ t=0 2ηt β2t 1− βt G2
so that
min t≤T F (wt)− F (w∗) ≤ 1∑T t=0 ηt
( ‖w0 −w∗‖2 +
T∑ t=0 η2t · 1 N νmaxσ 2 + T∑ t=0 ηtα 2 t · 32LE2G2 + T∑ t=0 2ηt β2t 1− βt G2
)
By setting the constant learning rates αt ≡ √ N T and βt ≡ c √ N T so that ηt = αt 1−βt = √ N T 1−c √
N T
≤ 2 √
N T , we have
min t≤T
F (wt)− F (w∗)
≤ 1 2 √ NT · ‖w0 −w∗‖2 + 2√ NT T · N T · 1 N νmaxσ 2 + 1√ NT T (
√ N
T )332LE2G2 + 2√ NT T (
√ N
T )3G2
= ( 1
2 ‖w0 −w∗‖2 + 2νmaxσ2) 1√ NT + N T (32LE2G2 + 2G2)
= O( νmaxσ
2
√ NT
+ NE2LG2
T )
Similarly, for partial participation, we have
min t≤T F (wt)− F (w∗) ≤ 1∑T t=0 αt
( ‖w0 −w∗‖2 +
T∑ t=0 α2t · ( 1 N νmaxσ 2 + C) + T∑ t=0 α3t · 6E2LG2 )
where C = 4KE 2G2 or N−KN−1 4 KE 2G2, so that with αt ≡ √ K T and βt ≡ c √ K T , we have
min t≤T
F (wt)− F (w∗) = O( νmaxσ
2
√ KT + E2G2√ KT + KE2LG2 T )

heading: F.2.1 DEFERRED PROOFS OF KEY LEMMAS
text: Proof of lemma 9. Define pt := βt 1−βt [wt −wt−1 + αtgt−1] = β2t
1−βt (vt − vt−1) for t ≥ 1 and 0 for t = 0. We can check that
wt+1 + pt+1 = wt + pt − αt
1− βt gt
Now we define zt := wt + pt and ηt = αt 1−βt for all t, so that we have the recursive relation
zt+1 = zt − ηtgt Now
‖zt+1 −w∗‖2 = ‖(zt − ηtgt)−w∗‖2
= ‖(zt − ηtgt −w∗)− ηt(gt − gt)‖2
= A1 +A2 +A3
where
A1 = ‖zt −w∗ − ηtgt‖2 A2 = 2ηt〈zt −w∗ − ηtgt,gt − gt〉 A3 = η 2 t ‖gt − gt‖2
where again EA2 = 0 and EA3 ≤ η2t ∑ k p 2 kσ 2 k. For A1 we have
‖zt −w∗ − ηtgt‖2 = ‖zt −w∗‖2 + 2〈zt −w∗,−ηtgt〉+ ‖ηtgt‖2
Using the convexity and L-smoothness of Fk,
− 2ηt〈zt −w∗,gt〉
= −2ηt N∑ k=1 pk〈zt −w∗,∇Fk(wkt )〉
= −2ηt N∑ k=1 pk〈zt −wkt ,∇Fk(wkt )〉 − 2ηt N∑ k=1 pk〈wkt −w∗,∇Fk(wkt )〉
= −2ηt N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉 − 2ηt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉 − 2ηt N∑ k=1 pk〈wkt −w∗,∇Fk(wkt )〉
≤ −2ηt N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉 − 2ηt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉+ 2ηt N∑ k=1 pk(Fk(w ∗)− Fk(wkt ))
≤ 2ηt N∑ k=1 pk [ Fk(w k t )− Fk(wt) + L 2 ‖wt −wkt ‖2 + Fk(w∗)− Fk(wkt ) ]
− 2ηt N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉
= ηtL N∑ k=1 pk‖wt −wkt ‖2 + 2ηt N∑ k=1 pk [Fk(w ∗)− Fk(wt)]− 2ηt N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉
which results in
E‖wt+1 −w∗‖2 ≤ E‖wt −w∗‖2 + ηtL N∑ k=1 pk‖wt −wkt ‖2 + 2ηt N∑ k=1 pk [Fk(w ∗)− Fk(wt)]
+ η2t ‖gt‖2 + η2t N∑ k=1 p2kσ 2 k − 2ηt N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉
As before, ‖gt‖2 ≤ 2L2 ∑ k pk‖wkt −wt‖2 + 4L(F (wt)− F (w∗)), so that η2t ‖gt‖2 + ηt N∑ k=1 pk [Fk(w ∗)− Fk(wt)] ≤ 2L2η2t ∑ k pk‖wkt −wt‖2 + ηt(1− 4ηtL)(F (w∗)− F (wt))
≤ 2L2η2t ∑ k pk‖wkt −wt‖2
for ηt ≤ 1/4L. Using ∑N k=1 pk‖wt−wkt ‖2 ≤ 16E2α2tG2 and ∑N k=1 p 2 kσ 2 k ≤ νmax 1N σ
2, it follows that E‖wt+1 −w∗‖2 + ηt(F (wt)− F (w∗)) ≤ E‖wt −w∗‖2 + (ηtL+ 2L2η2t ) N∑ k=1 pk‖wt −wkt ‖2 + η2t N∑ k=1 p2kσ 2 k
− 2ηt N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉
≤ E‖wt −w∗‖2 + 32LE2α2t ηtG2 + η2t νmax 1
N σ2
− 2ηt N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉
if ηt ≤ 12L . It remains to bound E ∑N k=1 pk〈zt − wt,∇Fk(wkt )〉. Recall that zt − wt =
βt 1−βt [wt −wt−1 + αtgt−1] = β2t 1−βt (vt−vt−1) and E‖vt−vt−1‖ 2 ≤ G2, E‖∇Fk(wkt )‖2 ≤ G2.
Cauchy-Schwarz gives
E N∑ k=1 pk〈zt −wt,∇Fk(wkt )〉 ≤ N∑ k=1 pk √ E‖zt −wt‖2 · √ E‖∇Fk(wkt )‖2
≤ β 2 t
1− βt G2
Thus
E‖wt+1 −w∗‖2 + ηt(F (wt)− F (w∗)) ≤ E‖wt −w∗‖2 + 32LE2α2t ηtG2 + η2t νmax 1
N σ2 + 2ηt β2t 1− βt G2

heading: G GEOMETRIC CONVERGENCE OF FEDAVG IN THE OVERPARAMETERIZED SETTING
text: Overparameterization is a prevalent machine learning setting where the statistical model has much more parameters than the number of training samples and the existence of parameter choices with zero training loss is ensured Allen-Zhu et al. (2018); Zhang et al. (2016). Due to the property of automatic variance reduction in overparameterization, a line of recent works proved that SGD and accelerated methods achieve geometric convergence Ma et al. (2018); Moulines & Bach (2011); Needell et al. (2014); Schmidt & Roux (2013); Strohmer & Vershynin (2009). A natural question is whether such a result still holds in the federated learning setting. In this section, we provide the first geometric convergence rate of FedAvg for the overparameterized strongly convex and smooth problems, and show that it preserves linear speedup at the same time. We then sharpen this result in the special case of linear regression. Inspired by recent advances in accelerating SGD Liu et al. (2020); Jain et al. (2017), we further propose a novel momentum-based FedAvg algorithm, which enjoys an improved convergence rate over FedAvg. Detailed proofs are deferred to Appendix Section H. In particular, we do not need Assumptions 3 and 4 and use modified versions of Assumptions 1 and 2 detailed in this section.

heading: G.1 GEOMETRIC CONVERGENCE OF FEDAVG IN THE OVERPARAMETERIZED SETTING
text: Recall the FL problem minw ∑N k=1 pkFk(w) with Fk(w) = 1 nk ∑nk j=1 `(w;x j k). In this section, we consider the standard Empirical Risk Minimization (ERM) setting where ` is non-negative, l-smooth, and convex, and as before, each Fk(w) is L-smooth and µ-strongly convex. Note that l ≥ L. This setup includes many important problems in practice. In the overparameterized setting, there exists w∗ ∈ arg minw ∑N k=1 pkFk(w) such that `(w ∗;xjk) = 0 for all x j k. We first show that FedAvg achieves geometric convergence with linear speedup in the number of workers. Theorem 5. In the overparameterized setting, FedAvg with communication every E iterations and constant step size α = O( 1E N lνmax+L(N−νmin) ) has geometric convergence:
EF (wT ) ≤ L
2 (1− α)T ‖w0 −w∗‖2 = O
( L exp ( − µ E
NT
lνmax + L(N − νmin)
) · ‖w0 −w∗‖2 ) .
Linear speedup and Communication Complexity The linear speedup factor is on the order of O(N/E) for N ≤ O( lL ), i.e. FedAvg with N workers and communication every E iterations provides a geometric convergence speedup factor of O(N/E), for N ≤ O( lL ). When N is above this threshold, however, the speedup is almost constant in the number of workers. This matches the findings in Ma et al. (2018). Our result also illustrates that E can be taken O(T β) for any β < 1 to achieve geometric convergence, achieving better communication efficiency than the standard FL setting. We emphasize again that compared to the single-server results in Ma et al. (2018), the difference of our result lies in the factor of N in the speedup, which cannot be obtained if one simply applied the single-server result to each device in our problem.

heading: G.2 OVERPARAMETERIZED LINEAR REGRESSION PROBLEMS
text: We now turn to quadratic problems and show that the bound in Theorem 5 can be improved to O(exp(− NEκ1 t)) for a larger range of N . We then propose a variant of FedAvg that has provable acceleration over FedAvg with SGD updates. The local device objectives are now given by the sum of squares Fk(w) = 12nk ∑nk j=1(w Txjk − z j k)
2, and there exists w∗ such that F (w∗) ≡ 0. Two notions of condition number are important in our results: κ1 which is based on local Hessians, and κ̃, which is termed the statistical condition number Liu & Belkin (2020); Jain et al. (2017). For their detailed definitions, please refer to Appendix Section H. Here we use the fact κ̃ ≤ κ1. Recall νmax = maxk pkN and νmin = mink pkN . Theorem 6. For the overparamterized linear regression problem, FedAvg with communication every E iterations with constant step size α = O( 1E N lνmax+µ(N−νmin) ) has geometric convergence:
EF (wT ) ≤ O ( L exp(− NT
E(νmaxκ1 + (N − νmin)) )‖w0 −w∗‖2
) .
When N = O(κ1), the convergence rate is O((1 − NEκ1 ) T ) = O(exp(− NTEκ1 )), which exhibits linear speedup in the number of workers, as well as a 1/κ1 dependence on the condition number κ1. Inspired by Liu & Belkin (2020), we propose the MaSS accelerated FedAvg algorithm (FedMaSS):
wkt+1 = { ukt − ηk1gt,k if t+ 1 /∈ IE ,∑ k∈St+1 [ ukt − ηk1gt,k ] if t+ 1 ∈ IE ,
ukt+1 = w k t+1 + γ k(wkt+1 −wkt ) + ηk2gt,k. When ηk2 ≡ 0, this algorithm reduces to the Nesterov accelerated FedAvg algorithm. In the next theorem, we demonstrate that FedMaSS improves the convergence to O(exp(− NT
E √ κ1κ̃ )). To our knowledge, this is the first acceleration result of FedAvg with momentum updates over SGD updates. Theorem 7. For the overparamterized linear regression problem, FedMaSS with communication
everyE iterations and constant step sizes η1 = O( 1E N lνmax+µ(N−νmin) ), η2 = η1(1− 1κ̃ ) 1+ 1√
κ1κ̃
, γ = 1− 1√ κ1κ̃
1+ 1√ κ1κ̃
has geometric convergence: EF (wT ) ≤ O ( L exp(− NT
E(νmax √ κ1κ̃+ (N − νmin))
)‖w0 −w∗‖2 ) .
Speedup of FedMaSS over FedAvg To better understand the significance of the above result, we briefly discuss related works on accelerating SGD. Nesterov and Heavy Ball updates are known to fail to accelerate over SGD in both the overparameterized and convex settings Liu & Belkin (2020); Kidambi et al. (2018); Liu et al. (2018); Yuan et al. (2016). Thus in general one cannot hope to obtain acceleration results for the FedAvg algorithm with Nesterov and Heavy Ball updates. Luckily, recent works in SGD Jain et al. (2017); Liu & Belkin (2020) introduced an additional compensation term to the Nesterov updates to address the non-acceleration issue. Surprisingly, we show the same approach can effectively improve the rate of FedAvg. Comparing the convergence rate of FedMass (Theorem 7) and FedAvg (Theorem 6), when N = O( √ κ1κ̃), the convergence rate is O((1− N E √ κ1κ̃ )T ) = O(exp(− NT E √ κ1κ̃
)) as opposed to O(exp(− NTEκ1 )). Since κ1 ≥ κ̃, this implies a speedup factor of √ κ1 κ̃ for FedMaSS. On the other hand, the same linear speedup in the number of workers holds for N in a smaller range of values.

heading: H PROOF OF GEOMETRIC CONVERGENCE RESULTS FOR OVERPARAMETERIZED PROBLEMS

heading: H.1 GEOMETRIC CONVERGENCE OF FEDAVG FOR GENERAL STRONGLY CONVEX AND SMOOTH OBJECTIVES
text: Theorem 5. For the overparameterized setting with general strongly convex and smooth objectives, FedAvg with local SGD updates and communication every E iterations with constant step size α = 12E N lνmax+L(N−νmin) gives the exponential convergence guarantee
EF (wt) ≤ L
2 (1− µα)t‖w0 −w∗‖2 = O(exp(−
µ
2E
N
lνmax + L(N − νmin) t) · ‖w0 −w∗‖2)
Proof. To illustrate the main ideas of the proof, we first present the proof for E = 2. Let t− 1 be a communication round, so that wkt−1 = wt−1. We show that
‖wt+1 −w∗‖2 ≤ (1− αtµ)(1− αt−1µ)‖wt−1 −w∗‖2
for appropriately chosen constant step sizes αt, αt−1. We have
‖wt+1 −w∗‖2 = ‖(wt − αtgt)−w∗‖2
= ‖wt −w∗‖2 − 2αt〈wt −w∗,gt〉+ α2t ‖gt‖2
and the cross term can be bounded as usual using µ-convexity and L-smoothness of Fk:
− 2αtEt〈wt −w∗,gt〉
= −2αt N∑ k=1 pk〈wt −w∗,∇Fk(wkt )〉
= −2αt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉 − 2αt N∑ k=1 pk〈wkt −w∗,∇Fk(wkt )〉
≤ −2αt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉+ 2αt N∑ k=1 pk(Fk(w ∗)− Fk(wkt ))− αtµ N∑ k=1 pk‖wkt −w∗‖2
≤ 2αt N∑ k=1 pk [ Fk(w k t )− Fk(wt) + L 2 ‖wt −wkt ‖2 + Fk(w∗)− Fk(wkt ) ] − αtµ‖ N∑ k=1 pk(w k t −w∗)‖2
= αtL N∑ k=1 pk‖wt −wkt ‖2 + 2αt N∑ k=1 pk [Fk(w ∗)− Fk(wt)]− αtµ‖wt −w∗‖2
= αtL N∑ k=1 pk‖wt −wkt ‖2 − 2αt N∑ k=1 pkFk(wt)− αtµ‖wt −w∗‖2
and so
E‖wt+1 −w∗‖2 ≤ E(1− αtµ)‖wt −w∗‖2 − 2αtF (wt) + α2t ‖gt‖2 + αtL N∑ k=1 pk‖wt −wkt ‖2
Applying this recursive relation to ‖wt −w∗‖2 and using ‖wt−1 −wkt−1‖2 ≡ 0, we further obtain E‖wt+1 −w∗‖2 ≤ E(1− αtµ) ( (1− αt−1µ)‖wt−1 −w∗‖2 − 2αt−1F (wt−1) + α2t−1‖gt−1‖2 ) − 2αtF (wt) + α2t ‖gt‖2 + αtL
N∑ k=1 pk‖wt −wkt ‖2
Now instead of bounding ∑N k=1 pk‖wt −wkt ‖2 using the arguments in the general convex case, we follow Ma et al. (2018) and use the fact that in the overparameterized setting, w∗ is a minimizer of each `(w, xjk) and that each ` is l-smooth to obtain ‖∇Fk(wt−1, ξkt−1)‖2 ≤ 2l(Fk(wt−1, ξkt−1) − Fk(w ∗, ξkt−1)), where recall Fk(w, ξ k t−1) = `(w, ξ k t−1), so that
N∑ k=1 pk‖wt −wkt ‖2 = N∑ k=1 pk‖wt−1 − αt−1gt−1 −wkt−1 + αt−1gt−1,k‖2
= N∑ k=1 pkα 2 t−1‖gt−1 − gt−1,k‖2
= α2t−1 N∑ k=1 pk(‖gt−1,k‖2 − ‖gt−1‖2)
= α2t−1 N∑ k=1 pk‖∇Fk(wt−1, ξkt−1)‖2 − α2t−1‖gt−1‖2
≤ α2t−1 N∑ k=1 pk2l(Fk(wt−1, ξ k t−1)− Fk(w∗, ξkt−1))− α2t−1‖gt−1‖2
again using wt−1 = wkt−1. Taking expectation with respect to ξ k t−1’s and using the fact that F (w∗) = 0, we have
Et−1 N∑ k=1 pk‖wt −wkt ‖2 ≤ 2lα2t−1 N∑ k=1 pkFk(wt−1)− α2t−1‖gt−1‖2
= 2lα2t−1F (wt−1)− α2t−1‖gt−1‖2
Note also that
‖gt−1‖2 = ‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2
while ‖gt‖2 = ‖ N∑ k=1 pk∇Fk(wkt , ξkt )‖2 ≤ 2‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2 + 2‖ N∑ k=1 pk(∇Fk(wt, ξkt )−∇Fk(wkt , ξkt ))‖2
≤ 2‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2 + 2 N∑ k=1 pkl 2‖wt −wkt ‖2
Substituting these into the bound for ‖wt+1 −w∗‖2, we have
E‖wt+1 −w∗‖2 ≤ E(1− αtµ)((1− αt−1µ)‖wt−1 −w∗‖2 − 2αt−1F (wt−1) + α2t−1‖gt−1‖2)
− 2αtF (wt) + 2α2t ‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2 + ( 2l2α2t−1α 2 t + αtα 2 t−1L ) ( 2lF (wt−1)− ‖gt−1‖2 ) = E(1− αtµ)(1− αt−1µ)‖wt−1 −w∗‖2
− 2αt(F (wt)− αt‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2)
− 2αt−1(1− αtµ) ( (1− lαt−1(2l 2α2t + αtL)
1− αtµ )F (wt−1)− αt−1 2 ‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2 )
from which we can conclude that
E‖wt+1 −w∗‖2 ≤ (1− αtµ)(1− αt−1µ)E‖wt−1 −w∗‖2
if we can choose αt, αt−1 to guarantee
E(F (wt)− αt‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2) ≥ 0
E ( (1− lαt−1(2l 2α2t + αtL)
1− αtµ )F (wt−1)− αt−1 2 ‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2 ) ≥ 0
Note that
Et‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2 = Et〈 N∑ k=1 pk∇Fk(wt, ξkt ), N∑ k=1 pk∇Fk(wt, ξkt )〉
= N∑ k=1 p2kEt‖∇Fk(wt, ξkt )‖2 + N∑ k=1 ∑ j 6=k pjpkEt〈∇Fk(wt, ξkt ),∇Fj(wt, ξ j t )〉
= N∑ k=1 p2kEt‖∇Fk(wt, ξkt )‖2 + N∑ k=1 ∑ j 6=k pjpk〈∇Fk(wt),∇Fj(wt)〉
= N∑ k=1 p2kEt‖∇Fk(wt, ξkt )‖2 + N∑ k=1 N∑ j=1 pjpk〈∇Fk(wt),∇Fj(wt)〉 − N∑ k=1 p2k‖∇Fk(wt)‖2 ≤ N∑ k=1 p2kEt‖∇Fk(wt, ξkt )‖2 + ‖ ∑ k pk∇Fk(wt)‖2 − 1 N νmin‖ ∑ k pk∇Fk(wt)‖2
= N∑ k=1 p2kEt‖∇Fk(wt, ξkt )‖2 + (1− 1 N νmin)‖∇F (wt)‖2
and so following Ma et al. (2018) if we let αt = min{ qN2lνmax , 1−q 2L(1− 1N νmin) } for a q ∈ [0, 1] to be optimized later, we have
Et(F (wt)− αt‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2)
≥ Et N∑ k=1 pkFk(wt)− αt [ N∑ k=1 p2kEt‖∇Fk(wt, ξkt )‖2 + (1− 1 N νmin)‖∇F (wt)‖2 ]
≥ Et N∑ k=1 pk(qFk(wt, ξ k t )− αt 1 N νmax‖∇Fk(wt, ξkt )‖2) + ((1− q)F (wt)− αt(1− 1 N νmin)‖∇F (wt)‖2)
≥ qEt N∑ k=1 pk(Fk(wt, ξ k t )− 1 2l ‖∇Fk(wt, ξkt )‖2) + (1− q)(F (wt)− 1 2L ‖∇F (wt)‖2)
≥ 0
again using w∗ optimizes Fk(w, ξkt ) with Fk(w ∗, ξkt ) = 0.
Maximizing αt = min{ qN2lνmax , 1−q 2L(1− 1N νmin) } over q ∈ [0, 1], we see that q = lνmaxlνmax+L(N−νmin) results in the fastest convergence, and this translates to αt = 12 N
lνmax+L(N−νmin) . Next we claim that αt−1 = c 1 2 N lνmax+L(N−νmin) also guarantees
E(1− lαt−1(2l 2α2t + αtL)
1− αtµ )F (wt−1)− αt−1 2 ‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2 ≥ 0
Note that by scaling αt−1 by a constant c ≤ 1 if necessary, we can guarantee lαt−1(2l 2α2t+αtL)
1−αtµ ≤ 1 2 ,
and so the condition is equivalent to
F (wt−1)− αt−1‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2 ≥ 0
which was shown to hold with αt−1 ≤ 12 N lνmax+L(N−νmin) .
For the proof of general E ≥ 2, we use the following two identities:
‖gt‖2 ≤ 2‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2 + 2 N∑ k=1 pkl 2‖wt −wkt ‖2
E N∑ k=1 pk‖wt −wkt ‖2 ≤ E2(1 + 2l2α2t−1) N∑ k=1 pk‖wt−1 −wkt−1‖2 + 8α2t−1lF (wt−1)− 2α2t−1‖gt−1‖2
where the first inequality has been established before. To establish the second inequality, note that N∑ k=1 pk‖wt −wkt ‖2 = N∑ k=1 pk‖wt−1 − αt−1gt−1 −wkt−1 + αt−1gt−1,k‖2
≤ 2 N∑ k=1 pk ( ‖wt−1 −wkt−1‖2 + ‖αt−1gt−1 − αt−1gt−1,k‖2 ) and ∑
k pk‖gt−1,k − gt−1‖2 = ∑ k pk(‖gt−1,k‖2 − ‖gt−1‖2)
= ∑ k pk‖∇Fk(wt−1, ξkt−1) +∇Fk(wkt−1, ξkt−1)−∇Fk(wt−1, ξkt−1)‖2 − ‖gt−1‖2
≤ 2 ∑ k pk ( ‖∇Fk(wt−1, ξkt−1)‖2 + l2‖wkt−1 −wt−1‖2 ) − ‖gt−1‖2
so that using the l-smoothness of `, E N∑ k=1 pk‖wt −wkt ‖2
≤ E2(1 + 2l2α2t−1) N∑ k=1 pk‖wt−1 −wkt−1‖2 + 4α2t−1 ∑ k pk‖∇Fk(wt−1, ξkt−1)‖2 − 2α2t−1‖gt−1‖2
≤ E2(1 + 2l2α2t−1) N∑ k=1 pk‖wt−1 −wkt−1‖2 + 4α2t−12l ∑ k pk(Fk(wt−1, ξ k t−1)− Fk(w∗, ξkt−1))− 2α2t−1‖gt−1‖2
= E2(1 + 2l2α2t−1) N∑ k=1 pk‖wt−1 −wkt−1‖2 + 8α2t−1lF (wt−1)− 2α2t−1‖gt−1‖2
Using the first inequality, we have
E‖wt+1 −w∗‖2 ≤ E(1− αtµ)‖wt −w∗‖2
− 2αtF (wt) + 2α2t ‖ N∑ k=1 pk∇Fk(wt, ξkt )‖2
+ (2α2t l 2 + αtL) N∑ k=1 pk‖wt −wkt ‖2
and we choose αt and αt−1 such that E(F (wt)− αt‖ ∑N k=1 pk∇Fk(wt, ξkt )‖2) ≥ 0 and (2α2t l2 + αtL) ≤ (1− αtµ)(2α2t−1l2 + αt−1L)/3. This gives E‖wt+1 −w∗‖2 ≤ E(1− αtµ)[(1− αt−1µ)‖wt−1 −w∗‖2 − 2αt−1F (wt−1) + 2α2t−1‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2
+ (2α2t−1l 2 + αt−1L)( N∑ k=1 pk‖wt−1 −wkt−1‖2 + N∑ k=1 pk‖wt −wkt ‖2)/3]
Using the second inequality N∑ k=1 pk‖wt −wkt ‖2 ≤ E2(1 + 2l2α2t−1) N∑ k=1 pk‖wt−1 −wkt−1‖2 + 8α2t−1lF (wt−1)− 2α2t−1‖gt−1‖2
and that 2(1 + 2l2α2t−1) ≤ 3, 2α2t−1l2 + αt−1L ≤ 1, we have
E‖wt+1 −w∗‖2 ≤ E(1− αtµ)[(1− αt−1µ)‖wt−1 −w∗‖2
− 2αt−1F (wt−1) + 2α2t−1‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2 + 8α2t−1lF (wt−1)
+ (2α2t−1l 2 + αt−1L)(2 N∑ k=1 pk‖wt−1 −wkt−1‖2)]
and if αt−1 is chosen such that
(F (wt−1)− 4αt−1lF (wt−1))− αt−1‖ N∑ k=1 pk∇Fk(wt−1, ξkt−1)‖2 ≥ 0
and
(2α2t−1l 2 + αt−1L)(1− αt−1µ) ≤ (2α2t−2l2 + αt−2L)/3
we again have E‖wt+1 −w∗‖2 ≤ E(1− αtµ)(1− αt−1µ)[‖wt−1 −w∗‖2 + (2α2t−2l2 + αt−2L) · (2 N∑ k=1 pk‖wt−1 −wkt−1‖2)/3]
Applying the above derivation iteratively τ < E times, we have
E‖wt+1 −w∗‖2 ≤ E(1− αtµ) · · · (1− αt−τ+1µ)[(1− αt−τµ)‖wt−τ −w∗‖2
− 2αt−τF (wt−τ ) + 2α2t−τ‖ N∑ k=1 pk∇Fk(wt−τ , ξkt−τ )‖2 + 8τα2t−τ lF (wt−τ )
+ (2α2t−τ l 2 + αt−τL)((τ + 1) N∑ k=1 pk‖wt−τ −wkt−τ‖2)]
as long as the step sizes αt−τ are chosen such that the following inequalities hold
(2α2t−τ l 2 + αt−τL)(1− αt−τµ) ≤ (2α2t−τ−1l2 + αt−τ−1L)/3
2(1 + 2l2α2t−τ ) ≤ 3 2α2t−τ l 2 + αt−τL ≤ 1
(F (wt−τ )− 4ταt−τ lF (wt−τ ))− αt−τ‖ N∑ k=1 pk∇Fk(wt−τ , ξkt−τ )‖2 ≥ 0
We can check that setting αt−τ = c 1τ+1 N lνmax+L(N−νmin) for some small constant c satisfies the requirements.
Since communication is done every E iterations, wt0 = w k t0 for some t0 > t−E , from which we can conclude that
E‖wt −w∗‖2 ≤ ( t−t0−1∏ τ=1 (1− µαt−τ ))‖wt0 −w∗‖2
≤ (1− c µ E
N
lνmax + L(N − νmin) )t−t0‖wt0 −w∗‖2
and applying this inequality to iterations between each communication round,
E‖wt −w∗‖2 ≤ (1− c µ
E
N
lνmax + L(N − νmin) )t‖w0 −w∗‖2
= O(exp( µ
E
N
lνmax + L(N − νmin) t))‖w0 −w∗‖2
With partial participation, we note that
E‖wt+1 −w∗‖2 = E‖wt+1 − vt+1 + vt+1 −w∗‖2
= E‖wt+1 − vt+1‖2 + E‖vt+1 −w∗‖2
= 1
K ∑ k pkE‖wkt+1 −wt+1‖2 + E‖vt+1 −w∗‖2
and so the recursive identity becomes
E‖wt+1 −w∗‖2 ≤ E(1− αtµ) · · · (1− αt−τ+1µ)[(1− αt−τµ)‖wt−τ −w∗‖2
− 2αt−τF (wt−τ ) + 2α2t−τ‖ N∑ k=1 pk∇Fk(wt−τ , ξkt−τ )‖2 + 8τα2t−τ lF (wt−τ )
+ (2α2t−τ l 2 + αt−τL+
1
K )((τ + 1) N∑ k=1 pk‖wt−τ −wkt−τ‖2)]
which requires
(2α2t−τ l 2 + αt−τL+
1
K )(1− αt−τµ) ≤ (2α2t−τ−1l2 + αt−τ−1L+
1
K )/3
2(1 + 2l2α2t−τ ) ≤ 3
2α2t−τ l 2 + αt−τL+
1
K ≤ 1
(F (wt−τ )− 4ταt−τ lF (wt−τ ))− αt−τ‖ N∑ k=1 pk∇Fk(wt−τ , ξkt−τ )‖2 ≥ 0
to hold. Again setting αt−τ = c 1τ+1 N lνmax+L(N−νmin) for a possibly different constant from before satisfies the requirements.
Finally, using the L-smoothness of F ,
F (wT )− F (w∗) ≤ L
2 E‖wT −w∗‖2 = O(L exp(−
µ
E
N
lνmax + L(N − νmin) T ))‖w0 −w∗‖2

heading: H.2 GEOMETRIC CONVERGENCE OF FEDAVG FOR OVERPARAMETERIZED LINEAR REGRESSION
text: We first provide details on quantities used in the proof of results on linear regression in Section G. The local device objectives are now given by the sum of squares Fk(w) = 12nk ∑nk j=1(w Txjk − z j k)
2, and there exists w∗ such that F (w∗) ≡ 0. Define the local Hessian matrix as Hk := 1nk ∑nk j=1 x j k(x j k) T , and the stochastic Hessian matrix as H̃kt := ξ k t (ξ k t ) T , where ξkt is the stochastic sample on the kth device at time t. Define l to be the smallest positive number such that E‖ξkt ‖2ξkt (ξkt )T lHk for all k. Note that l ≤ maxk,j ‖xjk‖2. Let L and µ be lower and upper bounds of non-zero eigenvalues of Hk. Define κ1 := l/µ and κ := L/µ.
Following Liu & Belkin (2020); Jain et al. (2017), we define the statistical condition number κ̃ as the smallest positive real number such that E ∑ k pkH̃ k tH −1H̃kt ≤ κ̃H. The condition numbers κ1 and κ̃ are important in the characterization of convergence rates for FedAvg algorithms. Note that κ1 > κ and κ1 > κ̃.
Let H = ∑ k pkH
k. In general H has zero eigenvalues. However, because the null space of H and range of H are orthogonal, in our subsequence analysis it suffices to project wt −w∗ onto the range of H, thus we may restrict to the non-zero eigenvalue of H.
A useful observation is that we can use w∗Txjk − z j k ≡ 0 to rewrite the local objectives as Fk(w) = 1 2 〈w −w ∗,Hk(w −w∗)〉 ≡ 12‖w −w ∗‖2Hk :
Fk(w) = 1
2nk nk∑ j=1 (wTxk,j − zk,j − (w∗Txk,j − zk,j))2 = 1 2nk nk∑ j=1 ((w −w∗)Txk,j)2
= 1 2 〈w −w∗,Hk(w −w∗)〉 = 1 2 ‖w −w∗‖2Hk
so that F (w) = 12‖w −w ∗‖2H . Finally, note that EH̃kt = 1nk ∑nk j=1 x j k(x j k) T = Hk and gt,k = ∇Fk(wkt , ξkt ) = H̃kt (wkt − w∗)
while gt = ∑N k=1 pk∇Fk(wkt , ξkt ) = ∑N k=1 pkH̃ k t (w k t −w∗) and gt = ∑N k=1 pkH k(wkt −w∗)
Theorem 6. For the overparamterized linear regression problem, FedAvg with communication every E iterations with constant step size α = O( 1E N lνmax+µ(N−νmin) ) has geometric convergence:
EF (wT ) ≤ O ( L exp(− NT
E(νmaxκ1 + (N − νmin)) )‖w0 −w∗‖2
) .
Proof. We again show the result first when E = 2 and t− 1 is a communication round. We have ‖wt+1 −w∗‖2 = ‖(wt − αtgt)−w∗‖2
= ‖wt −w∗‖2 − 2αt〈wt −w∗,gt〉+ α2t ‖gt‖2
and
− 2αtEt〈wt −w∗,gt〉
= −2αt N∑ k=1 pk〈wt −w∗,∇Fk(wkt )〉
= −2αt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉 − 2αt N∑ k=1 pk〈wkt −w∗,∇Fk(wkt )〉
= −2αt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉 − 2αt N∑ k=1 pk〈wkt −w∗,Hk(wkt −w∗)〉
= −2αt N∑ k=1 pk〈wt −wkt ,∇Fk(wkt )〉 − 4αt N∑ k=1 pkFk(w k t )
≤ 2αt N∑ k=1 pk(Fk(w k t )− Fk(wt) + L 2 ‖wt −wkt ‖2)− 4αt N∑ k=1 pkFk(w k t )
= αtL N∑ k=1 pk‖wt −wkt ‖2 − 2αt N∑ k=1 pkFk(wt)− 2αt N∑ k=1 pkFk(w k t )
= αtL N∑ k=1 pk‖wt −wkt ‖2 − αt N∑ k=1 pk〈(wt −w∗),Hk(wt −w∗)〉 − 2αt N∑ k=1 pkFk(w k t )
and
‖gt‖2 = ‖ N∑ k=1 pkH̃ k t (w k t −w∗)‖2
= ‖ N∑ k=1 pkH̃ k t (wt −w∗) + N∑ k=1 pkH̃ k t (w k t −wt)‖2
≤ 2‖ N∑ k=1 pkH̃ k t (wt −w∗)‖2 + 2‖ N∑ k=1 pkH̃ k t (w k t −wt)‖2
which gives E‖wt+1 −w∗‖2 ≤ E‖wt −w∗‖2 − αt N∑ k=1 pk〈wt −w∗,Hkwt −w∗〉+ 2α2t ‖ N∑ k=1 pkH̃ k t (wt −w∗)‖2
+ αtL N∑ k=1 pk‖wt −wkt ‖2 + 2α2t ‖ N∑ k=1 pkH̃ k t (w k t −wt)‖2 − 2αt N∑ k=1 pkFk(w k t )
following Ma et al. (2018) we first prove that
E‖wt −w∗‖2 − αt N∑ k=1 pk〈(wt −w∗),Hk(wt −w∗)〉+ 2α2t ‖ N∑ k=1 pkH̃ k t (wt −w∗)‖2
≤ (1− N 8(νmaxκ1 + (N − νmin)) )E‖wt −w∗‖2
with appropriately chosen αt. Compared to the rate O( µNlνmax+L(N−νmin) ) = O( N νmaxκ1+(N−νmin)κ ) for general strongly convex and smooth objectives, this is an improvement as linear speedup is now available for a larger range of N .
We have
Et‖ N∑ k=1 pkH̃ k t (wt −w∗)‖2
= Et〈 N∑ k=1 pkH̃ k t (wt −w∗), N∑ k=1 pkH̃ k t (wt −w∗)〉
= N∑ k=1 p2kEt‖H̃kt (wt −w∗)‖2 + N∑ k=1 ∑ j 6=k pjpkEt〈H̃kt (wt −w∗), H̃ j t (wt −w∗)〉
= N∑ k=1 p2kEt‖H̃kt (wt −w∗)‖2 + N∑ k=1 ∑ j 6=k pjpkEt〈Hk(wt −w∗),Hj(wt −w∗)〉
= N∑ k=1 p2kEt‖H̃kt (wt −w∗)‖2 + N∑ k=1 N∑ j=1 pjpkEt〈Hk(wt −w∗),Hj(wt −w∗)〉 − N∑ k=1 p2k‖Hk(wt −w∗)‖2
= N∑ k=1 p2kEt‖H̃kt (wt −w∗)‖2 + ‖ ∑ k pkH k(wt −w∗)‖2 − N∑ k=1 p2k‖Hk(wt −w∗)‖2
≤ N∑ k=1 p2kEt‖H̃kt (wt −w∗)‖2 + ‖ ∑ k pkH k(wt −w∗)‖2 − 1 N νmin‖ ∑ k pkH k(wt −w∗)‖2
≤ 1 N νmax N∑ k=1 pkEt‖H̃kt (wt −w∗)‖2 + (1− 1 N νmin)‖ ∑ k pkH k(wt −w∗)‖2
≤ 1 N νmaxl N∑ k=1 pk〈(wt −w∗),Hk(wt −w∗)〉+ (1− 1 N νmin)‖ ∑ k pkH k(wt −w∗)‖2
= 1
N νmaxl〈(wt −w∗),H(wt −w∗)〉+ (1−
1
N νmin)〈wt −w∗,H2(wt −w∗)〉
using ‖H̃kt ‖ ≤ l. Now we have
E‖wt −w∗‖2 − αt N∑ k=1 pk〈(wt −w∗),Hk(wt −w∗)〉+ 2α2t ‖ N∑ k=1 pkH̃ k t (wt −w∗)‖2 =
〈wt −w∗, (I − αtH + 2α2t ( νmaxl
N H + N − νmin N H2))(wt −w∗)〉
and it remains to bound the maximum eigenvalue of
(I − αtH + 2α2t ( νmaxl
N H + N − νmin N H2))
and we bound this following Ma et al. (2018). If we choose αt < N2(νmaxl+(N−νmin)L) , then
−αtH + 2α2t ( νmaxl
N H + N − νmin N H2) ≺ 0
and the convergence rate is given by the maximum of 1−αtλ+2α2t (νmaxlN λ+ N−νmin N λ 2) maximized over the non-zero eigenvalues λ of H. To select the step size αt that gives the smallest upper bound, we then minimize over αt, resulting in
min αt<
N 2(νmaxl+(N−νmin)L)
max λ>0:∃v,Hv=λv
{ 1− αtλ+ 2α2t ( νmaxl
N λ+ N − νmin N λ2) } Since the objective is quadratic in λ, the maximum is achieved at either the largest eigenvalue λmax of H or the smallest non-zero eigenvalue λmin of H.
When N ≤ 4νmaxlL−λmin + 4νmin, i.e. when N = O(l/λmin) = O(κ1), the optimal objective value is achieved at λmin and the optimal step size is given by αt = N4(νmaxl+(N−νmin)λmin) . The optimal convergence rate (i.e. the optimal objective value) is equal to 1 − 18 Nλmin
(νmaxl+(N−νmin)λmin) = 1 − 1 8 N (νmaxκ1+(N−νmin)) . This implies that when N = O(κ1), the optimal convergence rate has a linear speedup in N . When N is larger, this step size is no longer optimal, but we still have 1− 18 N (νmaxκ1+(N−νmin)) as an upper bound on the convergence rate.
Now we have proved
E‖wt+1 −w∗‖2 ≤ (1− 1
8
N
(νmaxκ1 + (N − νmin)) )E‖wt −w∗‖2
+ αtL N∑ k=1 pk‖wt −wkt ‖2 + 2α2t ‖ N∑ k=1 pkH̃ k t (w k t −wt)‖2 − 2αt N∑ k=1 pkFk(w k t )
Next we bound terms in the second line using a similar argument as the general case. We have
2α2t ‖ N∑ k=1 pkH̃ k t (w k t −wt)‖2 ≤ 2α2t l2 N∑ k=1 pk‖wt −wkt ‖2
and
E N∑ k=1 pk‖wt −wkt ‖2 ≤ E2(1 + 2l2α2t−1) N∑ k=1 pk‖wt−1 −wkt−1‖2 + 8α2t−1lF (wt−1)
= 4α2t−1l〈wt−1 −w∗,H(wt−1 −w∗)〉 and if αt, αt−1 satisfy
αtL+ 2α 2 t ≤ (1−
1
8
N
(νmaxκ1 + (N − νmin)) )(αt−1L+ 2α
2 t−1)/3
2(1 + 2l2α2t−1) ≤ 3 αtL+ 2α 2 t ≤ 1
we have
E‖wt+1 −w∗‖2
≤ (1− 1 8
N
(νmaxκ1 + (N − νmin)) )[E‖wt−1 −w∗‖2 − αt〈wt−1 −w∗,Hwt−1 −w∗〉+ 2α2t ‖ N∑ k=1 pkH̃ k t (wt −w∗)‖2
+ (αt−1L+ 2α 2 t−1) · 2 N∑ k=1 pk‖wt−1 −wkt−1‖2 + 4α2t−1l〈wt−1 −w∗,H(wt−1 −w∗)〉]
and again by choosing αt−1 = c N8(νmaxl+(N−νmin)λmin) for a small constant c, we can guarantee that
E‖wt−1 −w∗‖2 − αt−1〈wt−1 −w∗,Hwt−1 −w∗〉
+2α2t−1‖ N∑ k=1 pkH̃ k t−1(wt−1 −w∗)‖2 + 4α2t−1l〈wt−1 −w∗,H(wt−1 −w∗)〉
≤ (1− c N 16(νmaxl + (N − νmin)λmin) )E‖wt−1 −w∗‖2
For general E, we have the recursive relation
E‖wt+1 −w∗‖2 ≤ E(1− c 1
8
N (νmaxκ1 + (N − νmin)) ) · · · (1− c 1 8τ
N
(νmaxκ1 + (N − νmin)) )[‖wt−τ −w∗‖2
− αt−τ 〈wt−τ −w∗,Hwt−τ −w∗〉+ 2α2t−τ‖ N∑ k=1 pkH̃ k t−τ (wt−τ −w∗)‖2 + 4τα2t−1l〈wt−1 −w∗,H(wt−1 −w∗)〉
+ (2α2t−τ l 2 + αt−τL)((τ + 1) N∑ k=1 pk‖wt−τ −wkt−τ‖2)]
as long as the step sizes are chosen αt−τ = c N4τ(νmaxl+(N−νmin)λmin) such that the following inequalities hold
(2α2t−τ l 2 + αt−τL) ≤ (1− αt−τµ)(2α2t−τ−1l2 + αt−τ−1L)/3
2(1 + 2l2α2t−τ ) ≤ 3 2α2t−τ l 2 + αt−τL ≤ 1
and
‖wt−τ −w∗‖2 − αt−τ 〈wt−τ −w∗,Hwt−τ −w∗〉
+ 2α2t−τ‖ N∑ k=1 pkH̃ k t−τ (wt−τ −w∗)‖2 + 4τα2t−1l〈wt−1 −w∗,H(wt−1 −w∗)〉
≤ (1− c N 8(τ + 1)(νmaxκ1 + (N − νmin)) )E‖wt−τ −w∗‖2
which gives
E‖wt −w∗‖2 ≤ (1− c 1
8E
N
(νmaxκ1 + (N − νmin)) )t‖w0 −w∗‖2
= O(exp(− 1 E
N
(νmaxκ1 + (N − νmin)) t))‖w0 −w∗‖2
and with partial participation, the same bound holds with a possibly different choice of c.

heading: H.3 GEOMETRIC CONVERGENCE OF FEDMASS FOR OVERPARAMETERIZED LINEAR REGRESSION
text: Theorem 7. For the overparamterized linear regression problem, FedMaSS with communication
everyE iterations and constant step sizes η1 = O( 1E N lνmax+µ(N−νmin) ), η2 = η1(1− 1κ̃ ) 1+ 1√
κ1κ̃
, γ = 1− 1√ κ1κ̃
1+ 1√ κ1κ̃
has geometric convergence: EF (wT ) ≤ O ( L exp(− NT
E(νmax √ κ1κ̃+ (N − νmin))
)‖w0 −w∗‖2 ) .
Proof. The proof is based on results in Liu & Belkin (2020) which originally proposed the MaSS algorithm. Note that the update can equivalently be written as
vkt+1 = (1− αk)vkt + αkukt − δkgt,k
wkt+1 = { ukt − ηkgt,k if t+ 1 /∈ IE∑N k=1 pk [ ukt − ηkgt,k ] if t+ 1 ∈ IE
ukt+1 = αk
1 + αk vkt+1 +
1
1 + αk wkt+1
where there is a bijection between the parameters 1−α k
1+αk = γk, ηk = ηk1 , ηk−αkδk 1+αk
= ηk2 , and we further introduce an auxiliary parameter vkt , which is initialized at v k 0 . We also note that when δk = η k
αk , the update reduces to the Nesterov accelerated SGD. This version of the FedAvg algorithm
with local MaSS updates is used for analyzing the geometric convergence. As before, define the virtual sequences wt = ∑N k=1 pkw k t , vt = ∑N k=1 pkv k t , ut = ∑N k=1 pku k t , and
gt = ∑N k=1 pkEgt,k. We have Egt = gt and wt+1 = ut−ηtgt, vt+1 = (1−αk)vt+αkwt−δkgt, and ut+1 = α k
1+αk vt+1 + 1 1+αk wt+1.
We first prove the theorem with E = 2 and t− 1 being a communication round. We have
‖vt+1 −w∗‖2H−1 = ‖(1− α)vt + αut − δ ∑ k pkH̃ k t (u k t −w∗)−w∗‖2H−1
= ‖(1− α)vt + αut −w∗‖2H−1 + δ 2‖ ∑ k pkH̃ k t (u k t −w∗)‖2H−1
− 2δ〈 ∑ k pkH̃ k t (u k t −w∗), (1− α)vt + αut −w∗〉H−1 ≤ ‖(1− α)vt + αut −w∗‖2H−1︸ ︷︷ ︸ A + 2δ2‖ ∑ k
pkH̃ k t (ut −w∗)‖2H−1︸ ︷︷ ︸ B
+ 2δ2‖ ∑ k pkH̃ k t (ut − ukt )‖2H−1
−2δ〈 ∑ k pkH̃ k t (ut −w∗), (1− α)vt + αut −w∗〉H−1︸ ︷︷ ︸
C − 2δ〈 ∑ k pkH̃ k t (u k t − ut), (1− α)vt + αut −w∗〉H−1
Following the proof in Liu & Belkin (2020),
EA ≤ E(1− α)‖vt −w∗‖2H−1 + α‖ut −w ∗‖2H−1
≤ E(1− α)‖vt −w∗‖2H−1 + α µ ‖ut −w∗‖2
using the convexity of the norm ‖ · ‖H−1 and that µ is the smallest non-zero eigenvalue of H . Now
EB ≤ 2δ2(νmax 1
N κ̃+ N − νmin N )‖(ut −w∗)‖2H
using the folowing bound:
E (∑ k pkH̃ k t ) H−1 (∑ k pkH̃ k t ) = E ∑ k p2kH̃ k tH −1H̃kt + ∑ k 6=j pkpjH̃ k tH −1H̃jt
νmax 1 N E ∑ k pkH̃ k tH −1H̃kt + ∑ k 6=j pkpjH kH−1Hj
= νmax 1 N E ∑ k pkH̃ k tH −1H̃kt + ∑ k,j pkpjH kH−1Hj − ∑ k p2kH kH−1Hk
νmax 1 N E ∑ k pkH̃ k tH −1H̃kt + H− 1 N νmin ∑ k pkH kH−1Hk
νmax 1 N E ∑ k pkH̃ k tH −1H̃kt + H− 1 N νmin( ∑ k pkH k)H−1( ∑ k pkH k)
= νmax 1 N E ∑ k pkH̃ k tH −1H̃kt + N − νmin N H
νmax 1
N κ̃H + N − νmin N H
where we have used E ∑ k pkH̃ k tH −1H̃kt ≤ κ̃H by definition of κ̃ and the operator convexity of the mapping W →WH−1W . Finally,
EC = −E2δ〈 ∑ k pkH̃ k t (ut −w∗), (1− α)vt + αut −w∗〉H−1
= −2δ〈 ∑ k pkH k(ut −w∗), (1− α)vt + αut −w∗〉H−1 = −2δ〈(ut −w∗), (1− α)vt + αut −w∗〉 = −2δ〈(ut −w∗),ut −w∗ + 1− α α (ut −wt)〉 = −2δ‖ut −w∗‖2 + 1− α α δ(‖wt −w∗‖2 − ‖ut −w∗‖2 − ‖wt − ut‖2)
≤ 1− α α δ‖wt −w∗‖2 − 1− α α δ‖ut −w∗‖2
where we have used
(1− α)vt + αut = (1− α)((1 + α)ut −wt)/α+ αut
= 1
α ut − 1− α α wt
and the identity that −2〈a,b〉 = ‖a‖2 + ‖b‖2 − ‖a + b‖2. It follows that
E‖vt+1 −w∗‖2H−1
≤ (1− α)‖vt −w∗‖2H−1 + 1− α α δ‖wt −w∗‖2
+ ( α µ − 1− α α δ)‖ut −w∗‖2 + 2δ2(νmax 1 N κ̃+ N − νmin N )‖(ut −w∗)‖2H
+ 2δ2‖ ∑ k pkH̃ k t (ut − ukt )‖2H−1
− 2δ〈 ∑ k pkH̃ k t (u k t − ut), (1− α)vt + αut −w∗〉H−1
On the other hand, E‖wt+1 −w∗‖2 = E‖ut −w∗ − η ∑ k pkH̃ k t (ut −w∗)‖2
= E‖ut −w∗‖2 − 2η‖ut −w∗‖2H + η2‖ ∑ k pkH̃ k t (ut −w∗)‖2
≤ E‖ut −w∗‖2 − 2η‖ut −w∗‖2H + η2(νmax 1
N `+ L N − νmin N )‖ut −w∗‖2
where we use the following bound:
E (∑ k pkH̃ k t )(∑ k pkH̃ k t ) = E
∑ k p2kH̃ k t H̃ k t + ∑ k 6=j pkpjH̃ k t H̃ j t
νmax 1 N E ∑ k pkH̃ k t H̃ k t + ∑ k 6=j pkpjH kHj
= νmax 1 N E ∑ k pkH̃ k t H̃ k t + ∑ k,j pkpjH kHj − ∑ k p2kH kHk
νmax 1 N E ∑ k pkH̃ k t H̃ k t + H 2 − 1 N νmin ∑ k pkH kHk
νmax 1 N E ∑ k pkH̃ k t H̃ k t + H 2 − 1 N νmin( ∑ k pkH k)( ∑ k pkH k)
= νmax 1 N E ∑ k pkH̃ k t H̃ k t + N − νmin N H2
νmax 1
N lH + L N − νmin N H
again using that W →W 2 is operator convex and that EH̃kt H̃kt lHk by definition of l.
Combining the bounds for E‖wt+1 −w∗‖2 and E‖vt+1 −w∗‖2H−1 ,
E δ
α ‖wt+1 −w∗‖2 + ‖vt+1 −w∗‖2H−1
≤ (1− α)‖vt −w∗‖2H−1 + 1− α α δ‖wt −w∗‖2 + ( α µ − δ)‖ut −w∗‖2
+ (2δ2(νmax 1
N κ̃+ N − νmin N )− 2ηδ/α+ η2δ(νmax 1 N l + L N − νmin N )/α)‖ut −w∗‖2
+ 2δ2‖ ∑ k pkH̃ k t (ut − ukt )‖2H−1
+ δL ∑ k pk‖(ut − ukt )‖2H−1
Following Liu & Belkin (2020) if we choose step sizes so that
α µ − δ ≤ 0
2δ2(νmax 1
N κ̃+ N − νmin N )− 2ηδ/α+ η2δ(νmax 1 N l + L N − νmin N )/α ≤ 0
or equivalently
α/δ ≤ µ
2αδ(νmax 1
N κ̃+ N − νmin N ) + η(η(νmax 1 N l + L N − νmin N )− 2) ≤ 0
the second and third terms are negative. To optimize the step sizes, note that the two inequalities imply
α2 ≤ η(2− η(νmax 1
N l + L N − νmin N ))µ/2(νmax 1 N κ̃+ N − νmin N )
and maximizing the right hand side with respect to η, which is quadratic, we see that η ≡ 1/(νmax 1N l+ LN−νminN ) maximizes the right hand side, with
α ≡ 1√ 2(νmax 1 N κ1 + κ N−νmin N )(νmax 1 N κ̃+ N−νmin N )
δ ≡ α µ = η
α(νmax 1 N κ̃+ N−νmin N )
Note that α = 1√ 2(νmax 1 N κ1+κ N−νmin N )(νmax 1 N κ̃+ N−νmin N ) = O( N√ κ1κ̃ ) when N = O(min{κ̃, κ1/κ}).
Finally, to deal with the terms 2δ2‖ ∑ k pkH̃ k t (ut − ukt )‖2H−1 + δL ∑ k pk‖(ut − ukt )‖2H−1 , we can use Jensen
2δ2‖ ∑ k pkH̃ k t (ut − ukt )‖2H−1 + δL ∑ k pk‖(ut − ukt )‖2H−1
≤ (2δ2l2 + δL) ∑ k pk‖ut − ukt ‖2H−1
= (2δ2l2 + δL) ∑ k pk‖ α 1 + α vt + 1 1 + α wt − ( α 1 + α vkt + 1 1 + α wkt )‖2H−1
≤ (2δ2l2 + δL)(2( α 1 + α )2δ2 + 2( 1 1 + α )2η2) ∑ k pk‖H̃kt−1(ut−1 −w∗)‖2
≤ (2δ2l2 + δL)(2( α 1 + α )2δ2 + 2( 1 1 + α )2η2)l2‖(ut−1 −w∗)‖2
which can be combined with the terms with ‖(ut−1−w∗)‖2 in the recursive expansion of E δα‖wt− w∗‖2 + ‖vt −w∗‖2H−1 :
E δ
α ‖wt −w∗‖2 + ‖vt −w∗‖2H−1
≤ (1− α)‖vt−1 −w∗‖2H−1 + 1− α α δ‖wt−1 −w∗‖2 + ( α µ − δ)‖ut−1 −w∗‖2
+ (2δ2(νmax 1
N κ̃+ N − νmin N )− 2ηδ/α+ η2δ(νmax 1 N l + L N − νmin N )/α)‖ut−1 −w∗‖2
and the step sizes can be chosen so that the resulting coefficients are negative. Therefore, we have shown that
E‖wt+1 −w∗‖2 ≤ (1− α)2‖wt−1 −w∗‖2
where α = 1√ 2(νmax 1 N κ1+κ N−νmin N )(νmax 1 N κ̃+ N−νmin N ) = O( N νmax √ κ1κ̃+N−νmin ) when N = O(min{κ̃, κ1/κ}).
For general E > 1, choosing η = c/E(νmax 1N l + L N−νmin
N ) for some small constant c results in α = O( 1
E √ (νmax 1 N κ1+κ N−νmin N )(νmax 1 N κ̃+ N−νmin N ) ) and this guarantees that
E‖wt −w∗‖2 ≤ (1− α)t‖w0 −w∗‖2
for all t.

heading: I DETAILS ON EXPERIMENTS AND ADDITIONAL RESULTS
text: We describe the precise procedure to reproduce the results in this paper. As we mentioned in Section 5, we empirically verified the linear speed up on various convex settings for both FedAvg and its accelerated variants. For all the results, we set random seeds as 0, 1, 2 and report the best convergence rate across the three folds. For each run, we initialize w0 = 0 and measure the number of iteration to reach the target accuracy . We use the small-scale dataset w8a Platt (1998), which consists of n = 49749 samples with feature dimension d = 300. The label is either positive one or negative one. The dataset has sparse binary features in {0, 1}. Each sample has 11.15 non-zero feature values out of 300 features on average. We set the batch size equal to four across all experiments. In the next following subsections, we introduce parameter searching in each objective separately.

heading: I.1 STRONGLY CONVEX OBJECTIVES
text: We first consider the strongly convex objective function, where we use a regularized binary logistic regression with regularization λ = 1/n ≈ 2e− 5. We evenly distributed on 1, 2, 4, 8, 16, 32 devices and report the number of iterations/rounds needed to converge to −accuracy, where = 0.005. The optimal objective function value f∗ is set as f∗ = 0.126433176216545. This is determined numerically and we follow the setting in Stich (2019). The learning rate is decayed as the ηt = min(η0, nc 1+t ), where we extensively search the best learning rate c ∈ {2 −1c0, 2 −2c0, c0, 2c0, 2
2c0}. In this case, we search the initial learning rate η0 ∈ {1, 32} and c0 = 1/8.

heading: I.2 CONVEX SMOOTH OBJECTIVES
text: We also use binary logistic regression without regularization. The setting is almost same as its regularized counter part. We also evenly distributed all the samples on 1, 2, 4, 8, 16, 32 devices. The figure shows the number of iterations needed to converge to −accuracy, where = 0.02. The optiaml objective function value is set as f∗ = 0.11379089057514849, determined numerically. The learning rate is decayed as the ηt = min(η0, nc1+t ), where we extensively search the best learning rate c ∈ {2−1c0, 2−2c0, c0, 2c0, 22c0}. In this case, we search the initial learning rate η0 ∈ {1, 32} and c0 = 1/8.

heading: I.3 LINEAR REGRESSION
text: For linear regression, we use the same feature vectors from w8a dataset and generate ground truth [w∗, b∗] from a multivariate normal distribution with zero mean and standard deviation one. Then we generate label based on yi = xtiw
∗+b∗. This procedure will ensure we satisfy the over-parameterized setting as required in our theorems. We also evenly distributed all the samples on 1, 2, 4, 8, 16, 32 devices. The figure shows the number of iterations needed to converge to −accuracy, where = 0.02. The optiaml objective function value is f∗ = 0. The learning rate is decayed as the ηt = min(η0, nc1+t ), where we extensively search the best learning rate c ∈ {2−1c0, 2−2c0, c0, 2c0, 22c0}. In this case, we search the initial learning rate η0 ∈ {0.1, 0.12} and c0 = 1/256.

heading: I.4 PARTIAL PARTICIPATION
text: To examine the linear speedup of FedAvg in partial participation setting, we evenly distributed data on 4, 8, 16, 32, 64, 128 devices and uniformly sample 50% devices without replacement. All other hyperparameters are the same as previous sections.

heading: I.5 NESTEROV ACCELERATED FEDAVG
text: The experiments of Nesterov accelerated FedAvg (the update formula is given as follows) uses the same setting as previous three sections for vanilia FedAvg.
ykt+1 = w k t − αtgt,k
wkt+1 =
{ ykt+1 + βt(y k t+1 − ykt ) if t+ 1 /∈ IE∑
k∈St+1
( ykt+1 + βt(y k t+1 − ykt ) ) if t+ 1 ∈ IE
We set βt = 0.1 and search αt in the same way as ηt in FedAvg.

heading: I.6 THE IMPACT OF E .
text: In this subsection, we further examine how does the number of local steps (E) affect convergence. As shown in Figure 2, the number of iterations increases asE increase, which slow down the convergence in terms of gradient computation. However, it can save communication costs as the number of rounds decreased when the E increases. This showcase that we need a proper choice of E to trade-off the communication cost and convergence speed.

Review summaries:
This paper shows a linear speedup in FedAvg w.r.t. number of devices, mainly theoretically, while most prior works ignore it. The main convergence results are given for three cases: a) Strongly Convex+Smooth, b) Convex+Smoth, and c) Strongly/x convex+Smooth+0 training loss can be achieved. The paper is well-written and motivated with good discussions of the algorithm and the related works. 

This paper gives convergence analysis for FedAvg and its accelerated version under data heterogeneity and system heterogeneity. The main improvement comes from a more careful analysis for one-step descent where the authors make use of the term $\alpha_{t} \sum_{k=1}^{N} p_{k}\left[F_{k}\left(\mathbf{w}^{*}\right)-F_{k}\left(\overline{\mathbf{w}}_{t}\right)\right]$ that is ignored by previous work Li (2020b). The paper focuses on how FedAvg’s convergence scales with the number of participating devices. It improves previous analysis for FedAvg under more federated settings and shows that FedAvg has linear speedup for any number of participating devices.