Year: 2021

Abstract:
A central question in federated learning (FL) is how to design optimization algorithms that minimize the communication cost of training a model over heterogeneous data distributed across many clients. A popular technique for reducing communication is the use of local steps, where clients take multiple optimization steps over local data before communicating with the server (e.g., FedAvg, SCAFFOLD). This contrasts with centralized methods, where clients take one optimization step per communication round (e.g., Minibatch SGD). A recent lower bound on the communication complexity of first-order methods shows that centralized methods are optimal over highly-heterogeneous data, whereas local methods are optimal over purely homogeneous data (Woodworth et al., 2020b). For intermediate heterogeneity levels, no algorithm is known to match the lower bound. In this paper, we propose a multistage optimization scheme that nearly matches the lower bound across all heterogeneity levels. The idea is to first run a local method up to a heterogeneity-induced error floor; next, we switch to a centralized method for the remaining steps. We demonstrate the scheme’s practical utility in a variety of image classification tasks.

heading: 1 INTRODUCTION
text: In federated learning (FL) (McMahan et al., 2017; Kairouz et al., 2019; Li et al., 2020), distributed clients interact with a central server to learn a model without directly sharing their data with the server. The training objective is to solve the following minimization problem:
min x
[ F (x) = 1
N ∑ i∈[N ] Fi(x) ]
(1)
where variable x is the model parameter, i indexes the clients (or devices), N is the number of clients, and Fi(x) is a (strongly convex) loss that only depends on that client’s data. Typical FL deployments have two properties that make optimizing Eq. (1) challenging: (i) Data heterogeneity: Let Fi(x) = Ezi∼Pi [f(x; zi)] for some loss function f , where zi is the client’s data and Pi’s capture the heterogeneity of the clients’ data distributions. Formally, we define client heterogeneity ζ2 and heterogeneity at the optimum ζ2∗ as follows, where the optimum is defined as x ∗ := arg minx F (x):
ζ2 := sup x
1
N ∑ i∈[N ] ‖∇F (x)−∇Fi(x)‖2, ζ2∗ := (1/N) ∑ i∈[N ] ‖∇Fi(x∗)‖2, (2)
These definitions capture how different the local client gradients are from the overall global gradient (worst-case) (Woodworth et al., 2020a). (ii) Communication cost: In many deployments of FL, clients have limited bandwidths (e.g., mobile devices) and hence high communication cost. Due to these two challenges, most federated optimization algorithms alternate between local rounds of computation, where clients process only their own data to save communication, and global rounds, where clients synchronize with the central server to resolve disagreements due to heterogeneity. We assume clients communicate with the server in each global round (full participation). This commonly-studied setting (Woodworth et al., 2020a; Gorbunov et al., 2020; Yuan et al., 2021; Mitra et al., 2021; Karimireddy et al., 2020b; Reddi et al., 2020) applies to cross-silo FL (Kairouz et al., 2019).
Several federated algorithms navigate this trade-off between reducing communication and resolving data heterogeneity by modifying the amount and the nature of local and global computation (McMahan et al., 2017; Li et al., 2018; Wang et al., 2019b;a; Li et al., 2019; Karimireddy et al., 2020b;a; Al-Shedivat et al., 2020; Reddi et al., 2020; Charles & Konečnỳ, 2020; Mitra et al., 2021). Recently, Woodworth et al. (2020a) showed a lower bound on the convergence rate (with respect to number of communication rounds R ) of any first-order federated optimization algorithms for strongly-convex, smooth objectives as a function of client data heterogeneity ζ∗ (Section 2, Eq. (6)). This lower bound highlights two gaps in the literature: (1) Today, no single optimization algorithm has been shown to match this lower bound for all values of ζ∗. (2) Although different baseline optimization algorithms match the lower bound in different heterogeneity regimes, it remains unknown if any algorithm can match the lower bound when 0 < ζ2∗ < β
3/2∆/µ1/2, where ∆ is the initial function suboptimality (see Assumption 6), and the functions Fi are β-smooth and µ-strongly convex (see App. C for definitions of smoothness and strong convexity). ∆ can be sizable if one starts optimization far away from the optimum, and the gap between the upper and the lower bound in this regime is significant (§ 2). To explain these gaps more precisely, we first examine several baseline optimization algorithms.
To begin, consider the simple baseline of Minibatch Stochastic Gradient Descent (SGD) (Woodworth et al., 2020a, shortened as SGD). Here, in each global communication round, the central server sends each device the current model. Then each device calculates a single local stochastic gradient at this model over a minibatch of size K and sends it to the server. The server averages these to make a global minibatch gradient over NK samples and performs a model update (Algorithm 2). Further, this approach can be accelerated in the Nesterov sense (Ghadimi & Lan, 2012, AC-SA). Woodworth et al. (2020a) recently showed that Accelerated Minibatch SGD (shortened as AC-SA) is nearly optimal (up to exponentially shrinking condition numbers factors in communication rounds R) in the highly-heterogeneous regimes when ζ2∗ > β
3/2∆/µ1/2. However, in other regimes, AC-SA may be suboptimal as it does not exploit local computation to reduce communication.
A popular alternative exploiting local computation is Federated Averaging (FedAvg) (McMahan et al., 2017)1 (Algorithm 6). Here each client i independently runs SGD on its own objective Fi for K iterations (local updates), and periodically communicates their respective most recent iterate to the server, which then averages these iterates. Compared to (Accelerated) SGD, this approach allows clients to learn from local data in between communication rounds. However, on strongly convex objectives, FedAvg has only been shown to achieve the lower bound when the data is homogeneous, i.e. ζ2∗ = ζ
2 = 0. In fact, FedAvg’s known rates can only outperform AC-SA’s rates when heterogeneity is almost insignificant: ζ2 ≤ µ , where is the target sub-optimality gap for Eq. (1) (Woodworth et al., 2020a) (for more details on see Definition 1). Further, an algorithm specific lower bound shows that FedAvg is provably sub-optimal; improving the analysis will not bring FedAvg meaningfully closer to optimal rates or even allow FedAvg to beat AC-SA when heterogeneity is non-negligible Woodworth et al. (2020a). Several approaches have been proposed to improve FedAvg via acceleration in the homogeneous setting Yuan & Ma (2020) and cross-device variance reduction (Karimireddy et al., 2020b; Gorbunov et al., 2020) and adaptivity and momentum (Reddi et al., 2020) in the heterogeneous setting. Empirically, these variants outperform baselines like (Accelerated) SGD, though gains in the heterogeneous setting are not explained by existing analyses.
In summary, currently there are only two baseline optimization algorithms that (nearly, up to exponentially shrinking condition numbers) theoretically match the lower bound of Woodworth et al. (2020a) in some heterogeneity regime: FedAvg when ζ2 = 0 (i.e., ζ2∗ = 0) and AC-SA when ζ2∗ > β
3/2∆/µ1/2. Several natural questions follow. (1) Does there exist an algorithm that achieves the lower bound when 0 < ζ2∗ < β
3/2∆/µ1/2? In fact, Woodworth et al. (2020a) posed a weaker open question, asking whether it was possible to improve on AC-SA in the regime where data heterogeneity was bounded but not insignificant. (2) If so, can the same algorithm simultaneously achieve optimal convergence rates across all values of ζ∗?
This Work. We answer both questions affirmatively when ζ ' ζ∗ by drawing on multistage algorithms, in which learning rates decay during optimization. Such algorithms have achieved fast convergence rates in convex optimization (Aybat et al., 2019; Fallah et al., 2020). In the federated setting, they have improved practical performance (Reddi et al., 2020; Charles & Konečnỳ, 2020).
1This approach is also referred to as Local Stochastic Gradient Descent (Stich, 2018; Khaled et al., 2020; Woodworth et al., 2020b), LocalUpdate (Charles & Konečnỳ, 2020).
Our work exploits the following insight: far from the optimum, local computation can reduce communication; near the optimum, it is important to communicate frequently (per model update) to resolve heterogeneity. This intuition has been explored in the homogeneous setting (Wang & Joshi, 2019), where the theoretical benefits do not appear to give order-wise convergence improvements.
Contributions. In this paper, we first show that multistage optimization, when applied to baseline federated algorithms, achieves comparable or better worst-case convergence rates than known baselines (including SGD and FedAvg) in all heterogeneity regimes. In doing so, multistage algorithms resolve an open problem from (Woodworth et al., 2020b), which asked if one can design an algorithm that combines the advantages of both FedAvg and SGD and enjoys guarantees that dominate both. We propose multistage algorithms that (nearly) match the lower bound of Woodworth et al. (2020a) on convergence rates for first-order federated algorithms.
We show these theoretical gains by first analyzing an algorithm that (nearly) achieves the lower bound: a simple two-stage combination of ASG (a variant of AC-SA in Theorem 6) and FedAvg. Loosely, we first run FedAvg up to an “optimization error floor” that depends on client heterogeneity; then, we use ASG to complete the optimization (Thm. 2). Further, we show that similar gains can be obtained by a more general multistage algorithm which gradually moves from the FedAvg update to ASG update. This multistage variant (Thm. 3) has the added benefit of not requiring the knowledge of the noise bound σ2, the heterogeneity ζ2, or the initial distance to the optimum ∆. This intuition is not specific to the combination of ASG and FedAvg. We demonstrate the generality of multistaging by applying it to other pairs of FL optimization algorithms to obtain the same nearly-optimal worst-case convergence (but possibly much better average-case rates, § 4). Finally, we demonstrate the practical value of these theoretical insights through both convex and nonconvex experiments in § 4.

heading: 2 MODEL AND RELATED WORK
text: We assume that each Fi is µ-strongly convex (Assumption 1) and β-smooth (Assumption 2). We let κ = β/µ denote the condition number. N is the number of clients; all clients participate in each communication round. A discussion of when a subset of clients participate per round is in App. L. Heterogeneity ζ and heterogeneity at optimum ζ∗ are defined as in equation 2 (Assumptions 3 and 4, respectively). We assume unless otherwise specified that ζ2 > 0, that the problem is heterogeneous. We assume the client gradient variance is upper bounded as σ2, formally defined in Assumption 5. We bound the initial suboptimality gap as ∆ (see Assumption 6).
Optimization proceeds in rounds; K is the number of client iterations per communication round, and R is the number of communication rounds. In some cases, clients take minibatches of size B in their local steps; unless specified otherwise, B is 1. Our goal is to minimize the sub-optimality error EF (x̂)− F (x∗) (where x̂ is the solution returned by an optimization algorithm) after R rounds (or to minimize communication round complexity, the number of rounds R to a given error ). R is a standard proxy for the communication cost. Since communication costs in FL are typically much larger than computation cost (Kairouz et al., 2019), our goal is to minimize round complexity R for sufficiently large K. We use the notation Õ to hide polylogarithmic factors.

heading: 2.1 RELATED WORK
text: To date, the fastest (worst-case) convergence rate in terms of R for our setting is achieved either by FedAvg (McMahan et al., 2017) or SGD in different heterogeneity regimes.
FedAvg. The convergence of FedAvg for convex minimization was first studied in the homogeneous client setting (Stich, 2018; Wang & Joshi, 2018; Woodworth et al., 2020b). These rates were later extended to the heterogeneous client setting (Khaled et al., 2020; Karimireddy et al., 2020b; Woodworth et al., 2020a; Koloskova et al., 2020), including a lower bound (Woodworth et al., 2020a). The fastest known rate for the worst-case sub-optimality gap of FedAvg is due to Gorbunov et al. (2020) (Table 1):
Õ(κ∆ exp(−KR κ ) + σ2 µNKR + κσ2 µKR2 + κζ2 µR2 ). (3)
If we take enough local steps K, this error is dominated by the last term, βζ 2
µ2R2 . If heterogeneity is low (e.g., ζ2 = 0), then for large enough K, only one round of communication is needed. This explains why FedAvg outperforms AC-SA in the low-heterogeneity regime.
Accelerated SGD (AC-SA). With highly heterogeneous clients, AC-SA is (nearly, up to exponentially shrinking condition numbers) optimal, with a convergence rate of (Ghadimi & Lan, 2012; Woodworth et al., 2020a)
O(∆ exp(− R√ κ ) + σ2 µNKR ) . (4)
Recall the FedAvg upper bound in Eq 3. Consider the noiseless case: σ2 = 0. By solving for the round complexity needed for FedAvg’s upper bound and AC-SA’s upper bound to reach accuracy, we see FedAvg’s communication R complexity is not less than that of the basic minibatch algorithms unless ζ2 ≤ µ log(∆/ε), i.e., heterogeneity must scale with . This is not an artifact of loose analysis as Woodworth et al. (2020a) provide the following lower bound for the convergence rate of FedAvg:
Ω( min{∆ exp(−R κ ), κζ2∗ µR2 }+ σ 2 µNKR + min{∆, κσ 2 µK2R2 }). (5)
Considering the noiseless case, AC-SA has a convergence rate of ∆ exp(−R/ √ κ) and FedAvg has the lower bound min{∆ exp(−R/κ), κζ2∗/µR2}. ∆ exp(−R/κ) is always worse than ∆ exp(−R/ √ κ),
so FedAvg can only ever outperform AC-SA when ∆ exp(−R/ √ κ) requires a largerR than κζ2∗/µR 2 to reach accuracy. Solving for R, this means that FedAvg’s lower bound can only result in lower communication complexity than AC-SA when ζ2∗ ≤ µ . So even in the absolute best case scenario, where the lower bound of (Woodworth et al., 2020a) is the convergence rate of FedAvg, heterogeneity still must scale with for it to outperform minibatch algorithms.
Lower Bound. The lower bound for first-order distributed algorithms in the heterogeneous setting, from (Woodworth et al., 2020a), is (also shown in Table 1):
Ω(min{ ∆√ κ , µζ2∗ β2 } exp(− R√ κ ) + σ2 µNKR ) . (6)
Note the exponential difference between the lower bound in Eq. (6) and the tightest-known ζdependent upper bound (i.e., FedAvg, Eq. (3)) with respect to their dependency on heterogeneity ζ2 and rounds R. In the upper bound (Eq. (3)), ζ2 is scaled by 1/R2, whereas in the lower bound (Eq. (6)), it is scaled by a term decaying exponentially fast in R assuming ζ ' ζ∗. A natural question is whether this gap can be closed.
Algorithm 1 Multistage federated minimization framework
Input: Number of stages S, round lengths {R(s)}Ss=1, inner stepsizes {{η (s,r) l }R (s) r=1 }Ss=1, outer stepsizes {{η(s,r)g }R (s)
r=1 }Ss=1, momentums {γ(s)}Ss=1, initial point x(0), weights for calculating weighted average solution from iterates {{{w(s,r)k }Kk=1}R (s)
r=1 }Ss=1 for s = 1, . . . , S do x(s,0), x(s,1) ← x(s−1) for r = 1, . . . , R(s) do . Broadcast the iterate from current round x(s,r) ← x(s) . Perform server update x(s,r+1/2) ← (1 + γ(s))x(s,r) − γ(s)x(s,r−1) . Obtain pseudo-gradient from clients G(s,r) ← ClientOPT(η(s,r)l , x(s,r+1/2)) x(s,r+1) ← x(s,r+1/2) − η(s,r)g G(s,r)
end for x(s) ← StageWeightedAverage({{w(s,r)k }Kk=1}R (s)
r=1 ) end for return x(S)
Closing the gap. Currently, no algorithm is known to close this gap, or even achieve comparable or faster rates than both AC-SA and FedAvg across all heterogeneity regimes. This may be because prior theoretical analysis focused mainly on three hyperparameter regimes: (i) local update stepsizes scale as 1/K (e.g., S-Local-SVRG (Gorbunov et al., 2020), FedAdam (Reddi et al., 2020), FedLin (Mitra et al., 2021)), (ii) local update stepsizes are small or zero (e.g., SCAFFOLD), and (iii) local update stepsizes do not depend on K (e.g., FedAvg). Intuitively, regimes (i) and (ii) do not take sufficient advantage of local steps to excel in the low-heterogeneity regime, while algorithms in regime (iii) may converge too quickly to local optima, leading to a severe dependency on heterogeneity. The only work outside these regimes (to our knowledge) considers quadratic objective functions (Charles & Konečnỳ, 2020) and does not achieve the FedAvg rate.
We show that this theoretical gap can be closed by combining existing methods (in stepsize regimes with known analyses) in stages. This multi-stage approach sequentially leverages the large local stepsizes of (iii) and the small local stepsizes of (i) and (ii) (Table 1).2
Multistage Algorithms. Multistage algorithms have been employed in the classical convex optimization setting to obtain optimal rates (Aybat et al., 2019; Fallah et al., 2020; Ghadimi & Lan, 2012; Woodworth et al., 2020a). The idea is to exponentially decrease the learning rate of an algorithm (e.g., AC-SA) in stages to balance the optimizer’s error due to bias and variance. Our federated setting differs from the prior multistage optimization in the following ways: (1) there is optimization error due to client drift (Reddi et al., 2020; Karimireddy et al., 2020b) in addition to error due to bias and variance; this additional type of error cannot be accounted for by existing multistage algorithms, (2) our objective is to minimize communication rounds rather than iteration complexity, and (3) we “multistage” or schedule more hyperparameters than just a learning rate.
Wang & Joshi (2019) decrease the number (K) of local updates to transition from FedAvg to SGD in the homogeneous setting. This is conceptually similar to the multistage algorithms we study. However, they do not show order-wise improvements in convergence rate. Indeed, in the homogeneous setting, recent work showed that a simple variant of AC-SA achieves the optimal worst-case convergence rate without any learning rate decay or multistage (Woodworth et al., 2021).

heading: 3 MULTISTAGE FEDERATED MINIMIZATION
text: 2Note that in Table 1, min in the FedAvg →SGD and FedAvg →ASG rates comes from using a metaalgorithm that runs SGD and ASG respectively if the first part of the min{·} is smaller. The min{·} in the M-FedAvg→M-ASG rate comes from running M-ASG and M-FedAvg→M-ASG in parallel and then evaluating objective values of both solutions and choosing the best one (this requires one extra round of communication).
We introduce a framework (Algorithm 1) for multistage federated minimization in the strongly convex setting. It proceeds over S stages; between stages, we can change the optimization algorithm and the associated hyperparameters (e.g., stepsizes). Each stage s ≤ S requires a local (client) optimization algorithm (e.g., FedAvg), a set of inner (client) and outer (server) stepsizes, a round duration R, a number of local client steps K, and other hyperparameters (e.g., momentum, client weights).
Local updates: For a given stage s ≤ S, we iterate over rounds r ∈ {1, . . . , R}. In a round r, client i first initializes their iterate as x(s,r)i,0 = x (s,r), and then runs K local update steps of its local optimizer using the local stepsize η(s)l . Next we give the different update rules used by different algorithms at client i. Here ∇̃Fi(x) denotes the possibly random gradient evaluated at x. FedAvg (Algorithm 6) and (Accelerated) SGD (Algorithms 2 and 3) at step k use the update rule:
x (s,r) i,k ← x (s,r) i,k−1 − η (s,r) l ∇̃Fi(x (s,r) i,k−1) (7)
However, (Accelerated) SGD sets the local stepsize as η(s,r)l = 0. This implies that x (s,r) i,k = x (s,r) for all 0 ≤ k ≤ K. Client variance-reduced methods like SS-Local-SGD (Gorbunov et al., 2020) (which is a stateless variant of SCAFFOLD, defined in Algorithm 9) at step k use the update rule:
x (s,r) i,k ← x (s,r) i,k−1 − η (s,r) l (∇̃Fi(x (s,r) i,k−1) + ∇̃F (x (s,r))− ∇̃F (x(s,r)))
Global round update: All the clients send back their effective updates to the server. Then the server takes an average over N effective updates to compute the pseudo-gradient G(s,r). The local updates and the collection process are abstracted by the step G(s,r) ← ClientOPT(η(s,r)l , x(s,r)), where
ClientOPT(η(s,r)l , x (s,r)) =
1
N N∑ i=1 K∑ k=1 ∇̃Fi(x(s,r)i,k−1).
For (Accelerated) SGD (Algorithms 2 and 3) we set η(s,r)l = 0, and so x (s,r) i,k = x (s,r). Therefore, we can simplify ClientOPT as a simple re-scaled (multiplied by K) minibatched stochastic gradient of the global function, with a batchsize of K, i.e.
ClientOPT(η(s,r)l = 0, x (s,r)) =
1
N N∑ i=1 K∑ k=1 ∇̃Fi(x(s,r))
Finally, the server takes a (accelerated) stochastic pseudo-gradient descent step using its round iterate history, and the computed pseudo-gradient G(s,r) and global stepsize η(s,r)g .
Stage update: Once R rounds are complete, the server completes the stage by taking a weighted average of its global iterates with weights {{{w(s,r)i,k }i≤N}k≤K}r≤R:
StageWeightedAverage({{w(s,r)k }k≤K}r≤R) = 1
W (s) R∑ r=1 K∑ k=1 w (s,r) k ( 1 N N∑ i=1 x (s,r) i,k ),
where W (s) = ∑R r=1 ∑K k=1 w (s,r) k . Sums of local iterates ∑K k=1 w (s,r) k x (s,r) i,k are computed and sent by clients to the server, which does not require access to local iterates. Although this weighted sum is commonly used to facilitate theoretical analysis (Karimireddy et al., 2020b; Gorbunov et al., 2020), it is often unnecessary in practice (i.e., we can just output the last iterate). The weights are determined by strong convexity µ and stepsizes, and do not require further tuning.

heading: 3.1 TWO STAGE ALGORITHMS
text: We begin with a two-stage instantiation of Algorithm 1. We run FedAvg for a constant fraction of rounds until reaching the “heterogeneity floor”, followed by a minibatch algorithm until convergence.
Intuition: Fig. 1 shows a toy example. Consider two strongly convex client objectives: F1(x) = 1 2 (x − 1) 2 and F2(x) = (x + 1)2. The global objective F (x) = F1(x)+F2(x)
2 is their average. Fig. 1 (top) plots the objectives, and Fig. 1 (bottom) displays their gradients.
Far from the optimum, due to strong convexity, all client gradients point towards the optimal solution; specifically, this occurs when x ∈ (−∞,−1] ∪ [1,∞). In this regime, clients can use local steps (e.g., FedAvg) to save on communication while still ensuring that perclient local updates are consistent with each other. On the other hand, close to the optimum, i.e., when x ∈ (−1, 1), some client gradients may point away from the optimum. This suggests that clients should not take local steps to avoid driving the global solution away from the optimum. Therefore, when we are close to the optimum, we use an algorithm without local steps (e.g. SGD), since such algorithms are less affected by client gradient disagreement. This intuition holds more generally because we have assumed strong convexity.
Our observation that client gradients disagree near the optimum has been observed in practice even for nonconvex objective functions. Charles et al. (2021) observe that over the course of FedAvg execution on neural network StackOverflow next word prediction, client gradients become more orthogonal from each other.
Our technique parallels multistage algorithms for obtaining optimal dependence on noise in convex optimization (Aybat et al., 2019). There, a large initial stepsize reduces the bias term (i.e. exponential decay term) in the convergence rate; the stepsize is then decayed to reduce the noise term until convergence. It also relates to the approach of (Wang & Joshi, 2019), which reduces the number of local steps during training.
In this subsection, we find that two stages suffice to achieve low round complexities if ∆, ζ2, σ2 are known. In the next section, we show how to relx this assumption by using more stages. We first show how to combine FedAvg with SGD to obtain exponential round complexity gains over either baseline.
Theorem 1 (Informal version of Thm. 4 in Appendix). If we run Algorithm 1 over S = 2 stages with with large enough number of local rounds K (specified in Thm. 4), such that (i) the first stage runs FedAvg forR/2 rounds under some parameter settings, and (ii) the second stage runs SGD (γ(2) = 0) for R/2 rounds under other parameter settings, then output of the second stage x(2) satisfies
EF (x(2))− F (x∗) ≤ Õ(( β 2ζ2 µ3R2 ) exp(−R κ ) + σ2 µNKR ) (8)
(Proof in App. D) This algorithm exponentially improves the dependence of the round complexity on heterogeneity from the state of the art βζ 2
µ2R2 to ( β2ζ2 µ3R2 ) exp(− R κ ). Next, we obtain our main
result, which tightens the dependency on κ to match the lower bound in Eq. (6) (up to exponentially shrinking condition number factors and supposing ζ ' ζ∗) by replacing SGD with ASG. Theorem 2 (Informal version of Thm. 5 in Appendix). If we run Algorithm 1 over S = 2 stages with large enough number of local rounds K (specified in Thm. 5), such that (i) the first stage runs FedAvg as for R/2 rounds under some parameter settings, and (ii) the second stage runs ASG for R/2 rounds under some other parameter settings, then output of this second stage x(2) satisfies
EF (x(2))− F (x∗) ≤ Õ( βζ 2 µ2R2 exp(− R√ κ ) + σ2 µNKR +
βσ2
µ2NKR3 ). (9)
(Proof in Section E). Observe that this rate is strictly better than the FedAvg rate (Eq. (3)) and is also strictly better than the AC-SA rate (Eq. (4)) when βζ 2
µ2R2 ≤ ∆ (since the AC-SA rate requires that R ≥ √ κ). In practice, the initialization can be far from the optimum, leading to large ∆. The same rates can be achieved when using SS-Local-SGD instead of FedAvg as the federated optimization algorithm in the first stage of both of the above two-stage algorithms.3

heading: 3.2 MULTISTAGE ALGORITHMS
text: The previous algorithms require knowledge of several data- and problem-dependent parameters, such as β, µ, σ2, ζ2, and ∆. Although the FL literature commonly assumes knowledge of such parameters (Woodworth et al., 2020a; Karimireddy et al., 2020b; Gorbunov et al., 2020), this may be unrealistic in practice. In this section, we explore how to use multistage algorithms to remove dependencies on all such parameters except β, which can be estimated from line search techniques (Beck & Teboulle, 2009; Schmidt et al., 2015), and µ, which can be estimated from regularization (which is often used in machine learning problems). The key idea is to divide our total stages S into two superstages, which are evenly split amont the R total rounds: the first superstage uses one optimization algorithm (e.g., FedAvg, or SS-Local-SGD) with a stepsize decay schedule, and the second superstage uses a centralized algorithm with a stepsize decay schedule (e.g., SGD, or ASG).
Theorem 3 (Informal version of Theorem 7 in Appendix). Let Algorithm 1 run over two superstages for a total of R rounds and S stages, with R ≥ 4 √ κ and K = dκe. In the first superstage, run Multistage FedAvg (M-FedAvg) (see App. H for settings, pseudocode in Algorithm 7) for R/2 rounds. In the second superstage, run Multistage Accelerated Stochastic Gradient (M-ASG) (Aybat et al., 2019) (pseudocode in Algorithm 5) with some other parameter choices and batch-size B. Then the output after S stages x(S) has convergence rate
EF (x(S))− F (x∗) ≤ O((κ∆ exp(−R) + β 2ζ2
µ3R2 +
βσ2
µ2BR2 ) exp(− R√ κ ) +
σ2
µBNR .
(Proof in App. G) For simplicity, suppose ζ2 > 0, R ≥ log( µ 4∆
β3ζ2 ), and σ 2 = 0 (where R is large
enough for κ∆ exp(−R) to match β 2ζ2
µ3R2 ; this is a relatively few number of rounds). Then this multistage algorithm has a better convergence rate with respect to communication cost than both FedAvg and AC-SA when 0 < ζ2 < µ
2∆ β (note that both algorithms require R ≥ Ω(
√ κ)). This is
done without knowledge of ζ2, σ2, or ∆. However, note we do not achieve the optimal rate like in Thm. 2 due to the additional term of κ∆ exp(−R).
3Similar rates for both Thm. 1 and Thm. 2 can be obtained by running the first stage for one round and the second stage for R− 1 rounds. This removes the 1
R2 factor in front of the exponential terms while increasing
the exponential shrinkage in the exponential terms by a factor of 2 compared to splitting the rounds evenly.

heading: 4 EXPERIMENTS
text: We evaluate the utility of our framework on strongly convex and nonconvex settings. We compare the communication round complexities of four single-stage baselines: FedAvg, SGD, ASG, and SCAFFOLD. We compare these baselines to both two-stage variants and multistage variants. Due to computational constraints, we consider the low-round (R) regime, where we cannot match state-ofthe-art test accuracy (e.g., Reddi et al. (2020)); our goal is instead to show how multistage optimization compares to single-stage baselines for fixed R.
Convex Optimization (Logistic Regression) We first study federated regularized logistic regression, which is strongly convex (objective function in App. K.1). In this experiment, we use the MNIST dataset of handwritten digits (LeCun et al., 2010). We (roughly) control client heterogeneity by creating “clients” by shuffling samples from different digit classes. The details of this shuffling process are described in App. K.1; if a dataset is more shuffled (more homogeneous), it roughly corresponds to a lower heterogeneity dataset.
Fig. 2 compares the convergence of multi-stage algorithms to their single-stage counterparts in the stochastic gradient setting (minibatches are 1% of a client’s data), over R = 100 rounds (tuning details in App. K.1). We observe that in all heterogeneity levels, multistage algorithms (whether two or more stages) outperform other baselines. We also observe that SCAFFOLD→SGD outperforms all other curves, including SCAFFOLD→ASG. This can be explained by the fact that acceleration introduces additional noise terms in the stochastic setting as shown in Thm. 5, which can contribute to error if one does not take K large enough (we set K = 20 in the convex experiments). The effect of large K is elaborated upon in App. M.
Nonconvex Optimization (Neural Networks) We also evaluated nonconvex image classification tasks with convolutional neural networks. We started with digit classification over EMNIST (Cohen et al., 2017) (tuning details in App. K.2.1), where handwritten characters are partitioned by author. Figure 3 (Left) displays the training losses. The multistage (S > 2) algorithms do not perform as well as two-stage algorithms, possibly because we use the same stepsize schedule for both superstages to keep the number of hyperparameters low. While successful in the convex setting, this approach does not work as well in our nonconvex experiments, possibly due to increased task difficulty. We also considered image classification with ResNet-18 on CIFAR-100 (Krizhevsky, 2009) (tuning details in App. K.2.2). Figure 3 (Middle) displays the training losses on CIFAR-100; SCAFFOLD experiments are not included due to memory constraints. The two-stage algorithms again perform the best. In Fig. 3 (Right), we display the test accuracies on CIFAR-100; two-stage FedAvg→ SGD achieves the highest test accuracy (test accuracies for EMNIST in Table 2).

heading: 5 CONCLUSION
text: In this paper, we show that multistage parameter schedules improve FL communication complexity in theory and in practice. In App. A, we discuss why our theoretical results may help explain the benefits of multistaging in practice. Future work includes tightening our upper bounds from ζ2 to ζ2∗ and getting rid of the additional κ∆ exp(−R) term in Thm. 7.

heading: A DISCUSSION: STEPSIZE DECAY IN PRACTICE
text: In this section, we make the case that our strategy of switching from a local step-utilizing algorithm (like FedAvg) to a centralized algorithm (like SGD) is similar to the stepsize decay strategies in (Charles & Konečnỳ, 2020), thus potentially explaining the empirical gains they see in their papers.
We expand on the learning rate decay of LocalUpdateDecay (Charles & Konečnỳ, 2020). In LocalUpdateDecay, the decay scheme is roughly as follows: if the loss decreases too little, then reduce the outer stepsize by a factor of 0.9 and the inner stepsize by a factor of 0.1. It is easy to see that in a relatively few number of decay events, their algorithm will essentially have switched from a local step-utilizing algorithm to a centralized algorithm.
We plot the ratio between LocalUpdateDecay’s local and global learning rates over the course of execution and versus the same ratio for a centralized algorithm and FedAvg (Algorithm 6) in Fig. 4. ηl/ηg = 0 on the centralized algorithm line because ηl = 0 for a centralized algorithm. For FedAvg (Algorithm 6), ηl/ηg = 1 because ηl = ηg . We see that within 3 decay events the ratio is nearly that of the centralized algorithm. Therefore, it is plausible that the practical improvements they observe from their decay scheme can largely be explained by our results.

heading: B ALGORITHM PSEUDOCODE
text: Algorithm 2 SGD Input: Minibatch size K, rounds R, stepsize η, initial point x(0), number of devices N for r = 0, . . . , R− 1 do
Server sends x(r) to all clients . Run a step of SGD for client i = 1, . . . , N in parallel do
Sample K datapoints {zi,k}1≤k≤K Client communicates ∑K k=1∇f(x(r); zi,k) to server
end for x(r+1) ← x(r) − ηNK ∑N i=1 ∑K k=1∇f(x(r); zi,k)
end for return x(R)

heading: B.1 ALGORITHM 2 IN TERMS OF ALGORITHM 1
text: Set η(1,r)g = η/K η (1,r) l = 0, S = 1, γ (1) = 0, and R(1) = R. For StageWeightedAverage simply return the last iterate x(1,R (1)).
Algorithm 3 ASG Input: Minibatch size K, rounds R, stepsize η, initial point x(0) = x(1), strong convexity µ number of devices N for r = 0, . . . , R− 1 do
Server sends x(r) to all clients . Run a step of ASG for client i = 1, . . . , N in parallel do
Sample K datapoints {zi,k}1≤k≤K Client communicates ∑K k=1∇f(x(r); zi,k) to server
end for γ(r) ← 1− √ ηµ
1+ √ ηµ
y(r) ← (1 + γ(r))x(r) − γ(r)x(r−1) x(r+1) ← y(r) − ηNK ∑N i=1 ∑K k=1∇f(y(r); zi,k)
end for return x(R)

heading: B.2 ALGORITHM 3 IN TERMS OF ALGORITHM 1
text: Set η(1,r)g = η/K η (1,r) l = 0, S = 1, γ (1) = 1−√ηµ 1+ √ ηµ where µ is strong convexity, and R (1) = R. For StageWeightedAverage simply return the last iterate x(1,R (1)).
Algorithm 4 M-SGD Input: Minibatch size K, round sequence {R(s)}s≤S , stepsize schedule {η(s)}s≤S , initial point x(0), number of devices N for s = 1, . . . , S do
for r = 0, . . . , R(s) − 1 do Server sends x(s,r) to all clients . Run a step of SGD for client i = 1, . . . , N in parallel do
Sample K datapoints {zi,k}1≤k≤K Client communicates ∑K k=1∇f(x(s,r); zi,k) to server
end for x(s,r+1) ← x(s,r) − η (s)
NK ∑N i=1 ∑K k=1∇f(x(s,r); zi,k)
end for end for return x(S,R (S))

heading: B.3 ALGORITHM 4 IN TERMS OF ALGORITHM 1
text: Set η(s,r)g = η(s)/K, η (1,r) l = 0, same number of stages S, γ (s) = 0, and same round sequence R(s). For StageWeightedAverage simply return the last iterate x(S,R (S)).

heading: B.4 ALGORITHM 5 IN TERMS OF ALGORITHM 1
text: Set η(s,r)g = η(s)/K η (1,r) l = 0, same number of stages S, γ
(s) = 1− √ η(s)µ
1+ √ η(s)µ where µ is strong
convexity, and same round sequence R(s). For StageWeightedAverage simply return the last iterate x(S,R (S)).
Algorithm 5 M-ASG Input: Minibatch size K, round sequence {R(s)}s≤S , stepsize schedule {η(s)}s≤S , initial point x(0), strong convexity µ number of devices N for s = 1, . . . , S do x(s,0), x(s,1) ← x(s−1) for r = 1, . . . , R(s) do
Server sends x(s,r) to all clients . Run a step of ASG for client i = 1, . . . , N in parallel do
Sample K datapoints {zi,k}1≤k≤K Client communicates ∑K k=1∇f(x(s,r); zi,k) to server
end for γ(s) ← 1− √ η(s)µ
1+ √ η(s)µ
y(s,r) ← (1 + γ(s,r))x(s,r) − γ(s,r)x(s,r−1) x(s,r+1) ← y(s,r) − η (s)
NK ∑N i=1 ∑K k=1∇f(y(s,r); zi,k)
end for x(s) ← x(s,R(s)+1)
end for return x(S,R (s))

heading: B.5 ALGORITHM 6 IN TERMS OF ALGORITHM 1
text: Set η(s,r)g = η (s,r) l = η, S = 0, γ (s,r) = 0, R(1) = R. x(s,r)i,k = x (r) i,k . w (1,r) k = (1 − µη)rK+k+1. For ClientOPT,
ClientOPT(η(s,r)l , x (s,r)) =
1
N N∑ i=1 K∑ k=1 ∇̃Fi(x(s,r)i,k−1).
where the ∇̃Fi(x(s,r)i,k−1)’s come from
x (s,r) i,k ← x (s,r) i,k−1 − η (s,r) l ∇̃Fi(x (s,r) i,k−1) (10)
Algorithm 6 FedAvg Input: Number of local steps K, rounds R, stepsize η, initial point x(0), number of devices N for r = 0, . . . , R− 1 do
Server communicates x(r) to all clients for client i = 1, . . . , N in parallel do x
(r) i,0 ← x(r)
for k = 1, . . . , K do Randomly select sample zi,k−1 x
(r) i,k ← x (r) i,k−1 − η · ∇f(x (r) i,k−1; zi,k−1)
end for Client communicates ∑K k=1∇f(x (r) i,k−1; zi,k−1) to server
end for . Perform server update x(r+1) ← x(r) − ηN ∑ i∈[N ] ∑K k=1∇f(x (r) i,k−1; zi,k−1)
end for return x(R)

heading: B.6 ALGORITHM 7 IN TERMS OF ALGORITHM 1
text: Set η(s,r)g = η (s,r) l = η (s), same number of stages S, same round sequence {R(s)}s≤S , γ(s,r) = 0, R(1) = R. x(s,r)i,k = x (r) i,k . For StageWeightedAverage simply return the last iterate x
(S,R(S)). For ClientOPT,
ClientOPT(η(s,r)l , x (s,r)) =
1
N N∑ i=1 K∑ k=1 ∇̃Fi(x(s,r)i,k−1).
where the ∇̃Fi(x(s,r)i,k−1)’s come from
x (s,r) i,k ← x (s,r) i,k−1 − η (s,r) l ∇̃Fi(x (s,r) i,k−1) (11)
Algorithm 7 M-FedAvg Input: Number of local steps K, round sequence {R(s)}s≤S where R(0) = 0, stepsize schedule {η(s)}s≤S , stepsize η, initial point x(0,0), number of devices N for s = 1, . . . , S do x(s,r) ← x(s−1,R(s−1)) for r = 0, . . . , R(s) − 1 do
Server communicates x(s,r) to all clients for client i = 1, . . . , N in parallel do x
(s,r) i,0 ← x(s,r)
for k = 1, . . . , K do Randomly select sample zi,k−1 x
(s,r) i,k ← x (s,r) i,k−1 − η(s) · ∇f(x (s,r) i,k−1; zi,k−1)
end for Client communicates ∑K k=1∇f(x (s,r) i,k−1; zi,k−1) to server
end for . Perform server update
x(s,r+1) ← x(s,r) − η (s)
N ∑ i∈[N ] ∑K k=1∇f(x (s,r) i,k−1; zi,k−1)
end for end for return x(R)

heading: B.7 ALGORITHM 8 IN TERMS OF ALGORITHM 1
text: Set η(s,r)g = η (s,r) l = η, S = 0, γ (s,r) = 0, R(1) = R. w(1,r)k = (1 − µη)rK+k+1. x (s,r) i,k = x (r) i,k . For ClientOPT,
ClientOPT(η(s,r)l , x (s,r)) =
1
N N∑ i=1 K∑ k=1 ∇̃Fi(x(s,r)i,k−1).
where the ∇̃Fi(x(s,r)i,k−1)’s come from
x (s,r) i,k ← x (s,r) i,k−1 − η (s,r) l (∇̃Fi(x (s,r) i,k−1 + c− ci)) (12)
where ci = 1K ∑K i=1 ∇̃Fi(x (s,r−1) i,k ) if the previous round is in the same stage and ci =
1 K ∑K i=1 ∇̃Fi(x (s−1,R(s−1)) i,k ) if the previous round is in an earlier stage. c = 1 N ∑N i=1 ci. In Algorithm 1 we do not specify states (for c and ci), though one can easily include them.
Algorithm 8 SCAFFOLD Input: Number of local steps K, rounds R, stepsize ηl, ηg , initial point x(0), control variates c and ci, number of devices N for r = 0, . . . , R− 1 do
Server communicates x(r) to all clients Server communicates c to all clients for client i = 1, . . . , N in parallel do x
(r) i,0 ← x(r)
for k = 1, . . . , K do Randomly select sample zi,k−1 x
(r) i,k ← x (r) i,k−1 − ηl · (∇f(x (r) i,k−1; zi,k−1) + c− ci)
end for ci ← 1K ∑K−1 k=0 ∇f(x (r) i,k−1; zi,j)
Client communicates ∑K−1 k=0 ∇f(x (r) i,k−1; zi,k) to server
Client communicates ci to server end for . Perform server update x(r+1) ← x(r) − ηgN ∑ i∈[N ] ∑K−1 k=0 ∇f(x (r) i,k−1; zi,k)
c← 1N ∑N i=1 ci
end for return x(R)
Algorithm 9 SCAFFOLD (SS-Local-SGD) Input: Number of local steps K, rounds R, stepsize ηl, ηg , initial point x(0), number of devices N for r = 0, . . . , R− 1 do
Server communicates x(r) to all clients For client i ∈ [n] compute ci ← 1K ∑K i=1∇f(x(r); zi,k) for randomly drawn zi,k Clients communicate ci’s to server c← 1N ∑N i=1 ci Server communicates c to clients for client i = 1, . . . , N in parallel do x
(r) i,0 ← x(r)
for k = 1, . . . , K do Randomly select sample zi,k−1 x
(r) i,k ← x (r) i,k−1 − η · (∇f(x (r) i,k−1; zi,k−1) + c− ci)
end for ci ← 1K ∑K−1 k=0 ∇f(x (r) i,k−1; zi,j)
Client communicates ∑K−1 k=0 ∇f(x (r) i,k−1; zi,k) to server
end for . Perform server update x(r+1) ← x(r) − ηN ∑ i∈[N ] ∑K−1 k=0 ∇f(x (r) i,k−1; zi,k)
end for return x(R)

heading: B.8 ALGORITHM 9 IN TERMS OF ALGORITHM 1
text: Set η(s,r)g = η (s,r) l = η, S = 0, γ (s,r) = 0, R(1) = R. x(s,r)i,k = x (r) i,k . w (1,r) k = (1 − µη)rK+k+1. For ClientOPT,
ClientOPT(η(s,r)l , x (s,r)) =
1
N N∑ i=1 K∑ k=1 ∇̃Fi(x(s,r)i,k−1).
where the ∇̃Fi(x(s,r)i,k−1)’s come from
x (s,r) i,k ← x (s,r) i,k−1 − η (s,r) l (∇̃Fi(x (s,r) i,k−1 + c− ci)) (13)
where ci = 1K ∑K i=1 ∇̃Fi(x(s,r)), c = 1 N ∑N i=1 ci.

heading: C TECHNICALITIES
text: In this section we collect all the assumptions we make in the paper for convenience.

heading: C.1 ASSUMPTIONS
text: Assumption 1. Fi’s are µ-strongly convex for µ > 0, i.e.
〈∇Fi(x), y − x〉 ≤ −(Fi(x)− Fi(y) + µ
2 ‖x− y‖2) (14)
for any i, x, y. Assumption 2. Fi’s are β-smooth where
‖∇Fi(x)−∇Fi(y)‖ ≤ β‖x− y‖ (15) for any i, x, y. Assumption 3. Fi’s have a uniform upper bound on the gradient dissimilarity (heterogeneity), defined as
ζ2 := sup x
1
N ∑ i∈[N ] ‖∇F (x)−∇Fi(x)‖2, (16)
Assumption 4. Fi’s have an upper bound on the gradient dissimilarity (heterogeneity) at the optimum, defined as
ζ2∗ := (1/N) ∑ i∈[N ] ‖∇Fi(x∗)‖2 (17)
Assumption 5. We assume that the Fi’s have bounded variance of stochastic gradients, where zi is a randomly drawn data point from client i’s data:
Ezi‖∇f(x; zi)−∇Fi(x)‖2 ≤ σ2. (18) Assumption 6. We assume that the initial suboptimality gap is upper bounded by ∆, where x0 is the initial feasible iterate and x∗ is the global minimizer.
F (x0)− F (x∗) ≤ ∆ (19)
Assumption 7. We assume that the initial distance from the optimum is upper bounded by D, where x0 is the initial feasible iterate and x∗ is the global minimizer.
‖x0 − x∗‖2 ≤ D2 (20)

heading: C.2 DEFINITIONS
text: Definition 1. An algorithm is said to have converged, or to have satisfied the convergence criterion if
EF (x̂)− F (x∗) ≤ (21) where is some user-chosen value. Definition 2. A rate of convergence of an algorithm in our setting is an upper bound on the error of the returned solution x̂ in terms of N , σ2, ζ, R, K, ∆, β, µ, ζ∗, e.g.
EF (x̂)− F (x∗) ≤M(R,K,N,∆, ζ, ζ∗, β, µ, σ2) (22) where M is a function that returns a one-dimensional value.

heading: D PROOF OF THM. 1
text: Theorem 4 (Formal version of Thm. 1). Run Algorithm 1 over S = 2 stages. In the first stage, use FedAvg4 as for R/2 rounds under the following parameter settings: η = η(1,r)g = η (1,r) l = O(min{ 1β , 1 µKR}), S = 1, γ (1) = 0, and w(1,r)k = (1− µη)rK+k+1. In the second stage, run SGD (so γ(2) = 0) for R/2 rounds under the following settings:
1. If R > 4βµ and r < d R 2 e: η (2,r) l = 0, η (2,r) g = ηg = 1 4β , w (2,r) k = 0
2. IfR > 4βµ and r ≥ dR/2e: η (2,r) l = 0,η (2,r) g = ηg = 2 µ(κ−R−dR2 e) ,w(2,r)k = (κ+r−d R 2 e) 2
3. If R ≤ 4βµ : η (2,r) g = 1 4β , η (2,r) l = 0, w (2,r) k = (1− µηg)−(rK+k+1).
and
K ≥ κ R log( µ2R2D2 ζ2 ) K ≥ µσ 2R Nζ2 (23)
K ≥ βσ 2
ζ2 (24)
Denote the output of this second stage by x(2). For large enough number of local rounds K, the convergence rate of this two-stage algorithm is
EF (x(2))− F (x∗) ≤ Õ(( β 2ζ2 µ3R2 ) exp(−R κ ) + σ2 µNKR ) (25)
Proof. First, let the total budget for the number of rounds be 2R. We run the two stages of the algorithm with R rounds.
In the first stage, we run Algorithm 1 with
η = η(s,r)g = η (s,r) l = min{
1
4β ,
log(max{2,min{E‖x (0)−x∗‖2Nµ2R2K2 σ2 , E‖x(0)−x∗‖2µ2R3K3
6βK(σ2+2Kζ2)) }} µKR
S = 1, γ(s) = 0, w(s,r)k = (1−µη)rK+k+1. Note this is just running FedAvg Gorbunov et al. (2020) with the above stepsizes. Therefore, from the proof of Gorbunov et al. (2020)(Eq.82), we know that we get the following convergence rate for the output of the algorithm
EF (x(1))− F (x∗) ≤ 8β exp(−KR 8κ )E‖x(0) − x∗‖2 (26)
+ 2(Φ + 1)σ2
µNKR +
6β(Φ2 + 1)σ2
µ2KR2 +
12β(Φ2 + 1)ζ2
µ2R2 (27)
where Φ = log(max{2,min{E‖x (0)−x∗‖2Nµ2R2K2 σ2 , E‖x(0)−x∗‖2µ2R3K3 6βK(σ2+2Kζ2)) }}.
In the second stage, we run Algorithm 1 on x(1) with β(s) = 0 and
1. If R > 4βµ and r < d R 2 e: η (2,r) l = 0, η (2,r) g = ηg = 1 4β , w (2,r) k = 0
2. If R > 4βµ and r ≥ dR/2e: η (2,r) l = 0, η (2,r) g = ηg = 2 µ(κ−R−dR2 e) ,
w (2,r) k = (κ+ r − d R 2 e) 2
3. If R ≤ 4βµ : η (2,r) l = 0, η (2,r) g = 1 4β , w (2,r) k = (1− µηg)−(r+1).
4This is the FedAvg detailed in (Gorbunov et al., 2020), which differs from usual FedAvg in that it returns a weighted average of the iterates rather than the final iterate. This is mostly a proof convenience, and in practice one would return the final iterate as we do in the experiments.
Note that this is running SGD with the above stepsize schedule. From the proof from (Woodworth et al., 2020a), the output of this algorithm is
EF (x(2))− F (x∗) ≤ 128κ(EF (x(1))− F (x∗)) exp(− R 8κ ) + 72σ2 µNKR (28)
Now we set the terms that shrink with K in Eq. (26))large enough to be of the same order as the last term in Eq. (26) or to be of order in the final convergence rate. We start by setting
βD2 exp(−KR κ ) ≤ βζ 2 µ2R2 (29)
=⇒ exp(−KR κ ) ≤ ζ 2 µ2R2D2 (30)
=⇒ K ≥ κ R log( µ2R2D2 ζ2 ) (31)
Next we set σ2
µNKR ≤ ζ
2
µ2R2 (32)
=⇒ K ≥ µσ 2R
Nζ2 (33)
Finally we set
βσ2
µ2KR2 ≤ ζ
2
µ2R2 (34)
=⇒ K ≥ βσ 2
ζ2 (35)
So we need that
K ≥ κ R log( µ2R2D2 ζ2 ) K ≥ µσ 2R Nζ2 (36)
K ≥ βσ 2
ζ2 (37)
This implies that
EF (x(1))− F (x∗) ≤ Õ( βζ 2
µ2R2 ) (38)
Plugging this into Eq. (28), we get
EF (x(2))− F (x∗) ≤ Õ( β 2ζ2 µ3R2 exp(−R κ ) + σ2 µNKR ) (39)

heading: E PROOF OF THM. 2
text: Theorem 5 (Formal version of Thm. 2). Run Algorithm 1 over S = 2 stages. In the first stage, use FedAvg for for R/2 rounds under the parameter settings of Stage 1 from Thm. 4.
In the second stage, run ASG for R/2 rounds with the following parameter settings:
η (2,r) l = 0, η (2,r) g = ηg = O(min{ 1Kβ , 1 µKR2 }) γ
(2) = 1− √ µηgK
1+ √ µηgK
w (1,r) k = w (2,r) k = 0 for all k, r except for w (2,R/2) 1 = 1 (i.e. return the last iterate). Also let
K ≥ κ R log( µ2R2D2 ζ2 ) K ≥ µσ 2R Nζ2 (40)
K ≥ βσ 2
ζ2 (41)
Denote the output of this second stage by x(2). For large enough number of local rounds K, the expected convergence rate of this two-stage algorithm is
EF (x(2))− F (x∗) ≤ Õ( βζ 2 µ2R2 exp(− R√ κ ) + σ2 µNKR +
βσ2
µ2NKR3 ). (42)
Proof. First, let the total budget for the number of rounds be 2R. We run the two stages of the algorithm with R rounds.
In the first stage, we run Algorithm 1 with
η = η(s,r)g = η (s,r) l = min{
1
4β ,
log(max{2,min{E‖x (0)−x∗‖2Nµ2R2K2 σ2 , E‖x(0)−x∗‖2µ2R3K3
6βK(σ2+2Kζ2)) }} µKR
S = 1, γ(s) = 0, w(s,r)k = (1 − µη)rK+k+1. Note this is just running FedAvg (Gorbunov et al., 2020) with the above stepsizes. Therefore, from the proof of (Gorbunov et al., 2020), we know that we get the following convergence rate for the output of the algorithm
EF (x(1))− F (x∗) ≤ 8β exp(−KR 8κ )E‖x(0) − x∗‖2 (43)
+ 2(Φ + 1)σ2
µNKR +
6β(Φ2 + 1)σ2
µ2KR2 +
12β(Φ2 + 1)ζ2
µ2R2 (44)
where Φ = log(max{2,min{E‖x (0)−x∗‖2Nµ2R2K2 σ2 , E‖x(0)−x∗‖2µ2R3K3 6βK(σ2+2Kζ2)) }}.
Next we run Algorithm 1 on x(1) with η(s,r)l = 0 and
η(s,r)g = ηg = min{ 1
Kβ ,
log(max{2,min{ 2∆1NµKRσ2 , 2∆1µ
2NKR3
βσ2 })} µKR2 } (45)
γ(s) = 1− √ µηgK
1+ √ µηgK , S = 1, and w(s,r)k = 0 for all k, r, s except for w (1,R1) 1 = 1 and ∆1 ≤
8β exp(−κR8 )E‖x (0) − x∗‖2 + 2(Φ+1)σ
2 µNKR + 6β(Φ2+1)σ2 µ2KR2 + 12β(Φ2+1)ζ2 µ2R2 . From Thm. 6 we know that
EF (x(2))− F (x∗) ≤ 2∆1 exp(− R√ κ ) + σ2(Φ 1/2 2 + 1) µNKR + βσ2(Φ 3/2 2 + 1) µ2NKR3 (46)
where Φ2 = log(max{2,min{ 2∆1µKσ2 , 2∆µ2K3
βσ2 })}. Now we set the terms that shrink with K in Eq. (43))large enough to be of the same order as the last term in Eq. (43). We start by setting
βD2 exp(−KR κ ) ≤ βζ 2 µ2R2 (47)
=⇒ exp(−KR κ ) ≤ ζ 2 µ2R2D2 (48)
=⇒ K ≥ κ R log( µ2R2D2 ζ2 ) (49)
Next we set σ2
µNKR ≤ ζ
2
µ2R2 (50)
=⇒ K ≥ µσ 2R
Nζ2 (51)
Finally we set
βσ2
µ2KR2 ≤ ζ
2
µ2R2 (52)
=⇒ K ≥ βσ 2
ζ2 (53)
So we need that
K ≥ κ R log( µ2R2D2 ζ2 ) K ≥ µσ 2R Nζ2 (54)
K ≥ βσ 2
ζ2 (55)
This implies that
EF (x(1))− F (x∗) ≤ Õ( βζ 2
µ2R2 ) (56)
Plugging this into Eq. (46), we get
EF (x(2))− F (x∗) ≤ Õ( βζ 2 µ2R2 exp(− R√ κ ) + σ2 µNKR +
βσ2
µ2NKR3 ) (57)

heading: F THM. 6
text: Theorem 6. Consider Nesterov’s Accelerated Stochastic Gradient (ASG) method with the iteration sequence
xk+1/2 = (1 + γ)xk − γxk−1, xk+1 = xk+1/2 − α∇̃F (xk+1/2) (58)
with γ = 1− √ αµ
1+ √ αµ and
α = min{ 1 β ,
log(max{2,min{ 2∆NµKσ2 , 2∆µ2NK3
βσ2 })} µK2 } , then (59)
EF (xK)− F (x∗) ≤ 2∆ exp(− K√ κ ) + σ2(Φ1/2 + 1) µK + βσ2(Φ3/2 + 1) µ2K3 (60)
Proof. From (Aybat et al., 2019) (Eq.13), we know that
EVk+1 ≤ (1− √ αµ)(EVk) + σ2α
2 (1 + αβ) (61)
where F (xk)− F (x∗) ≤ Vk ≤ 2(F (xk)− F (x∗)). If we unroll this, we get
EVK ≤ (1− √ αµ)KEV0 + σ2α1/2 2µ1/2 + σ2α3/2β 2µ1/2 ≤ exp(−√αµK)EV0 + σ2α1/2 2µ1/2 + σ2α3/2β 2µ1/2
(62)
Now if
α = min{ 1 β ,
log(max{2,min{ 2∆NµKσ2 , 2∆µ2NK3
βσ2 })} µK2 } (63)
observe that if 1β ≤ log(max{2,min{ 2∆NµK σ2 , 2∆µ
2NK3
βσ2 })}
µK2 we know that letting Φ =
log(max{2,min{ 2∆NµKσ2 , 2∆µ2NK3 βσ2 })}
EVK ≤ EV0 exp(− K√ κ ) + σ2Φ1/2 2µK + βσ2Φ3/2 2µ2K3 (64)
and if 1β > log(max{2,min{ 2∆µK σ2 , 2∆µ
2K3
βσ2 })}
µK2
EVK ≤ σ2
µK +
βσ2 µ2K3 + σ2Φ1/2 2µK + βσ2Φ3/2 2µ2K3 (65)
where we used that EV0 ≤ 2∆. So altogether, we have that
EF (xK)− F (x∗) ≤ 2∆ exp(− K√ κ ) + σ2(Φ1/2 + 1) µK + βσ2(Φ3/2 + 1) µ2K3 (66)

heading: G PROOF OF THM. 3
text: Theorem 7 (Formal version of Thm. 3 in Appendix). Let Algorithm 1 be run over two superstages for a total of R rounds over S stages, with R ≥ 4 √ κ. In the first superstage, run Multistage FedAvg (M-FedAvg) (see App. H for settings, pseudocode in Algorithm 7) for R/2 rounds. In the second superstage, run Multistage Accelerated Stochastic Gradient (M-ASG) (Aybat et al., 2019) (pseudocode in Algorithm 5), using the following parameters (for legibility, we reset s = 1 for the second superstage): η
(s,r) l = 0, η (1,r) g = 1 Kβ with R
(1) ≥ 1 R(1) = d(p+ 1) √ κ log(12(p+ 1)κ)e
η (s,r) g = 1 K22sβ with R
(s) = 2sd √ κ log(2p+2)e for any p ≥ 1 and s ≥ 2 γ(s) = 1−
√ µη (s,r) g K
1+ √ µη (s,r) g K
For B large enough, this achieves expected convergence rate of
EF (x(S))− F (x∗) ≤ O((κ∆ exp(−R) + β 2ζ2
µ3R2 +
βσ2
µ2BR2 ) exp(− R√ κ ) +
σ2
µBNR
Proof. Let x(s ′) be the output of multistage FedAvg from the first half of the rounds (s′ is the index of the final stage). Then we know that the output of this has guarantee (from Thm. 8)
EF (x(s ′))− F (x∗) ≤ O(κ∆ exp(−KR
κ ) +
βσ2
µ2BNKR +
β2σ2 µ3BKR2 + β2ζ2 µ3R2 ) (67)
Where in this convergence rate K is limited to be at most κ. Next, observe that the second half of the rounds are running M-ASG from (Aybat et al., 2019). That has a guarantee of (Corollary 3.8), (with no limit on K)
EF (x(S))− F (x∗) ≤ O(∆1 exp(− R√ κ ) + σ2 µBNKR ) (68)
where
∆1 ≤ O(κ∆ exp(− KR
κ ) +
βσ2
µ2BNKR +
β2σ2 µ3BKR2 + β2ζ2 µ3R2 ) (69)
If we set K = κ as per the limit, then we have
∆1 ≤ Õ(κ∆ exp(−R) + σ2
µBNR +
βσ2 µ2BR2 + β2ζ2 µ3R2 ) (70)
Plugging ∆1 in and K = κ, we get
EF (x(S))− F (x∗) ≤ O((κ∆ exp(−R) + σ 2
µBNR +
βσ2 µ2BR2 + β2ζ2 µ3R2 ) exp(− R√ κ ) + σ2 βBNR )
Which can be simplified to (noting that µ ≤ β)
EF (x(S))− F (x∗) ≤ O((κ∆ exp(−R) + β 2ζ2 µ3R2 ) exp(− R√ κ ) + σ2 µBNR + βσ2 µ2BR2 (71)

heading: H MULTISTAGE FEDAVG (M-FEDAVG)
text: Theorem 8. Instantiate Algorithm 1 as
η(1,r)g = η (1,r) l =
1
4β ,R(1)K ≥ 1
η(1,r)g = η (1,r) l =
1
β2s+2 , R(s)K = dp2s+2κ log(2)e
where p ≥ 3 is some arbitrary positive number. Also set γ(s) = 0. Let ClientOPT be one round of either SS-Local-SGD or FedAvg, except each local step is minibatched with a minibatch of size B, and have StageWeightedAverage simply return the last iterate. Furthermore, require that K ≤ dpκ log(2)e, and dpκ log(2)e mod K = 0. Then if R(1) = R/C for some constant and KR ≥ 2κ,
EF (x(S))− F (x∗) ≤ O(κ∆ exp(−KR κ ) + βσ2 µ2BNKR + β2σ2 µ3BKR2 + β2ζ2 µ3R2 ) (72)
Proof. We start by proving that
errs ≤ exp(−(R1K)/(4κ))
2p(s−1) err0 +
σ2
2sµβN +
6
22sµβ (Kσ2 + 2K2ζ2) (73)
Where errs = E‖x(s,Rs) − x∗‖2. We prove this by induction on s. Recall that from Thm. 9 we have
E‖x(s,r) − x∗‖2 ≤ (1− µη)rKE‖x(s,0) − x∗‖2 + ησ 2
µN + 6κη2(Kσ2 + 2K2ζ2) (74)
Starting with the base case, with η = 14β , R1K ≥ 1, we have
err1 ≤ exp(− R1K
4κ )err0 +
σ2
4βµN +
3
2βµ (Kσ2 + 2K2ζ2) (75)
Now let us assume the result holds for s and prove it for s+ 1. Using Thm. 9 for stage s+ 1 with η = 1β2s+3 and Rs+1K = dp2 s+3κ log(2)e gets us
errs+1 ≤ exp(−p log(2))errs + σ2
2s+3µβN +
6
22(s+3)µβ (Kσ2 + 2K2ζ2) (76)
Now using the induction hypothesis, we know
errs ≤ exp(−(R(1)K)/(4κ))
2p(s−1) err0 +
σ2
2sµβN +
6
22sµβ (Kσ2 + 2K2ζ2) (77)
And substituting it back into (Eq. (73)) while having p ≥ 3 gets us
errs+1 ≤ exp(−(R(1)K)/(4κ))
2ps err0 +
σ2
2s+1µβN +
6
22s+2µβ (Kσ2 + 2K2ζ2) (78)
Using the following lemma, we can translate this into a convergence rate:
Lemma 1. Let R(1)K ≥ 1 and R(s)K = dp2s+2κ log(2)e. Then we know that 2k ≥ Θ(1)(RK−R
(1)K pκ ).
Proof. Comes from a simple unrolling of the stage lengths; similar to the proof in Fallah et al. (2020).
So if we plug in R(1) = RC for some constant C ≥ 2 and KR ≥ 2κ, we get the convergence rate in the theorem statement by applying smoothness to transform the bound on distance to optimum to function value.

heading: I PROOF OF THM. 9
text: Theorem 9. Suppose we ran Algorithm 1 for one stage with stepsizes η(s,r)l = η (s,r) g = η, with γ(s) = 0 and return the last iterate. Furthermore, require that η ≤ 1β . Then
E‖x(s,r) − x∗‖2 ≤ (1− µη)rKE‖x(s,0) − x∗‖2 + ησ 2
µN + 6κ(Kη2σ2 + 2K2η2ζ2) (79)
From (Gorbunov et al., 2020) (Eq.82), we have the following for FedAvg:
E‖xt+1 − x∗‖2 ≤ (1− µη)E‖xt − x∗‖2 + 2βηEVt + η2 σ2
N (80)
where Vt = 1N ∑N i=1 E‖xti − xt‖2, where t = Kr + k where r is current round, k is step on current round. xti is client i’s iterate at the t-th step. x t = 1N ∑N i=1 x t i, and η is both the inner stepsize and outer stepsize (i.e. η(r,s)l = η, η (r,s) g = η).
From (Woodworth et al., 2020a) we know that for FedAvg
Vt ≤ 3Kη2σ2 + 6K2η2ζ2 (81)
So together,
E‖xt+1 − x∗‖2 ≤ (1− µη)E‖xt − x∗‖2 + 6βη(Kη2σ2 + 2K2η2ζ2) + σ 2η2
N (82)
And unrolling this,
E‖xT − x∗‖2 ≤ (1− µη)TE‖x0 − x∗‖2 + ησ 2
µN + 6κ(Kη2σ2 + 2K2η2ζ2) (83)

heading: J USING SS-LOCAL-SGD INSTEAD OF FEDAVG
text: In this section we show how one can use SS-Local-SGD instead of FedAvg as ClientOPT. Theorem 10. Let Vt = 1N ∑N i=1 E‖xti − xt‖2, where t = Kr+ k where r is current round, k is step on current round, xti is client i’s iterate at the t-th step. Then SS-Local-SGD has
Vt ≤ 48K2η2ζ2 + 9η2σ2 (84)
Proof.
Vt ≤ 1
N2 N∑ i=1 N∑ j=1 E‖xti − xtj‖2 (85)
By Jensen’s. Let gi(xti) be the SS-Local-SGD local update for client i at iteration t. Let Et denote expectation conditioned on everything up to the t− 1’th iteration. Let yt be the last synchronized iterate before iteration t. In the following we use the fact that the control iterates are computed with
minibatches of size K at the clients.
E‖xti − xtj‖2
≤ E‖xt−1i − x t−1 j − ηEt−1gi(x t−1 i ) + ηEt−1gj(x t−1 j )‖ 2 + η2σ2
= E‖xt−1i − x t−1 j − η(∇̃F (y t−1)− ∇̃Fi(yt−1)−∇Fi(xt−1i ))
+ η(∇̃F (yt−1)− ∇̃Fj(yt−1) +∇Fj(xt−1j ))‖ 2 + η2σ2
= E‖xt−1i − x t−1 j ± η∇F (x t−1 i )± η∇F (x t−1 j )± η∇F (y t−1)
− η(−∇̃Fi(yt−1)−∇Fi(xt−1i )) + η(−∇̃Fj(y t−1) +∇Fj(xt−1j ))‖ 2 + η2σ2
≤ (1 + 1 K − 1 )E‖xt−1i − x t−1 j − η∇F (x t−1 i ) + η∇F (x t−1 j )‖ 2
+ η2KE‖∇F (xt−1i )−∇Fi(x t−1 j )−∇F (x t−1 j ) +∇Fj(x t−1 j ) +∇F (yt−1)−∇Fi(yt−1)−∇F (yt−1) +∇Fj(yt−1)‖2
+ 3η2σ2
≤ (1 + 1 K − 1 )E‖xt−1i − x t−1 j ‖ 2 + 16Kη2ζ2 + 3η2σ2
So altogether,
1
N2 N∑ i=1 N∑ j=1 E‖xti − xtj‖2 ≤ 1 N2 N∑ i=1 N∑ j=1 (1 + 1 K − 1 )E‖xt−1i − x t−1 j ‖ 2 + 16Kη2ζ2 + 3η2σ2
(86)
Unrolling this until the last time the iterates were synchronized, we get
1
N2 N∑ i=1 N∑ j=1 E‖xti − xtj‖2 ≤ 3K(16Kη2ζ2 + 3η2σ2) (87)
Which implies
Vt ≤ 48K2η2ζ2 + 9η2σ2 (88)
This can be substituted into Thm. 9 to get the multistage rate for Thm. 7, or it can be substituted into Theorem G.3 in (Gorbunov et al., 2020) to get a convergence rate that differs from (Eq. (26)) by only constants in the terms. Thus we can use SS-Local-SGD instead of Local-SGD in any of the theorems presented in the main paper with only minor constant changes.

heading: K EXPERIMENTAL SETUP DETAILS

heading: K.1 CONVEX OPTIMIZATION
text: We empirically evaluate multistaging on federated regularized logistic regression. Let (xi,j , yi,j) be the jth datapoint of the ith client. We minimize Eq. (1) where
Fi(w) = 1
ni ( ni∑ j=1 −yi,j log(w>xi,j)− (1− yi,j) log(1− w>xi,j))) + µ 2 ‖w‖2.
Dataset. We use the MNIST dataset of handwritten digits (LeCun et al., 2010). We model a federated setting with five clients by partitioning the data into groups. First, we take 500 images from each digit class (total 5,000). Each client’s local data is a mixture of data drawn from two digit classes (leading to heterogeneity), and data sampled uniformly from all classes. We call a federated dataset X% homogeneous if the first X% of each class’s 500 images is shuffled and evenly partitioned to
each client. The remaining (100−X)% is partitioned as follows: client i ∈ {1, . . . , 5} receives the remaining non-shuffled data from classes 2i− 2 and 2i− 1. For example, in a 50% homogeneous setup, client 3 has 250 samples from digit 4, 250 samples from digit 5, and 500 samples drawn uniformly from all classes. Note that 100% homogeneity is not the same thing as setting heterogeneity ζ = 0 due to sampling randomness; we use this technique for lack of a better control over ζ. To model binary classification, we let even classes represent 0’s and odd classes represent 1’s, and set K = 20.
Hyperparameters. All experiments are initialized at 0 with regularization µ = 0.1. We fix the total number of rounds R (differs across experiments). For algorithms using Nesterov momentum, we calculate the momentum parameter from strong convexity µ and the chosen stepsizes. For the two-stage methods, we tune the stage 1 stepsize η = η(1,r)l = η (1,r) g in the range below:
{10−3, 10−2.5, 10−2, 10−1.5, 10−1} (89)
For the second stage we let the stepsizes be η(2,r)l = 0, η (2,r) g = η/K. We also tune the percentage of rounds before switching to the second stage in the range as
{10−2, 10−1.625, 10−1.25, 10−0.875, 10−0.5} (90)
For M-(FedAvg→SGD), M-(FedAvg→ASG), M-(SCAFFOLD→SGD), M-(SCAFFOLD→ASG), the number of total stages is problem-dependent. In superstage 1, we first tune the Stage 1 stepsize η (as well as the fraction of total rounds allocated to Stage 1) in the range (89). Let R(1) be the number of rounds in the first stage. For subsequent stages s ≥ 2, we use step sizes η(s,r)l = η (s,r) g = η 2s−1 and R(s) = 2s−1R(1). Once the stepsize decreases to η/K, we enter the second superstage. Since we evaluate only minibatch algorithms in superstage 2, we use the same parameters and decay sequence as in superstage 1, except the starting stepsize is η(1,r)g = ηK while η (s,r) l = 0. The SCAFFOLD (SS-Local-SGD) used in this section is Algorithm 9.
All algorithms are tuned for the lowest average gradient norm evaluated over the whole dataset on 1000 trials.

heading: K.2 NONCONVEX OPTIMIZATION

heading: K.2.1 EMNIST
text: In this section we detail how we use the EMNIST (Cohen et al., 2017) dataset for our nonconvex experiments, where the handwritten characters are partitioned by their author. In this dataset, there are 3400 clients, with 671,585 images in the train set and 77,483 in the test set. The task is to train a convolutional network with two convolutional layers, max-pooling, dropout, and two dense layers as in the Federated Learning task suite from (Reddi et al., 2020). The SCAFFOLD used in this section is Algorithm 8.
Hyperparameters. We set R = 500, and rather than fixing the number of local steps, we fix the number of local client epochs (the number of times a client performs gradient descent over its own dataset) to 20. The number of clients sampled per round is 10. These changes were made in light of the imbalanced dataset sizes across clients, and is standard for this dataset-task (Reddi et al., 2020; Charles & Konečnỳ, 2020; Charles et al., 2021). For all algorithms aside from Multistage Minibatch SGD (M-SGD), we tune the stage 1 stepsize η = η(1,r)l = η (1,r) g in the range {0.05, 0.1, 0.15, 0.2, 0.25}. The second stage of the two-stage algorithms we let the stepsizes be η(2,r)l = 0, η (2,r) g = η/K̄ where K̄ is the average number of local steps performed across clients. For two-stage algorithms, we tune the percentage of rounds before switching to the second stage in the range of {0.1, 0.3, 0.5, 0.7, 0.9}. We next detail the way the multistage algorithms (S > 2) are run. In superstage 1, we first tune the Stage 1 stepsize η {0.1, 0.2, 0.3, 0.4, 0.5}. Let R(1) be the number of rounds in the first stage. We tune R(1) ∈ d{0.1R, 0.275R, 0.45R}e. For subsequent stages s ≥ 2, we use step sizes η
(s,r) l = η (s,r) g = η 2s−1 and R (s) = 2s−1R(1). After R′ rounds have passed, we enter the second superstage. We tune R′ ∈ d{0.5R, 0.7R, 0.9R}e. Let the last stage of superstage 1 be s′. In the subsequent stages, we let R(s) = 2s−s
′−1R(1) and η(s,r)g = ηK̄2s−s′−1 and η (s,r) l = 0. Since we
evaluate only minibatch algorithms in superstage 2, we use the same parameters as superstage 1, except the starting stepsize is η(1,r)g = ηK while η (s,r) l = 0. For multistage SGD, we tune
η ∈ {0.1, 0.2, 0.3, 0.4, 0.5} and R(1) ∈ d{10−2R, 10−1.625R, 10−1.25R, 10−0.875R, 10−0.5R}e We then set η(s,r)g = ηK̄2s−1 , η (s,r) l = 0, and R
s = R(1)2s−1. For multistage SCAFFOLD and FedAvg, we tune η ∈ {0.1, 0.2, 0.3, 0.4, 0.5} and R(1) ∈ d{0.1R, 0.275R, 0.45R}e. We then set η (s,r) g =
η K̄2s−1 , η(s,r)l = 0, and R s = R(1)2s−1. When evaluating a particular metric (training loss
or test accuracy) the algorithms are tuned according to the mean of the last 5 evaluated iterates of that particular metric. Reported accuracies are also the mean of the last 5 evaluated iterates.
K.2.2 CIFAR-100
As stated before, for our nonconvex experiments, we use the CIFAR-100 (Krizhevsky, 2009) dataset, where the handwritten characters are partitioned via the method proposed in (Reddi et al., 2020). In this dataset, there are 500 train clients, with a total of 50,000 images in the train set. There are 100 test clients and 10,000 images in the test set. The task is to train a ResNet-18, replacing batch normalization layers with group normalization as in the Federated Learning task suite from (Reddi et al., 2020). We leave out SCAFFOLD in all experiments because of out of memory issues.
Hyperparameters. We set R = 500, and rather than fixing the number of local steps, we fix the number of local client epochs (the number of times a client performs gradient descent over its own dataset) to 20. The number of clients sampled per round is 10. These changes were made in light of the imbalanced dataset sizes across clients, and is standard for this dataset-task (Reddi et al., 2020; Charles & Konečnỳ, 2020; Charles et al., 2021). For all algorithms aside from Multistage Minibatch SGD (M-SGD), we tune the stage 1 stepsize η = η(1,r)l = η (1,r) g in the range {0.05, 0.1, 0.15, 0.2, 0.25}. The second stage of the two-stage algorithms we let the stepsizes be η(2,r)l = 0, η (2,r) g = η/K̄ where K̄ is the average number of local steps performed across clients. For two-stage algorithms, we tune the percentage of rounds before switching to the second stage in the range of {0.1, 0.3, 0.5, 0.7, 0.9}. When evaluating a particular metric (training loss or test accuracy) the algorithms are tuned according to the mean of the last 5 evaluated iterates of that particular metric. Reported accuracies are also the mean of the last 5 evaluated iterates.

heading: K.3 NONCONVEX OPTIMIZATION
text: Table 2 describes the final test accuracy of the learned models for each optimization technique. Table 2 includes the test accuracy for our experiments on EMNIST classification, with the experimental tuning details in App. K.2.1. We observe that SCAFFOLD→ SGD has the best test accuracy.

heading: L DISCUSSION ON PARTIAL CLIENT PARTICIPATION
text: Algorithm 10 (Simplified) Partial Participation FedAvg→SGD Input: Number of local steps K, rounds R, stepsize η(1), stepsize sequence {η(2)r }1≤r≤R−1, initial point x(0), maximum number of clients to sample per round S Server chooses a client j ∈ [N ] uniformly randomly Server communicates x(0) to client j Client j sets x(0)j,0 = 0 for k = 0, . . . ,K − 1 do x
(0) j,k+1 ← x (0) j,k − η(1) · ∇Fi(x (0) j,k)
end for Client j sets x̄j,K = 1WK ∑K k=1 wkx (0) j,k where wk = (1− η(1)µ 4 ) −k and WK = ∑K k=1 wk Client j communicates x̄j,K to server Server sets x(1) ← x̄j,K for r = 1, . . . , R− 1 do
Sample clients S ⊆ [N ] of size S x(r+1) ← x(r) − η(2)r · 1SK ∑ i∈S,1≤k≤K ∇Fi(x(r))
end for
In this section we discuss the rates our framework can achieve if only a subsample of clients is selected. For simplicity’s sake, we will assume that σ2 = 0, e.g. there is no stochasticity in the client gradient computations. The algorithm we will study is as follows.
1. We run SGD on one client for K local steps in the first round. 2. We take the solution from the end of step (1) and run SGD from the server. Precisely, each
round we use the average of K gradients taken from each of the S uniformly sampled clients as the unbiased gradient for SGD.
We write out the algorithm explicitly in Algorithm 10.
We choose to run the FedAvg component of the algorithm with one client sampled and for only one round because analyzing this simple case already conveys the nature of the gain we receive from our framework in the partial participation setting.
We start by analyzing an upper bound on suboptimality after step (1). For our proof, we need the following definitions:
ζ̂2 := sup i,x ‖∇F (x)−∇Fi(x)‖2, ζ̂2∗ := sup i ‖∇Fi(x∗)‖2, (91)
Theorem 11. Let x̄j,K = 1WK ∑K k=1 wkx (0) j,k where wk = (1 − η(1)µ 4 ) −k and WK = ∑K k=1 wk, where x(0)j,k are generated in the following way:
x (0) j,k ← x (0) j,k−1 − η1 · ∇Fj(x (0) j,k−1) (92)
where x(0)j,0 = x (0). For any j ∈ [N ], x̄j,K satisfies
F (x̄ (0) j,K)− F (x
∗) ≤ 4β 2 µ ‖x(0) − x∗‖2 exp(− K 4κ2 ) + 2ζ̂2 µ (93)
as long as η(1) ≤ µ8β2 .
Proof. Shorthand x(0)j,k as xk and η (1) as η1 for the duration of this proof. Then
‖xk − x∗‖2 = ‖xk−1 − η1∇Fj(xk)− x∗‖2 (94) = ‖xk−1 − x∗‖2 − 2η1〈∇Fj(xk), xk − x∗〉+ η21‖∇Fj(xk)‖2 (95)
We start by bounding the middle term of Eq. (94).
−2η1〈∇Fj(xk), xk − x∗〉 = −2η1(〈∇F (xk), xk − x∗〉+ 〈∇Fj(xk)− F (xk), xk − x∗〉) (96) The first term of Eq. (96) can be upper bounded using strong convexity Assumption 1:
−〈∇F (xk), xk − x∗〉 ≤ −(F (xk)− F (x∗) + µ
2 ‖xk − x∗‖2) (97)
The second term can be upper bounded using the Fenchel-Young inequality:
−〈∇Fj(xk)− F (xk), xk − x∗〉 ≤ 1
µ ‖Fj(xk)− F (xk)‖2 +
µ 4 ‖xk − x∗‖2 (98)
≤ ζ̂ 2 µ + µ 4 ‖xk − x∗‖2 (99)
Where in the last inequality we used the definition of ζ̂2.
Next, we upper bound the third term of Eq. (94).
η21‖∇Fj(xk)‖2 = 2η21‖∇Fj(xk)− Fj(x∗)‖2 + 2η21‖∇Fj(x∗)‖2 (100) ≤ 2η21β2‖xk − x∗‖2 + 2η21 ζ̂2∗ (101)
where in the last step we used smoothness Assumption 2 for the first term and the definition of ζ̂2∗ for the second term.
We can substitute all our upper bounds into Eq. (94) to get
‖xk − x∗‖2 (102)
≤ (1− η1µ+ η1µ
2 + 2η21β
2)‖xk−1 − x∗‖2 − 2η1(F (xk)− F (x∗)) + 2η1ζ̂
2
µ + 2η21 ζ̂ 2 ∗ (103)
Setting η1 ≤ µ8β2 , we know that 2η 2 1β 2 ≤ η1µ4 and that 2η 2 1 ζ̂ 2 ∗ ≤
η1µζ̂ 2 ∗
4β2 ≤ η1ζ̂
2 ∗
4µ ≤ η1ζ̂
2
4µ . So we have
‖xk − x∗‖2 (104)
≤ (1− η1µ 4 )‖xk−1 − x∗‖2 − 2η1(F (xk)− F (x∗)) + 3η1ζ̂
2
µ (105)
And rearranging,
F (xk)− F (x∗) ≤ (1− η1µ4 )‖xk−1 − x ∗‖2 − ‖xk − x∗‖2
2η1 +
2ζ̂2
µ (106)
Using convexity of F , and letting x̄K := x̄j,K , we know if wk = (1− η1µ4 ) −k, then
F (x̄K)− F (x∗) ≤ 1
WK K∑ k=1 wk(F (xk)− F (x∗)) (107)
≤ 1 2η1WK K∑ k=1 wk[(1− η1µ 4 )‖xk−1 − x∗‖2 − ‖xk − x∗‖2] + 2ζ̂2 µ (108)
≤ 1 2η1WK (‖x0 − x∗‖2) + 2ζ̂2 µ (109)
Noting again that η1 ≤ µ8β2 and that 1 WK ≤ 1wK ≤ (1− η1µ 4 ) K , we have
F (x̄K)− F (x∗) ≤ 4β2
µ ‖x0 − x∗‖2(1−
η1µ
4 )K +
2ζ̂2
µ (110)
≤ 4β 2
µ ‖x0 − x∗‖2 exp(−
K
4κ2 ) +
2ζ̂2
µ (111)
Remark 1. Surprisingly, Thm. 11 is the only known rate (to our knowledge) for FedAvg (even running for one round, only sampling one client, and assuming noiseless gradients) under partial participation in the heterogeneous client setting that is tight in the following sense: given large enough K, only one communication round is needed to get convergence to the global solution if ζ2 = 0 (which implies ζ̂2 = 0, e.g. completely homogeneous clients). Other rates for FedAvg in partial participation, such as those in (Karimireddy et al., 2020b; Stich & Karimireddy, 2019) do not have this property. The work on federated optimization in the partial participation setting is fairly immature, and is an important future direction.
Remark 2. One might expect the ζ̂ 2 µ term to shrink with S, e.g. be something like ζ̂2
µS , as this is true for partial-participation SGD Woodworth et al. (2020a). However, the evidence points towards this not being true for partial participation FedAvg, because in the rate for full participation FedAvg Eq. (3), the heterogeneity term does not shrink with N .
Using this, we get a convergence rate for FedAvg→SGD (Algorithm 10) in the partial participation setting:
Theorem 12. Let x(1) = x̄j,K , as specified from Thm. 11. Let the total round budget be R + 1. Then if we have the following (i.e. running partial participation SGD with x(1) as the initial point) (Woodworth et al., 2020a))
1. If R > 4βµ and r − 1 < d R 2 e: η (2) r = ηg = 1 4β , wr = 0
2. If R > 4βµ and r − 1 ≥ dR/2e: η (2) r = 2 µ(κ−R−dR2 e) , wr = (κ+ r − dR2 e) 2
3. If R ≤ 4βµ : η (2) r = 1 4β , wr = (1− µηg) −(r+1).
Where Algorithm 10 returns
x̂ = R+1∑ r=2 wrx (r) (112)
Then we have the following convegence rate for Algorithm 10 if K ≥ κ2 log β 2D2
ζ̂2 :
E[F (x̂)− F (x∗)] ≤ O( ζ̂ 2 µ exp(− R√ κ ) + (1− S N ) ζ2 µSR ) (113)
Proof. The rate for running partial participation SGD with x(1) as the initial point which samples S clients per round from (Woodworth et al., 2020a) (eq. 18) is
E[F (x(r))− F (x∗)] ≤ O(κ[EF (x(1))− F (x∗)] exp(−R κ ) + (1− S N ) ζ2∗ µSR ) (114)
We can substitute for EF (x(1))− F (x∗) the final suboptimality from Thm. 11, which gives
E[F (x(r))− F (x∗)] ≤ O(κ(β 2 µ ‖x(0) − x∗‖2 exp(−K κ2 ) + ζ̂2 µ ) exp(−R κ ) + (1− S N ) ζ2∗ µSR )
If we set K ≥ κ2 log β 2D2 ζ̂2 where D2 ≤ ‖x(0) − x∗‖2, then the convergence rate becomes
E[F (x(r))− F (x∗)] ≤ O(κζ̂ 2 µ exp(−R κ ) + (1− S N ) ζ2∗ µSR ) (115)
It is easy to check that if we had used Nesterov acceleration in the second phase (like in (Woodworth et al., 2020a)), the rate would instead be
E[F (x(r))− F (x∗)] ≤ O( ζ̂ 2 µ exp(− R√ κ ) + (1− S N ) ζ2 µSR ) (116)
For both partial SGD and partial FedAvg→ SGD (Eq. (114) vs Eq. (115)), we see that the rate in Eq. (116) is better if ζ̂µ ≤ ∆. Similar is true with acceleration.
On the other hand, our new rates Eq. (115) and Eq. (116) improves on the rate for FedAvg Thm. 11 because our errors shrink with S and with R, supposing all the measures of heterogeneity ζ, ζ̂, ζ∗, ζ̂∗ are of the same order. However, the magnitude of improvement is not as substantial as in the full participation case: here we only improve by a factor of R, while in the full participation case we improved exponentially.
It is unclear if our algorithm/analysis is suboptimal here, because there is no known lower bound for partial participation FL in the heterogeneous setting (unlike for full participation). Understanding the limits of optimization in FL when clients are sampled and client data is heterogeneous is important future work.

heading: M THE EMPIRICAL EFFECT OF LARGE K
text: In this section, we verify the claim in § 4 that with higher K we would see the accelerated variants of our algorithms perform better.
Plots are in Fig. 5. We evaluate in the setting of App. K.1. SGD (Algorithm 2), ASG (Algorithm 3), FedAvg (Algorithm 6), SCAFFOLD (Algorithm 8) are run for 100 rounds, and tuned for smallest final full gradient norm after 100 rounds according to the following learning rate grid:
η ∈ {10−3, 10−2.5, 10−2, 10−1.5, 10−1, 10−0.5, 100} (117)
For algorithms that are “1-X→Y”, we run algorithm X for one round and algorihtm Y for the rest of the rounds. These algorithms are tuned in the following way. Algorithm X has its stepsize tuned in Eq. (117), and Algorithm Y has its stepsize tuned in Eq. (117). We require these to be tuned separately because local update algorithms and centralized algorithms have stepsizes that depend on K differently if K is large (see Thm. 5, Thm. 6). These are tuned for the lowest gradient norm on the full dataset across clients.
Overall, we see that in the large K regime, algorithms that start with a local update algorithm and then finish with an accelerated centralized algorithm have the best communication complexity. This is because larger K means the error due to variance decreases, as shown in the rate of Thm. 2.
With large enough K, the exponential terms are what dominate the error, which means using ASG (Algorithm 3) as the centralized algorithm in the second superstage should be better than using SGD (Algorithm 2).

heading: N EMPIRICAL EVALUATION OF LEARNING RATE DECAY OF ONE SUPERSTAGE ALGORITHMS
text: In this section, we compare the performance of the algorithms against learning rate decay on SGD (M-SGD, Algorithm 4), ASG (M-ASG, Algorithm 5), FedAvg (M-FedAvg, Algorithm 7), and SCAFFOLD (M-SCAFFOLD, Algorithm 8 where the learning rates decay over rounds).
For all four algorithms, we tune
η ∈ {10−3, 10−2.5, 10−2, 10−1.5, 10−1} (118)
We also tune R(1) as
R(1) ∈ d{10−2R, 10−1.625R, 10−1.25R, 10−0.875R, 10−0.5R}e (119)
For all four algorithms, R(s) = R(1)2s−1. For M-SGD and M-ASG, η(s,r)g = ηK2s−1 and η (s,r) l = 0. For M-FedAvg and M-SCAFFOLD, η(s,r)g = η2s−1 and η (s,r) l = η 2s−1 . All are algorithms are tuned for the best average (across 1000 trials) final gradient norm evaluated on the full dataset across clients.
Across all heterogeneity levels, M-(SCAFFOLD→SGD) outperforms M-SCAFFOLD and M(FedAvg→SGD) outperforms M-FedAvg. All of the algorithms produced by our framework outperform all the baselines that do not have learning rate decay. As mentioned in § 4, the low value of K (20) causes algorithms that utilize M-ASG in the second superstage to underperform. A discussion on the effect of larger K on ASG is in App. M.

Review summaries:
This paper introduces the multistage optimization technique for federated learning applications. Specifically, multistage optimization first uses federated optimization algorithms like FedAvg and SCAFFOLD and converges to some budget, and then uses minibatch algorithms like SGD or accelerated SGD in order to converge faster to a point with very small error. Because centralized methods are optimal when data are heterogeneous and local methods are optimal when data are purely homogeneous, using a multistage optimization technique can incorporate the benefits from both sides.  The theoretical part is relatively easy. The proof is to choose an appropriate error budget to which federated optimization algorithms converge, and then choose the hyperparameters for the federated optimization algorithms in the first stage and the minibatch algorithms in the second stage, e.g. learning rate, momentum, etc. The theoretical part only includes the convergence results for the strongly convex case.  The empirical part includes two experiments: logistic regression and neural network, which belongs to strongly convex case and nonconvex case respectively. For each experiment, the authors compare different minibatch algorithms like SGD, AGD, different local methods like FedAvg, SCAFFOLD, and some multistage procedures that combine local methods with minibatch methods. For the convex setting, multistage procedures perform the best, and for the nonconvex setting, multistage algorithms also perform generally the best.  

It is known that if the level of heterogeneity is sufficiently high, then accelerated minibatch SGD is optimal for federated optimization matching the known lower bound of (Woodworth et al., 2020a). On the other hand, when the level of heterogeneity is very low, then FedAvg/LocalSGD outperforms the former in terms of communication complexity and needs only a few communication rounds given enough local computation. This paper *proposes* a multi-stage optimization procedure and *claims* that it nearly matches the lower bound for all heterogeneity levels.