R EDUCING THE C OMMUNICATION C OST OF F EDERATED L EARNING THROUGH M ULTISTAGE O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst-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 ﬁrst run a local method up to a heterogeneity-induced error ﬂoor; next, we switch to a centralized method for the remaining steps.
We demonstrate the scheme’s practical utility in a variety of image classiﬁcation tasks.
I NTRODUCTION
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:
[Equation]
[Equation] ) where variable x is the model parameter, i indexes the clients (or devices), N is the number of clients, and F i ( 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 F i ( x ) =
E z i ∼ P
i
[ f ( x ; z i )] for some loss function f , where z i is the client’s data and P
i ’s capture the heterogeneity of the clients’ data distributions.
Formally, we deﬁne client heterogeneity ζ 2 and heterogeneity at the optimum ζ 2 ∗ as follows, where the optimum is deﬁned as x ∗ := arg min x F ( x ) :
[Equation]
These deﬁnitions 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 & Koneˇcn`y, 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 ﬁrst-order federated optimization algorithms for strongly-convex, smooth objectives as a function of client data heterogeneity ζ ∗ (Section 2, Eq.
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 ζ ∗ .
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 F i are β -smooth and µ -strongly convex (see App.
C for deﬁnitions 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 signiﬁcant (§ 2).
To explain these gaps more precisely, we ﬁrst 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 F i 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 insigniﬁcant: ζ 2 ≤ µ(cid:15) , where (cid:15) is the target sub-optimality gap for
Eq.
(1) (Woodworth et al., 2020a)
(for more details on (cid:15) see Deﬁnition 1).
Further, an algorithm speciﬁc 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 insigniﬁcant.
(2) If so, can the same algorithm simultaneously achieve optimal convergence rates across all values of ζ ∗ ?
This Work.
We answer both questions afﬁrmatively when ζ (cid:39) ζ ∗ 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 & Koneˇcn`y, 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 beneﬁts do not appear to give order-wise convergence improvements.
Contributions.
In this paper, we ﬁrst 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 ﬁrst-order federated algorithms.
We show these theoretical gains by ﬁrst 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 ﬁrst run FedAvg up to an “optimization error ﬂoor” 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 beneﬁt of not requiring the knowledge of the noise bound σ 2 , the heterogeneity ζ 2 , or the initial distance to the optimum ∆ .
This intuition is not speciﬁc 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.
M ODEL AND R ELATED W ORK
We assume that each F i 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 deﬁned as in equation 2 (Assumptions 3 and 4, respectively).
We assume unless otherwise speciﬁed that ζ 2 > 0 , that the problem is heterogeneous.
We assume the client gradient variance is upper bounded as σ 2 , formally deﬁned 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 speciﬁed otherwise, B is 1.
Our goal is to minimize the sub-optimality error E
F (ˆ 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 (cid:15) ).
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 sufﬁciently large K .
We use the notation ˜
O to hide polylogarithmic factors.
2. R ELATED W ORK
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 ﬁrst 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; [Table] [Table] [Table]
[Table] [Table] a lower [Table]
[Table]
The fastest known rate for the worst-case sub-optimality gap of FedAvg is due to Gorbunov et al.
(2020) (Table 1):
[Equation]
[Table] enough local steps K , this error is dominated by the last term, βζ 2 µ 2 R 2 .
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)
[Equation]
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 (cid:15) accuracy, we see FedAvg’s communication R complexity is not less than that of the basic minibatch algorithms unless ζ 2 ≤ µ(cid:15) log(∆ /ε ) , i.e., heterogeneity must scale with (cid:15) .
This is not an artifact of loose analysis as Woodworth et al.
(2020a) provide the following lower bound for the convergence rate of FedAvg:
[Equation]
Considering the noiseless case, AC-SA has a convergence rate of ∆ exp( − R/ √ κ ) and FedAvg has the lower bound min { ∆ exp( − R/κ ) , κζ 2 ∗ /µR 2 } .
∆ exp( − R/κ ) is always worse than ∆ exp( − R/ √ κ ) , so FedAvg can only ever outperform AC-SA when ∆ exp( − R/ √ κ ) requires a larger R than κζ 2 ∗ /µRto reach (cid:15) accuracy.
Solving for R , this means that FedAvg’s lower bound can only result in lower communication complexity than AC-SA when ζ 2 ∗ ≤ µ(cid:15) .
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 (cid:15) for it to outperform minibatch algorithms.
Lower Bound.
The lower bound for ﬁrst-order distributed algorithms in the heterogeneous setting, from (Woodworth et al., 2020a), is (also shown in Table 1):
[Equation]
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 /R 2 , whereas in the lower bound (Eq.
, it is scaled by a term decaying exponentially fast in R assuming ζ (cid:39) ζ ∗ .
A natural question is whether this gap can be closed.
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 = for s = 1 , . .
, S do x ( s, 0) , x ( s, 1) ← x ( s − 1) for r = 1 , . . .
, R ( s ) do (cid:46) Broadcast the iterate from current round x
( s,r ) ← x ( s ) (cid:46)
Perform server update x ( s,r +1 / 2) ← (1 + γ ( s ) ) x ( s,r ) − γ ( s ) x ( s,r − 1) (cid:46) 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 sufﬁcient 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 & Kone ˇ cn ` y, 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).
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)
M ULTISTAGE F EDERATED M INIMIZATION
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 ﬁrst initializes their iterate as x ( s,r )
i,= 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 ˜ ∇ F i ( 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,
[Equation]
i
[Equation]
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, deﬁned in Algorithm 9) at step k use the update rule:
[Equation]
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
[Equation]
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.
[Equation]
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
[Equation]
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 ) =
W ( s )
R (cid:88) r =
K (cid:88) k
= w ( s,r ) k
N
N (cid:88) i =
x ( s,r ) i,k ) , where W ( s ) =
(cid:80)
Rr =1 (cid:80)
Kk =1 w ( s,r ) k .
Sums of local iterates (cid:80)
Kk =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.
3. T WO STAGE ALGORITHMS
We begin with a two-stage instantiation of Algorithm 1.
We run FedAvg for a constant fraction of rounds until reaching the “heterogeneity ﬂoor”, followed by a minibatch algorithm until convergence.
Intuition: Fig.
1 shows a toy example.
Consider two strongly convex client objectives:
F 1 ( x ) = 12 ( x − 1) 2 and F 2 ( x ) =
( x + 1) 2 .
The global objective F ( x ) = F 1 ( x )+ F 2 ( x )is their average.
Fig.
1 (top) plots the objectives, and Fig.
1 (bottom) displays their gradients.
x ∈ ( , − 1] ∪ [1 , ∞ )
In this .g. , FedAvg) to save on communication while still ensuring that perclient local updates are consistent with each other.
the hand, the optimum i.e., when x ∈ ( − 1 , 1) , gradients may the optimum.
optimum the optimum, 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 StackOverﬂow 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 ﬁnd that two stages sufﬁce 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 ﬁrst 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 (speciﬁed in Thm. 4), such that (i) the ﬁrst stage runs FedAvg for R/ 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) satisﬁes
[Equation]
(Proof in App.
D)
This algorithm exponentially improves the dependence of the round complexity on heterogeneity from the state of the art βζ {+
+} µ 2 R 2 to ( β 2 ζ 2 µ 3 R 2 ) 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 ζ (cid:39) ζ ∗ ) 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 (speciﬁed in Thm. 5), such that (i) the ﬁrst 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) satisﬁes
[Equation]
(Proof in Section E).
Observe that this rate is strictly better than the FedAvg rate (Eq.
and is also strictly better than the AC-SA rate (Eq.
when βζ 2 µ 2 R 2 ≤ ∆ (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
3. M ULTISTAGE ALGORITHMS
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 ﬁrst 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 = (cid:100) κ (cid:101) .
In the ﬁrst 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
[Equation]
(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 µ 3 R 2 ; 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 ) . [Figure] [Table] [Figure]
E XPERIMENTS
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.
; our goal is instead to show how multistage optimization compares to single-stage baselines for ﬁxed R .
Convex Optimization (Logistic Regression)
We ﬁrst 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 shufﬂing samples from different digit classes.
The details of this shufﬂing process are described in App.
K.1; if a dataset is more shufﬂed (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 classiﬁcation tasks with convolutional neural networks.
We started with digit classiﬁcation 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 difﬁculty.
We also considered image classiﬁcation 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).
C ONCLUSION
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 beneﬁts 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.
A D ISCUSSION : S TEPSIZE D ECAY IN P RACTICE
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 & Koneˇcn`y, 2020), thus potentially explaining the empirical gains they see in their papers. We expand on the learning rate decay of LocalUpdateDecay (Charles & Koneˇcn`y, 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.
η l /η g = 0 on the centralized algorithm line because η l = 0 for a centralized algorithm.
For FedAvg (Algorithm 6), η l /η g = 1 [-N (cid:80) Ni =
 δ ri and then updates the global model with outer stepsize-]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. decay events (the precise number is not constant and is based on how much error decreases between rounds).
X axis is the number of learning rate decay events. Y axis is the ratio between local and global learning rates.
We see that within 2 decay events LocalUpdateDecay is closely approaching the centralized learning rate ratio.
B A LGORITHM P SEUDOCODE Algorithm
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 (cid:46) Run a step of SGD for client i = 1 , . . .
, N in parallel do
Sample K datapoints { z i,k } 1 ≤ k ≤ K
Client communicates (cid:80)
Kk
=1 ∇ f ( x ( r ) ; z i,k ) to server end for x ( r +1) ← x ( r ) − ηNK (cid:80) Ni =1 (cid:80)
Kk =1 ∇ f ( x ( r ) ; z i,k ) end for return x ( R )
B. A LGORITHM 2 IN TERMS OF A LGORITHM
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
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 (cid:46) Run a step of ASG for client i = 1 , . . .
, N in parallel do
Sample K datapoints { z i,k } 1 ≤ k ≤ K
Client communicates (cid:80)
Kk
=1 ∇ f ( x ( r ) ; z i,k ) to server end for γ ( r ) ← 1 −√ ηµ 1+ √ ηµ y ( r ) ← (1 + γ ( r ) ) x
( r ) − γ ( r ) x ( r − 1) x ( r +1) ← y ( r ) − ηNK (cid:80) Ni =1 (cid:80)
Kk =1 ∇ f ( y ( r ) ; z i,k ) end for return x ( R )
B.
A LGORITHM 3 IN TERMS OF A LGORITHM
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) ) .
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 (cid:46) Run a step of SGD for client i = 1 , . . .
, N in parallel do
Sample K datapoints { z i,k } 1 ≤ k ≤ K
Client communicates (cid:80)
Kk
=1 ∇ f ( x ( s,r ) ; z i,k ) to server end for x ( s,r +1)
← x ( s,r ) − η ( s ) NK (cid:80) Ni =1 (cid:80)
Kk =1 ∇ f ( x ( s,r ) ; z i,k ) end for end for return x ( S,R ( S ) )
B. A LGORITHM 4 IN TERMS OF A LGORITHM
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 ) ) .
B. A LGORITHM 5 IN TERMS OF A LGORITHM
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 ) ) .
) , 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 (cid:46) Run a step of ASG for client i = 1 , . . .
, N in parallel do
Sample K datapoints { z i,k } 1 ≤ k ≤ K
Client communicates (cid:80)
Kk
=1 ∇ f ( x ( s,r ) ; z i,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 (cid:80) Ni =1 (cid:80) Kk =1 ∇ f ( y ( s,r ) ; z i,k ) end for x ( s ) ← x ( s,R ( s )+1 ) end for return x ( S,R ( s ) ) B A LGORITHM 6 IN TERMS OF A LGORITHM
Set
[Equation]
For ClientOPT , [Equation] where the ˜ ∇ F i ( x ( s,r ) i,k − 1 ) ’s come from
[Equation]
Algorithm
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 z i,k −x ( r ) i,k ← x ( r ) i,k − 1 − η · ∇ f ( x ( r ) i,k − 1 ; z i,k − 1 ) end for
Client communicates (cid:80)
Kk
=1 ∇ f ( x ( r ) i,k − 1 ; z i,k − 1 ) to server end for (cid:46) Perform server update x ( r +1) ← x ( r ) − ηN (cid:80) i ∈
[ N ] (cid:80)
Kk =1 ∇ f ( x ( r ) i,k − 1 ; z i,k − 1 ) end for return x ( R )
B. A LGORITHM 7 IN TERMS OF A LGORITHM
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 (
[Equation]
[Equation] where the ˜ ∇ F i ( x ( s,r ) i,k − 1 ) ’s come from x [Equation]
←
[Equation]
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,← x ( s,r ) for k = 1 , ...,
K do
Randomly select sample z i,k −x ( s,r ) i,k ← x ( s,r ) i,k − 1 − η ( s ) · ∇ f ( x ( s,r ) i,k − 1 ; z i,k − 1 ) end for Client communicates (cid:80) Kk =1 ∇ f ( x ( s,r ) i,k − 1 ; z i,k − 1 ) to server end for (cid:46) Perform server update x ( s,r +1)
← x ( s,r ) − η ( s ) N (cid:80) i ∈ [ N ] (cid:80) K k =1 ∇ f ( x ( s,r ) i,k − 1 ; z i,k − 1 ) end for end for return x ( R ) B. A LGORITHM 8 IN TERMS OF A LGORITHM
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 (cid:88)
[Equation]
where the ˜ ∇ F i ( x ( s,r ) i,k − 1 )
’s come from
[Equation]
where c i =
[Equation]
) if the previous round is in the same stage and c i =
[Equation]
) if the previous round is in an earlier stage.
c = 1 N (cid:80) Ni =1 c i .
In Algorithm 1 we do not specify states (for c and c i ), though one can easily include them.
Input: Number of local steps K , rounds R , stepsize η l , η g , initial point x (0) , control variates c and c i , 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 z i,k −x ( r ) i,k ← x ( r ) i,k − 1 − η l · ( ∇ f ( x ( r ) i,k − 1 ; z i,k − 1 ) + c − c i ) end for c i ← 1 K (cid:80) K − 1 k =0 ∇ f
( x ( r ) i,k − 1 ; z i,j )
Client communicates (cid:80) K − 1 k =0 ∇ f ( x ( r ) i,k − 1 ; z i,k ) to server
Client communicates c i to server end for (cid:46) Perform server update x ( r +1) ← x ( r ) − η g N (cid:80) i ∈
[ N ] (cid:80) K − 1 k =0 ∇ f
( x ( r ) i,k − 1 ; z i,k ) c ← 1 N (cid:80) Ni =1 c i end for return x
( R )
Input: Number of local steps K , rounds R , stepsize η l , η g , initial point x (0) , number of devices N for r = 0 , . . .
R − do
Server communicates x ( r ) to all clients
For client i ∈ [ n ] compute c i ← 1 K (cid:80) Ki =1 ∇ f ( x ( r ) ; z i,k ) for randomly drawn z
i,k
Clients communicate c i ’s to server c ← 1 N (cid:80) Ni
=1 c i
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 z i,k −x ( r ) i,k ← x ( r ) i,k − 1 − η · ( ∇ f ( x ( r ) i,k − 1 ; z i,k − 1 ) + c − c i ) end for c i ← 1 K (cid:80) K − 1 k =0 ∇ f
( x ( r ) i,k − 1 ; z i,j )
Client communicates (cid:80) K − 1 k =0 ∇ f ( x ( r ) i,k − 1 ; z i,k ) to server end for (cid:46) Perform server update x ( r +1) ← x ( r ) − ηN (cid:80) i ∈
[ N ] (cid:80) K − 1 k =0 ∇ f
( x ( r ) i,k − 1 ; z i,k ) end for return x ( R )
B. A LGORITHM 9 IN TERMS OF A LGORITHM
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 (cid:88) i = K (cid:88) k
= ˜ ∇ F i ( x ( s,r ) i,k − 1 ) . where the ˜ ∇ F i ( x ( s,r ) i,k − 1 ) ’s come from
[Equation]
[Equation]
where c i = 1 K (cid:80)
Ki =1 ˜ ∇
F i ( x ( s,r ) ) , c = 1 N (cid:80) Ni
=1 c i .
C T ECHNICALITIES
In this section we collect all the assumptions we make in the paper for convenience.
C. A SSUMPTIONS
Assumption 1.
F i ’s are µ -strongly convex for µ > 0 , i.e.
[Equation]
for any i , x , y .
Assumption 2.
F i ’s are β -smooth where
[Equation]
for any i , x , y .
Assumption 3.
F i ’s have a uniform upper bound on the gradient dissimilarity (heterogeneity), deﬁned as
[Equation]
Assumption 4.
F i ’s have an upper bound on the gradient dissimilarity (heterogeneity) at the optimum, deﬁned as
[Equation]
Assumption 5.
We assume that the F i ’s have bounded variance of stochastic gradients, where z i is a randomly drawn data point from client i ’s data: [Equation]
[Equation]
Assumption 6.
We assume that the initial suboptimality gap is upper bounded by ∆ , where x 0 is the initial feasible iterate and x ∗ is the global minimizer. [Equation]
Assumption 7.
We assume that the initial distance from the optimum is upper bounded by D , where x 0 is the initial feasible iterate and x ∗ is the global minimizer.
[Equation]
C. D EFINITIONS
Deﬁnition 1.
An algorithm is said to have converged, or to have satisﬁed the convergence criterion if
[Equation]
where (cid:15) is some user-chosen value.
Deﬁnition 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.
[Equation]
where M is a function that returns a one-dimensional value.
D P ROOF OF T HM .
Theorem 4 (Formal version of Thm. 1) .
Run Algorithm 1 over S = 2 stages.
In the ﬁrst stage, use FedAvg 4 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:
If R > 4 βµ and r < (cid:100) R 2 (cid:101) : η (2 ,r ) l = η (2 ,r ) g [Equation]
η g [Equation] 14 β , w (2 ,r ) k
[Equation]
If R > 4 βµ and r ≥ (cid:100) R/ 2 (cid:101) :
η (2 ,r ) l
= 0 , η (2 ,r ) g
[Equation]
[Equation]
R −(cid:100) R 2 (cid:101) )
[Equation]
w (
[Equation]
r ) k
[Equation]
If R ≤ 4 βµ : η (2 ,r ) g = 14 β , η (2 ,r ) l
[Equation]
and [Equation] [Equation] 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
[Equation]
Proof.
First, let the total budget for the number of rounds be 2 R .
We run the two stages of the algorithm with R rounds.
In the ﬁrst stage, we run Algorithm 1 with
[Equation]
S
[Equation]
[Equation]
=
[Equation]
) rK
[Equation]
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) [Equation] we know that we get the following convergence rate for the output of the algorithm
[Equation]
where Φ = log(max { 2 , min { E (cid:107)
[Equation]
In the second stage
[Equation]
we run Algorithm 1 on x (1) with β ( s ) = 0 and
If
R > βµ
and
r
< (cid:100) R 2 (cid:101) η (2 ,r ) l = 0 , η (2 ,r ) g = η g =
14 β , w (2 ,r ) k
=2.
If R > 4 βµ and r ≥ (cid:100) R/ 2 (cid:101) :
η (2 ,r ) l = 0 , η (2 ,r ) g
= η g = 2 µ ( κ − R −(cid:100) R 2 (cid:101) ) , w (2 ,r ) k
= ( κ + r − (cid:100) R 2 (cid:101) )3.
If R ≤ 4 βµ : η (2 ,r ) l = 0 , η (2 ,r ) g = 14 β , w (2 ,r ) k
= (1 − µη g ) − ( r +1) .
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 [Equation]
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 (cid:15) in the ﬁnal convergence rate.
We start by setting
[Equation]
Next we set
[Equation]
[Equation]
Finally we set
[Equation]
[Equation]
So we need that
[Equation] This implies that [Equation] Plugging this into Eq.
, we get [Equation]
E P ROOF OF T HM .
Theorem 5 (Formal version of Thm. 2) .
Run Algorithm 1 over S = 2 stages.
In the ﬁrst 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:
[Equation]
= w [Equation] = 0 for all k, r except for w (2 ,R/ 2)= 1 (i.e. return the last iterate). Also let
[Equation]
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
[Equation]
Proof.
First, let the total budget for the number of rounds be 2 R .
We run the two stages of the algorithm with R rounds.
In the ﬁrst stage [Equation] we run Algorithm 1 with
[Equation]
= 1 , [Equation] =
[Equation]
= (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
[Equation]
where Φ = log(max { 2 , min { E
[Equation]
Next we run Algorithm 1 on x (1) with η ( s,r ) l = 0 and [Equation] γ
[Equation]
1+
[Equation]
g K , S = 1 , and w ( s,r ) k = 0 for all k, r, s except for w (1 ,R 1 )= 1 and ∆ 1 ≤
[Equation]
From Thm.
6 we know that
[Equation]
46) where Φ 2 = log(max { 2 , min { 2∆ 1 µKσ, 2∆ µ
[Equation]
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. We start by setting
[Equation]
Next we set
[Equation]
Finally we set [Equation] So we need that [Equation]
This implies that
[Equation]
Plugging this into Eq.
(46), we get
[Equation]
Theorem 6.
Consider Nesterov’s Accelerated Stochastic Gradient (ASG) method with the iteration sequence
[Equation] [Equation]
with γ = 1 −√ αµ 1+ √ αµ and
[Equation]
} , then
[Equation] [Equation]
Proof.
From (Aybat et al., 2019) (Eq.13), we know that
[Equation]
(61) where F ( x k ) − F ( x ∗ ) ≤ V k ≤ 2( F ( x k [Equation] − F ( x ∗ )) .
If we unroll this, we get
E
[Equation]
Now if
[Equation]
63) observe that if 1 β
[Equation]
we know that letting Φ = log
[Equation]
64) and if 1 β > log(max [Equation] 2 , min
[Equation]
where we used that EV 0 ≤ 2∆ .
So altogether, we have that
[Equation]
G P ROOF OF T HM .
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 ﬁrst 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): η [Equation] s [Equation]
r ) l = 0 , η (1 ,r ) g
= 1 Kβ with R (1) ≥
R (1) =
(cid:100) ( p + 1) √ κ log(12( p + 1) κ ) (cid:101) η
( s,r ) g
= 1 K 2
s β with R ( s ) =
2 s (cid:100)√ κ log(2 p +2 ) (cid:101) for any p ≥ 1 and s ≥γ
[Equation]
s
[Equation]
s
[Equation]
r
[Equation]
g
K
[Equation]
µη
[Equation] s [Equation] r
[Equation]
g
K
For B large enough, this achieves expected convergence rate of
[Equation]
Proof.
Let x ( s (cid:48) ) be the output of multistage FedAvg from the ﬁrst half of the rounds ( s (cid:48) is the index of the ﬁnal stage).
Then we know that the output of this has guarantee (from Thm. [Equation]
[Equation]
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 )
[Equation]
where
[Equation]
[Equation]
If we set K = κ as per the limit, then we have
[Equation]
Plugging ∆ 1 in and K
= κ
[Equation]
we get [Equation] Which can be simpliﬁed to (noting that µ ≤ β ) [Equation]
H M ULTISTAGE F ED A VG (M-F ED A VG )
Theorem 8.
Instantiate Algorithm
as
[Equation]
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 ≤ (cid:100) pκ log(2) (cid:101) , and (cid:100) pκ log(2) (cid:101) mod K = 0 .
Then if R (1) = R/C for some constant and
KR ≥ 2 κ ,
[Equation]
Proof.
We start by proving that
[Equation]
Where err s = E (cid:107) x
( s,R s ) − x ∗ (cid:107) 2 .
We prove this by induction on s .
Recall that from Thm.
9 we have
[Equation]
Starting with the base case, with η = 14 β , R 1 K ≥ 1 , we have
[Equation]
Now let us assume the result holds for s and prove it for s + 1 .
Using Thm. 9 for stage s + 1 with η = 1 β 2 s +3 and R s +1 K =
(cid:100)
p 2 s +3 κ log(2) (cid:101) gets us
[Equation]
Now using the induction hypothesis, we know
[Equation]
And substituting it back into (Eq. (73)) while having p ≥ 3 gets us
[Equation]
Using the following lemma, we can translate this into a convergence rate:
Lemma 1.
Let R (1) K ≥ 1 and R ( s ) K
=
(cid:100) p 2 s
+2 κ log(2) (cid:101) .
Then we know that 2 k ≥ Θ(1)( RK − R (1) K pκ ) .
Proof.
Comes from a simple unrolling of the stage lengths; similar to the proof in Fallah et al.
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.
I P ROOF OF T HM .
Theorem 9.
Suppose we ran Algorithm 1 for one stage with stepsizes η ( s,r ) l = η ( s,r ) g
= η , with γ ( s ) =
and return the last iterate.
Furthermore, require that η ≤ 1 β .
Then
[Equation]
From (Gorbunov et al., 2020) (Eq.82), we have the following for FedAvg:
[Equation]
where V t = 1 N (cid:80) Ni =1 E (cid:107) x ti − x t (cid:107) 2 , where t =
Kr
+ k where r is current round, k is step on current round.
x t i is client
i ’s iterate at the t -th step.
x t = 1 N (cid:80) 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
[Equation]
So together,
[Equation]
And unrolling this,
[Equation]
J U SING SS-L OCAL -SGD INSTEAD OF F ED A VG
In this section we show how one can use SS-Local-SGD instead of FedAvg as ClientOPT.
Theorem 10.
Let V t = 1 N (cid:80)
Ni =1 E (cid:107) x ti − x t (cid:107) 2 , where t =
Kr
+ k where r is current round, k is step on current round, x ti is client
i ’s iterate at the t -th step.
Then SS-Local-SGD has
[Equation]
Proof.
[Equation]
By Jensen’s.
Let g i ( x ti ) be the SS-Local-SGD local update for client i at iteration t .
Let E t denote expectation conditioned on everything up to the t − 1 ’th iteration.
Let y t be the last synchronized iterate before iteration t .
In the following we use the fact that the control iterates are computed with
[Equation]
[-K − 1) E (cid:107) x t − 1 i − x t − 1 j (cid:107) 2 + 16 Kη 2 ζ 
-]
So altogether,
[Equation]
Unrolling this until the last time the iterates were synchronized [Equation] we get
[Equation]
Which implies
[Equation]
This can be substituted into Thm.
9 to get the multistage rate for Thm. [Equation] or it can be substituted into
Theorem G.3 in (Gorbunov et al., 2020) to get a convergence rate that differs from (Eq.
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.
K E XPERIMENTAL S ETUP D ETAILS K. C ONVEX O PTIMIZATION
We empirically evaluate multistaging on federated regularized logistic regression.
Let ( x i,j , y i,j ) be the j th datapoint of the i th client.
We minimize Eq.
(1) where
[Equation]
[Table]
Dataset.
We use the MNIST dataset of handwritten digits (LeCun et al., 2010).
We model a federated setting with ﬁve 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 ﬁrst X % of each class’s 500 images is shufﬂed and evenly partitioned to each client.
The remaining (100 − X )% is partitioned as follows: client i ∈ { 1 , . .
, 5 } receives the remaining non-shufﬂed data from classes 2 i − 2 and 2 i − 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 ζ =
due to sampling randomness; we use this technique for lack of a better control over ζ .
To model binary classiﬁcation, 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 .
We ﬁx 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
[Equation]
in the range below:
[Equation]
For the second stage we let the stepsizes be η (2 ,r ) l
[Equation] [Equation]
/K .
We also tune the percentage of rounds before switching to the second stage in the range as
[Equation]
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 ﬁrst 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 ﬁrst stage.
For subsequent stages s ≥ 2 , we use step sizes η ( s,r ) l
= η ( s,r ) g =
η 2 s − 1 and
R ( s ) =
2 s − 1 R (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.
K. N ONCONVEX O PTIMIZATION K.2. EMNIST
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 ﬁxing the number of local steps, we ﬁx 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 & Koneˇcn`y, 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 ﬁrst 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 ﬁrst stage.
We tune R (1) ∈ (cid:100){ 0 .
1 R, 0 . 275 R, 0 . 45 R }(cid:101) . For subsequent stages s ≥ 2 , we use step sizes η ( s,r ) l
= η ( s,r ) g = η 2 s − 1 and R ( s ) = 2 s − 1 R (1) . After R (cid:48) rounds have passed, we enter the second superstage.
We tune R (cid:48) ∈ (cid:100){ 0 . 5 R, 0 . 7 R, 0 .
9 R }(cid:101) . Let the last stage of superstage 1 be s (cid:48) . In the subsequent stages, we let R ( s ) = 2 s − s (cid:48)
− 1 R (1) and η ( s,r ) g
=
η ¯ K 2 s − s (cid:48) − 1 and η ( s,r ) l = 0 . Since we
[Table] 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 . and R (1) ∈ (cid:100){ 10 − 2 R, 10 − 1 .
625 R, 10 − 1 .
25 R, 10 − 0 .
875 R, 10 − 0 .
5 R }(cid:101) We then set η ( s,r ) g = η
¯ K 2 s − 1 , η ( s,r ) l = 0 , and R s =
R (1) 2 s − 1 . For multistage SCAFFOLD and FedAvg, we tune η ∈ { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 } and R (1) ∈ (cid:100){ 0 . 1 R, 0 .
275 R, 0 . 45 R }(cid:101) . We then set η ( s,r ) g =
η ¯ K 2 s − 1 , η ( s,r ) l = 0 , and R s = R (1) 2 s − 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. CIFAR-
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 ﬁxing the number of local steps, we ﬁx 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 & Koneˇcn`y, 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.
K. N ONCONVEX O PTIMIZATION
Table 2 describes the ﬁnal test accuracy of the learned models for each optimization technique.
Tableincludes the test accuracy for our experiments on EMNIST classiﬁcation, with the experimental tuning details in App. K.2.1.
We observe that SCAFFOLD → SGD has the best test accuracy.
L D ISCUSSION ON P ARTIAL C LIENT P ARTICIPATION
j
Client j sets x j,
=for k = 0 , . . . , K − 1 do x (0) j,k +1 ← x (0) j,k − η (1) · ∇ F i ( x (0) j,k ) end for
Client j sets x = W K (cid:80) = do
S x ( r +1)
← x ( r ) − η (2) r · 1 SK (cid:80) i ∈S , 1 ≤ k ≤ K ∇ F i ( 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.
We run SGD on one client for K local steps in the ﬁrst round.
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 deﬁnitions:
[Equation]
Theorem 11.
Let ¯ x j,K =
[Equation]
where w k =
[Equation]
[Equation]
and
[Equation]
where x (0) j,k are generated in the following way:
[Equation]
{+where x (0) j, 
+}
= x (0) .
For any j ∈ [ N ] , ¯ x j,K satisﬁes
[Equation]
[Equation]
as long as η (1) ≤ µ 8 β 2 .
Proof.
Shorthand x (0) j,k as x k and η (1) as η 1 for the duration of this proof.
Then
[Equation]
[Equation]
The ﬁrst term of Eq.
(96) can be upper bounded using strong convexity Assumption 1:
[Equation]
The second term can be upper bounded using the Fenchel-Young inequality:
[Equation]
Where in the last inequality we used the deﬁnition of ˆ ζ 2 .
Next, we upper bound the third term of Eq.
[Equation]
where in the last step we used smoothness Assumption 2 for the ﬁrst term and the deﬁnition of ˆ ζ 2 ∗ for the second term.
We can substitute all our upper bounds into Eq.
(94) to get
[Equation]
[Equation]
Setting η 1 ≤ µ 8 β 2 , we know that 2 η 21 β 2 ≤ η 1 µ 4 and that 2 η 21 ˆ ζ 2 ∗ ≤ η
[Equation]
ˆ ζ 2 ∗ 4 β 2 ≤ η 1 ˆ ζ 2 ∗ 4 µ ≤ η 1 ˆ ζ 2 4 µ .
So we have
[Equation]
[Equation]
And rearranging,
[Equation]
Using convexity of F , and letting ¯ x K :
=
¯ x j,K , we know if w k =
[Equation]
1 − η 1 µ 4 ) − k , then
[Equation]
[Equation]
Noting again that η 1 ≤ µ 8 β 2 and that 1 W K ≤ 1 w K ≤ (1 − η 1 µ 4 ) K , we have
[Equation]
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 speciﬁed from Thm.
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))
If R > 4 βµ and r − 1 < (cid:100) R 2 (cid:101) :
η (2) r = η g = 14 β , w r =2. If R > 4 βµ and r − 1 ≥ (cid:100) R/ 2 (cid:101) : η (2) r = 2 µ ( κ − R −(cid:100) R 2 (cid:101) ) , w r = ( κ + r −
(cid:100) R 2 (cid:101) )3.
If R ≤
[Equation]
Where Algorithm 10 returns
[Equation]
Then we have the following convegence rate for Algorithm 10 if K ≥ κ 2 log β 2 D 2 ˆ ζ 2 :
[Equation]
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
[Equation]
We can substitute for E F ( x (1) )
− F ( x ∗ ) the ﬁnal suboptimality from Thm.
11, which gives
[Equation]
If we set K ≥ κ 2 log β 2 D 2 ˆ ζ 2 where D 2 ≤ (cid:107) x (0) − x ∗ (cid:107) 2 , then the convergence rate becomes
[Equation]
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
[Equation] 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.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.
M T HE E MPIRICAL E FFECT OF L ARGE K
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.
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 ﬁnal full gradient norm after 100 rounds according to the following learning rate grid:
[Equation]
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.
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 ﬁnish 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).
N E MPIRICAL E VALUATION OF L EARNING R ATE D ECAY OF O NE S UPERSTAGE A LGORITHMS
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
[Equation]
We also tune R (1) as
[Equation]
(cid:101) (119)
For all four algorithms, R ( s ) =
R (1) 2 s − 1 . For M-SGD and M-ASG, η ( s,r ) g = η K 2 s − 1 and η ( s,r ) l = 0 . For M-FedAvg and M-SCAFFOLD, η ( s,r ) g =
η 2 s − 1 and η ( s,r ) l =
η 2 s − 1 .
All are algorithms are tuned for the best average (across 1000 trials) ﬁnal 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.