STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning
Authors: Prashant Khanduri
PRANAY SHARMA
Haibo Yang
Mingyi Hong
Jia Liu
Ketan Rajawat
Pramod Varshney
Venue: NeurIPS 2021 Poster (2021)
Keywords: Distributed Non-Convex Optimization
Minibatch SGD
Local SGD

Abstract: Federated Learning (FL) refers to the paradigm where multiple worker nodes (WNs) build a joint model by using local data. Despite extensive research, for a generic non-convex FL problem, it is not clear, how to choose the WNs' and the server's update directions, the minibatch sizes, and the local update frequency, so that the WNs use the minimum number of samples and communication rounds to achieve the desired solution. This work addresses the above question and considers a class of stochastic algorithms where the WNs perform a few local updates before communication. We show that when both the WN's and the server's directions are chosen based on certain stochastic momentum estimator, the algorithm requires $\tilde{\mathcal{O}}(\epsilon^{-3/2})$ samples and $\tilde{\mathcal{O}}(\epsilon^{-1})$ communication rounds to compute an $\epsilon$-stationary solution. To the best of our knowledge, this is the first FL algorithm that achieves such {\it near-optimal} sample and communication complexities simultaneously.  Further, we show that there is a trade-off curve between local update frequencies and local minibatch sizes, on which the above sample and communication complexities can be maintained. {Finally,   we show that for the classical FedAvg (a.k.a. Local SGD, which is a momentum-less special case of the STEM), a similar trade-off curve exists, albeit with worse sample and communication complexities. Our insights on this trade-off provides guidelines for choosing the four important design elements for FL algorithms, the update frequency, directions, and minibatch sizes to achieve the best performance.} 

STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning
Prashant Khanduri University of Minnesota
Pranay Sharma
Carnegie Mellon University
khand095@una sarad0995umm....m edu
pranaysh@andrew. cmu. edu
Haibo Yang
Mingyi Hong"
Jia Liu
The Ohio State University
University of Minnesota mhongSumn. edu
The Ohio State University liu@ece. osu. osu.edu
yang. hangg55922bbuukyymalll.. osu. edu
Ketan Rajawat Indian Institute of Technology Kanpur ketan@iitk ac. in
Pramod K. Varshney
Syracuse University
varshneydsyr edu edu
Abstract
Federated Learning (FL) refers to the paradigm where multiple worker nodes (WNs) build a joint model by using local data. Despite extensive research, for a generic non-convex FL problem, it is not clear, how to choose the WNs' and the server's 's update directions, the minibatch sizes, and the number of local updates, so that the WNs use the minimum number of samples and communication rounds achieve the desired solution. This work addresses the above question and considers class of stochastic algorithms where the WNs perform few local updates before communication. We show that when both the WN's and the server's 's directions are chosen based on certain stochastic momentum estimator, the algorithm requires ((22) samples and Õ communication rounds compute e-stationary solution. To the best of our knowledge, this is the first FL algorithm that achieves such netr-optimal sample and communication complexities simultaneously. Fur- ther. we show that there is trade-off curve between the number of local updates and the minibatch sizes, on which the above sample and communication complexi- ties can be maintained. Finally, we show that for the classical FedAvg (a.k.a. Local SGD, which is a momentum-less special case of the STEM), a similar trade-off curve exists, albeit with worse sample and communication complexities. Our insights on this trade-off provides guidelines for choosing the four important design elements for FL algorithms, the number of local updates, WNs' and server's update directions, and minibatch sizes to achieve the best performance.
Introduction
In Federated Learning (FL), multiple worker nodes (WNs) collaborate with the goal of learning
joint model, by only using local data. Therefore it has become popular for machine learning problems where datasets are massively distributed M In FL, the data is often collected at or lof-ed to aim lear centralized model
multiple WNs which collaboration with a server node
"Correspondnng Author: Mingyi Hong.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Trade-off: Batch Sizes and Local Updates
Batch Sizes vs of Local Updates
30
Comm. complexity Ofe 32)
Communication Complexity
x
x100*
25
SASTEM
#f Local Updatess
20
Sample Complexity -O,- $2)
Comm. complexity - O((")
15
2
Minibatch STEM
10-
1
FedSTEM
0
Minibatch STEM
Comm. complexity: Ot " 2-20
10
Batch Size,
30
of Local Updates, I
40
40
5
Batch sizes,
15
20
30
(a) Communication complexity.
(b) Minibatch sizes vs Local Updates.
Figure The 3D surface (a) plats the communication complexity the proposed STEM for different minibatch sizes and number of local updates. The surface is generated such that each point represents STEM with a particular choice of (b, I). so that requires (3/2) samples to achieve e-stationarity. Plot (b) shows the optimal trade off between the minibatch sizes and the number of local updates at each WN (i.e.. achieving the lowest communication and sample complexities). Both plots are generated for an accuracy of = 10 and all the constants dependent system parameters (variance of stochastic gradients, heterogencity parameter, optimality gap, Lipschitz constants, etc.) assumed to be 1. Fed STEM is a special case of STEM where (1) minibatch is used; Minibatch STEM a special case of STEM where ((1) local updates are used.
The local WNs share the computational load and since the data is local to each WN, FL. also provides some level of data privacy A classical distributed optimization problem that K WNs aim to solve:
K
K
(z) f(x) = 1 Σ K
= K Σ B (f(0) (x: (x; (0))
min ER
(I)
where f(A) ; _d RR denotes the smooth (possibly non-convex) objective function and g(b) ~ D()) represents the sample/s drawn from distribution D(k) at the 1tth WN with h E [K]. When the distributions (A)) are different across the WNs, it is referred to as the heterogeneous data setting.
The optimization performance of non-convex FL algorithms is typically measured by the total number of samples accessed (cf. Definition 2.2) and the total rounds of communication (cf. Definition 2.3 required by each WN achieve an e-stationary solution (cf. Definition 2.1 To minimize the sample and the communication complexities, FL algorithms rely on the following four key design elements: (i) the WNs' local model update directions, (ii) Minibatch size compute each local direction, (iii) the number of local updates before WNs share their parameters, and (iv) the SN's update direction. How to find effective FL algorithms by (optimally) designing these parameters has received sigmificant research interest recently.
Contributions. The main contributions of this work are listed below:
1) We propose the Stochastic Two-Sided Momentum (STEM) algorithm, that utilizes certain momentum coosentum-assistee stochastic gradient directions for both the WNs and SN updates. We show that there exists an optimal trade off between the minibatch sizes and number of local updates, such that on the trade-off curve STEM requires (32-3/2 2 samples and ((--1) communication rounds to reach an e-stationary solution; see Figure for an illustration. These complexity results are the best achievable for first-order stochastic FL algorithms (under certain assumptions, cf. Assumption []; see (5-8) and 99 I0,, as well as Remark of this paper for discussions regarding optimality. To the best of our knowledge, STEM is the first algorithm which - (i) simultaneously achieves the optimal sample and communication complexities for FL and (ii) can optimally trade off the minibatch sizes and the number of local updates.
2) A momentum-less special case of our STEM result further reveals some interesting insights of the classical FedAvg algorithm (a. the Local SGD) [T. Specifically, we show that for FedAvg, there also exists a trade-off between the minibatch sizes and the number of local updates, such that it
requires (E samples and 3322) communication rounds to achieve an e-stationary solution.
The notation Ö(-) hides the logarithmic factors.

---

Algorithm
Work
Sample
Comm.
Minibatch (6)
Local Updates (4) /round
12
Ole 0/4)))
((1)
O(E 1/2)
FedAvg"
15
O
3)
0(1)
0(1)
20
this work
O(e
(4
O(< 2ा m)
SCAFFOLD'
15
0(c-9)
O(6-2)
0(1)
((1)
FedPD/FedProx/
O(e-2)
0(c-))
0(1)
O(c=))
NIME//FeGH
17 I
O(e 3/2)
0()7)1))
0(1)
((1)
STEM
O(E
O Ta-05
Fed STEM
this work
Ōcc-8/2)
Ö(e-1)
O(1)
O ((-1//2
Minibatch STEM"
013
O(1)
Table 1:
Comparison of FedAvg and STEM with different FL algorithms for various choices of the minibatch
sizes (b) and the number of per node local updates between two rounds of communication (I).
Pv [0, 1] trades off b and I: = 1 (resp. multiple) samples. Fed STEM and Minibatch STEM two variants of the proposed STEM.
0) uses multiple (resp. (()) local updates and ( (resp.
"The data heterogeneity assumption is weaker than Assumption 2
(please see [9] for details).
RRequires bounded Hessian dissimilarity to model data heterogeneity across WNs.
"Guarantees for Minibatch STEM with / = 1 and SCAFFOLD are independent of the data heterogeneity.
Collectively, insights on the trade-off's provide practical guidelines for choosing different design elements for FL algorithms.
Related Works. FL. algorithms were first proposed in the form of FedAvg where the local update directions at each WN were chosen to be the SGD updates. Earlier works analyzed these algorithms in the homogeneous data setting T 25 while many recent studies have focused on designing new algorithms to deal with heterogencous data settings, as well as problems where the local loss functions non-convex p
32 In [2]. the authors showed that Parallel Restarted SGD (Local SGD or FedAvg III achieves linear speed up while requiring
JC
12
16
26
O(e=2) samples and ( 3/2)) rounds of communication to reach an e-stationary solution. In [14), Momentum SGD proposed, which achieved the same sample and communication complexities as Parallel Restarted SGD [12. without requiring that the second moments of the gradients be bounded. Further, it was shown that under the homogeneous data setting, the communication complexity can be improved to O(e while maintaining the same sample complexity. The works in 15 T6 conducted tighter analysis for FedAvg with partial WN participation with (1) local updates and batch sizes. Their analysis showed that FedAvg's sample and communication complexities are both O Oe--)). Additionally, SCAFFOLD was proposed in [15.. which utilized variance reduction based local update directions
33 to achieve
the same sample and communication complexities as FedAvg. Similarly, VRL-SGD proposed in 19 also utilized variance reduction and showed improved communication complexity of (Oe=") "). while requiring the same computations as FedAvg. Importantly, both SCAFFOLD and VRL-SGD's guarantees were independent of the data heterogeneity. The FedProx proposed in I0 used penalty based method to improve the communication complexity of FedAvg (i.e., the Parallel Restarted and Momentum SGD [14 [2) to O(--)), FedProx used gradient similarity assumption to model data
heterogeneity which can stringent for many practical applications. This assumption was relaxed
by FedPD proposed in y
Recently, the works MIME algorithm 17 matched the optimal sample complexity (under certain smoothness assump- tions) of (3/2 of the centralized non-convex stochastic optimization algorithms [5| Similarly, Fed-GLOMO [8] achieved the same sample complexity while employing compression further reduce communication. Both MIME and Fed-GLOMO required O -3/2) communication rounds to achieve an c-stationary solution. Please see Table for a summary of the above discussion.
proposed utilize hybrid momentum gradient estimators g The
The comparison of Local SGD (FedAvg) to Minibatch SGD for convex and strongly convex problems with homogeneous data setting was first conducted in [99 and later extended to heterogeneous setting in K It was shown that Minibatch SGD almost always dominates the Local SGD. In contrast, it was shown in (44 that Local SGD dominates Minibatch SGD terms of generalization performance. Although existing FL results are rich, but they somehow ad hoc and there is a lack of principled
3

---

understanding of the algorithms. We note that the proposed STEM algorithmic framework provides theoretical framework that unifies all existing FL results on sample and communication complexities.
Notations. The expected value of a random variable X is denoted by [X[] and its expectation denoted as [ Event A]. We denote by (and IRd) the real line
conditioned on an Event (and the d-dimensional Euclidean space). The set of natural numbers is denoted by N. Given positive integer K N, we denote [K] _ {1, 2,, ..... J. Notation || denotes the fz-norm and (.-) the Euclidean inner product. For a discrete set B, [B] denotes the cardinality of the set. Uniform distribution over discrete set (1,..., } is denoted as
2 Preliminaries
Before we proceed to the algorithms, we make the following assumptions about problem 1 Assumption (Sample Gradient Lipschitz Smoothness). The stochastic functions (( with g(h) ~ (() for all E [K], satisfy the mean squared smoothness property, i.e, we have
V (+) (2; v f )12 < L - y/l² for all x, R4.
Assumption (Unbiased gradient and Variance Bounds). (i) Unbiased Gradient. The stochastic gradients computed at each WN are unbiased
E V y (1) (x:g g( (x), A gl6) (DA), A k E [K].
(ii) Intra- and inter- node Variance Bound. The following bounds hold:
E (x;) (x) v (x)) < 0², (1V) ( (z) (1)
v
c² A g(h) ~ D(A), wk, E [K].
Note that Assumption 1 stronger than directly assuming ((8)ss 's are Lipschitz smooth (which will refer to as the averaged gradient Lipschitz smooth condition), but is still a rather standard assumption in SGD analysis. For example it has been used in analyzing centralized SGD algorithms such as SPIDER [5). SNVRG O STORM 7 (and many others) as well as in FL. algorithms such
as MIME 17 and Fed- heterogeneity, and we call > 0 as the heterogeneity parameter. This is typical assumption required evaluate the performance of FL algorithms. If data distributions across individual WNs identical, i.e., (k)) D for all k, l [K] then have 0.
18. The second relation in Assumption 2 (ii) quantifies the data
Next, we define the e-stationary solution for non-convex optimization problems, as well as quantify the computation and communication complexities to achieve an esstationary point.
Definition 2.1 (e-Stationary Point). A point & is called e-stationary if [=))) < Moreover, stochastic algorithm is said to achieve an e-stationary point in / iterations [ where the expectation is over the stochasticity of the algorithm until time instant t.
Definition 2.2 (Sample complexity). We assume an Incremental First-order Oracle (IFO) frame- work 3 where, given a sample ((t) D(R) at the 1tth node and iterate æ, the oracle returns single IFO operation. We measure the sample (and computational) complexity in terms of the total number of calls to the IFO by all WNs to achieve an e-stationary point given in Definition
(F(R) (z; (k)), V f(k) (z; (.)))). Each access the oracle is counted
Definition 2.3 (Communication complexity). We define communication round as one back-and- forth sharing parameters between the WNs and the SN, Then the communication complexity is defined to be the total number of communication rounds between any WN and the SN required to achieve an e-stationary point given in Definition
The STEM algorithm and the trade-off analysis
In this section, discuss the proposed algorithm and present the main results. The key in the algorithm design is to carefully balance all the four design elements mentioned in Sec. | so that sufficient and useful progress can be made between two rounds of communication.
Let us discuss the key steps of STEM, listed in Algorithm In Step 10, each node locally updates its model parameters using the local direction di, computed by using b stochastic gradients two

---

Algorithm The Stochastic Two-Sided Momemtum (STEM) Algorithm
1: Input: Parameters: 0, the number of local updates I, batch size b, stepsizes {mo
2: Initialize: Iterate 2 = σι = k
EE k=
k (k) descent direction dfa) = d1 = 1 K =k=1 df
with dlt)
Σ 41
( [ :((h) and [B|+] for l E [K].
2 EB!
3: Perform: 22
[k)
z mid A E
(M)
[K]
4: for t = 1 to T'do do
5: for k 1 to K do
#at the WN (1) (8) ((H1
6:
(
Σ
(k)
(-1 (1 (1- (()) Σ
(A) 1111
where we choose [6|4]] b, and ax- = cc
7:
if t mod I = 0 then
#at the SN
d = d++1 = K k d+
(N)
8:
9:
2+2
(k) Ft+1 1/4
1111
K
-==1 3t+1 1441d441
(k)
#server-side momentum #worker-side momentum
(k) (k)
(k)
10:
end 2++2 P(+1
else
1/111041
11:
12: end for
13:
end for
14: Return: Fa where ((1....,.T T}.
consecutive iterates 1241 and After every I local steps, the WNs share their current local models {x(+1 1==1 and directions {d 1K with the SN. The SN aggregates these quantities, and performs a server-side momentum step, before returning To+1 and de+1 to all the WNs. Because both the WNs and the SN perform momentum based updates, we call the algorithm stochastic two-sided momentum algorithm. The key parameters are: b the minibatch size, the local update steps between two communication rounds, 7} the stepsizes, and a, the momentum parameters.
1K
One key technical innovation of our algorithm design is to identify the most suitable way to incorporate momentum based directions in FL algorithms. Although the momentum-base based gradient estimator itself not new and has been used in the literature before (see e.g.. in [ S and [7] to improve the sample complexities of centralized and decentralized stochastic optimization problems, respectively), is by no means clear if and how it can contribute to improve the communication complexity of FL algorithms. We show that in the FL setting, the local directions together with the local models have be aggregated by the SN so to avoid being influenced too much by the local data. More importantly, besides the WNs, the SN also needs to perform updates using the (aggregated) momentum directions. Finally, such two-sided momentum updates have to be done carefully with the correct choice of minibatch size and the number of local updates Overall, it is the judicious choice of all these design elements that results in the optimal sample and communication complexities.
Next, present the convergence guarantees of the STEM algorithm.
3.1 Main results: convergence guarantees for STEM
In this section, we analyze the performance STEM. We first present our main result, and then provide discussions about a few parameter choices. In the next subsection, discuss special case of STEM related to the classical FedAvg and minibatch SGD algorithms.
Theorem 3.1. Under the Assumptions
and suppose the stepsize sequence is chosen as:
171 (w, + 22111/3
(2)
where we define
R (bK)2/302/3
max 202 4.2.4096..3 - -02t 4096 1
L
5

---

Further, let us set c
64L bK
+
24
o²
L2
24(bK)9T and set the initial batch size as
B = bT; set the local updates and minibatch size b as follows:
T= (=0((T/k2)))) b= ((0
(3)
where u satisfies E [0, 1.. Then for STEM the following holds:
(i) For To chosen according to Algorithm 7
we have:
B( = ( {(11)-5 O
2² K2v/3T1-u88
+ +0
(4)
K20/3T1-0/3
K /220/371-0/
(ii) For any [0, 1}, we have Sample Complexity: The sample complexity of STEM is Ō(e-3/2). This implies that each WN requires at most Ō(K-le-3/2) gradient computations. thereby achieving linear speedup with the number of WNs present in the network.
Communication Complexity: The communication complexity of STEM is Ö
The proof of this result is relegated to the Supplemental Material. A few remarks are in order.
Remark (Near- Opttimal sample and communication complexities). Theorem 3.1 suggests that when / and are selected appropriately, then STEM achieves (E and (E sample and communication complexities. Taking them separately, these complexity bounds are the best achievable by the existing FL algorithms (upto logarithmic factors regardless of sample or batch Lipschitz smooth assumption) [35]: see Table We note that the (/23 complexity is the best possible that can be achieved by centralized SGD with the sample Lipschitz gradient assumption; see [5]. On the other hand, the (ee complexity bound is also likely to be the optimal, since in [9] the authors showed that even when the local steps use a class of (deterministic) first-order algorithms, ( is the best achievable communication complexity. The only difference is that [9] does not explicitly assume the inter-node variance bound (i.e., the second relation Assumption (ii). We leave the precise characterization of the communication lower bound with inter-node variance as future work. D
Remark 2 (Large Batch Sizes and/or Local Updates). At first glance, it may seem that the requirement of STEM to compute large mini-batches and/or local updates (cf. Table 1 to achieve this (near) optimal performance is drawback, however, note that it is in fact advantage of STEM that it allows the WNs to perform larger number of local updates (or compute large minibatches) without communicating often. This follows from the fact that irrespective of the number of local updates (or batch sizes) STEM achieves near-optimal communication complexity while attaining optimal overall sample complexity. Moreover, note that even with b / O(1) (i.e., b and are chosen as constants), STEM achieves the same (optimal) sample and communication complexities as achieved by FedGLOMO [8] and MIME [77. We further note that the best of our knowledge the algorithms that achieve the communication complexity of O(E either require the number of local updates or the batch-sizes that depend on the solution accuracy For example, FedProx [T] FedPD (1) and FedDDyn [36] rely on solving the "local problems" to achieve an e-accuracy, which implies that the number of local updates (or the batch sizes) implicitly depends on the desired solution accuracy c, as is the case for STEM. Similarly, as shown in [12 and [4] the communication complexity of FedAvg
and its momentum version can be improved from (cc-2) to O(e 3/2) when the number of local updates (or batch size) is chosen as ( 0(2-1/2) (cf. Section 3.2 for more detailed discussion). O
Remark (The Optimal Batch Sizes and Local Updates Trade-off). The parameter E [0, 1] is used to balance the local minibatch sizes b, and the number of local updates I. Eqs. in B suggest that when increases from 0 to 1, decreases and I increases. Specifically, if v 1, then b is constant but / (T11/3 (O(T1/3/K2// (2/3)_ In this case, each WN chooses small minibatch while executing multiple local updates, and STEM resembles a FedAvg (a.k.a. Local SGD) algorithm but with double-sided momentum update directions, and is referred to as Fed STEM. In contrast, if 0, then b = (T /2/////K)))) /) but / is constant. In this case, each WN chooses a large batch size while executing only a few, or even one, local updates, and STEM resembles the Minibatch SGD, but again with different update directions, and is referred to as Minibatch STEM. Such trade-off can be seen in Fig. [bb] Due to space limitation, these two special cases will be precisely stated in the supplementary materials as corollaries of Theorem
D

---

Algorithm 2 The FedAvg Algorithm
1: Input: {nr] 7, the # of local updates per communication round; b, the minibatch sizes. to do
2:
3:
for 1 K do
d. = 2 (x(
with | b
4:
(8) = z;
(8) n.
5:
" F++
6:
if l mod 1 = 0 then
x 1t R
Σ k (k)
7:
1++1
8:
end if
9:
end for
10:
end for
11:
Return: To where [I,
Remark (The Sub-Optimal Batch Sizes and Local Updates Trade-off). From our proof (Theo-
rem
C.10 included in the supplemental material),
see that STEM requires O max { (bb K E 2/23) samples and Ó( max { (b- (
and communication rounds. According to the above expressions, if b increases beyond O0K--1-111 then the sample com- plexity will increase from the optimal Ō (;32)); otherwise, the optimal sample complexity (/23 is maintained. On the other hand, if b- decreases beyond O(K-1(_112) the communication complexity increases from Ō(e-1). For instance, if we choose b = O(1 and I = O(1) the communi- cation complexity becomes ((e-3/2) while the optimal sample complexity ((c-3/2) is maintained. This trade-off is illustrated in Figure Ta where we maintain the optimal sample complexity, while changing b and I generate the trade-off surface.
e)
4 captures the effect of the heterogeneity of data across WNs, where Ç is the parameter characterizing the intra-node variance and has been defined in Assumption (i)). Highly heterogeneous data with large Ç² can adversely impact the performance of STEM. Note that such dependency on also appears in other existing FL algorithms, such as 9 14 18. However, there is one special case of STEM that does not depend on the parameter Ç. This is the case where = i.e., the minibatch SGD counterpart of STEM where only single local iteration is performed between two communication rounds. We have the following corollary.
Remark (Data Heterogeneity). The term Ö
( K 573 in the gradient bound
Corollary (Minibatch STEM). Under Assumptions ! andd 2 and choose the algorithm parameters
as in Theorem / At each WN, choose I 1, b= (T 1/2 and the initial batch size B = b I.
Then STEM satisfies:
(i) For To chosen according to Algorithm 7 we have
E V =0
cooff(an f((1)) ((( + 3
02
T
T
(ii) Minibatch STEM achieves Ō(E-3/9) sample and Ö(E-1) communication complexity.
Next, show that FedAvg also exhibits trade-off similar to that of STEM but with worse sample and communication complexities.
3.2 Special cases: The FedAvg algorithm
We briefly discuss another interesting special case of STEM, where the local momentum update is replaced by the conventional SGD (i.e., a, 1, (), while the server does not perform the momentum update (i.c.. di 0, t). This is essentially the classical FedAvg algorithm, just that balances the number of local updates and the minibatch size b. We show that this algorithm also exhibits trade-off between b and and on the trade-off curve it achieves O(E -2) sample complexity and -3-33/2 communication complexity.
7

---

Algorithm
Training Acc.
Testing Acc.
Algorithm
Training Acc.
Testing Acc.
FedAvg
78.2
74.1
FedAvg
73.6
75.4
FedProx
79.2
74.8
FedProx
80.0
75.2
FedDyn
68.9
66.0
FedDyn
76.1
71.3
SCAFFOLD
71.9
74.0
SCAFFOLD
72.5
73.7
MIME
82.6
76.8
MIME
61.5
58.6
FedGLOMO
76.1
72.8
FedGLOMO
10.0
10.0
STEM
80.1
78.8
STEM
81.1
78.5
(a) Mild heterogeneity, = 64, and I 7.
(b) Moderate heterogeneity, 8, and / 61.
Table 2: Training and testing accuracy of different algorithms on CIFAR-10 dataset for different batch-sizes, number of local updates, and heteregeneity settings.
Theorem 3.2 (The FedAvg Algorithm). Under Assumptions and
suppose the stepsize is chosen
as:
bK T
Let us set:
T= ((T())O))))) b ((T
(5)
where u E [0, is a constant. Then for FedAvg with 81 81122bb,, the following holds
(i) For To chosen according to Algorithm
we have
f(
+0 Ku/2222/3-116 11)
S² ku/222/33v100
B[V.((=))1== K2/2T2//-u/6
+ +0
(ii) For any choice of E [0,1] we have:
Sample Complexity: The sample complexity of FedAvg is O(e This implies that each WN requires most OK--te- gradient compulations, thereby achieving linear speedup with the number of WNs in the network.
Communication Complexity: The communication complexity of FedAvg is O(E 3/2).
Note that the requirement T being lower bounded is only relevant for theoretical purposes, similar requirement was also imposed in [4] to prove convergence. Again, the parameter v E [0, in the statement of Theorem 3.2 balances / and b each WN while maintaining state-of-the-art sample and communication complexities; please see Table | for comparison of those bounds with existing FedAvg bounds. For 1, FedAvg (cf. Theorem 3.2 reduces FedAvg proposed in and for P 0, the algorithm can be viewed as a large batch FedAvg with constant local updates Note that similar STEM. it known that for
the Minibatch SGD's performance is independent of the heterogeneity parameter, (13). We also point out that if Algorithm [ uses Nesterov's or Polyak's momentum
15 [6
[4] local WNs instead of the recursive momentum estimator
we get the same guarantees as in Theorem 3.2
In summary, this section established that once the WN's and the SN's update directions (SGD in FedAvg and momentum based directions in STEM) are fixed, there exists a sequence of optimal choices the number of local updates I, and the batch sizes b, which guarantees the best possible sample and communication complexities for the particular algorithm. The trade-off analysis presented in this section provides some useful guidelines for how best select i and in practice. Our subsequent numerical results will also verify that if b or I are not chosen judiciously, then the practical performance of the algorithms can degrade significantly.
Numerical results
In this section, we validate the proposed STEM algorithm and compare its performance with the de facto standard FedAvg [I], and the algorithms stated in Table [ Note that instead of FedPD we include the performance comparison with FedDyn [6] since they are known to be very closely

---

Algorithm
Training Acc.
Testing Acc.
Algorithm
Training Acc.
Testing Acc.
FedAvg
57.6
57.1
FedAvg
40.1
39.2
FedProx
59.1
58.5
FedProx
43.5
43.2
FedDyn
51.2
51.3
FedDyn
43.7
43.2
SCAFFOLD
53.1
54.7
SCAFFOLD
40.3
41.3
MIME
56.1
55.1
MIME
32.1
32.1
FedGLOMO
56.8
56.1
FedGLOMO
40.3
40.1
STEM
58.5
57.4
STEM
44.5
43.8
Table 3: Training and testing accuracy on CIFAR-10 dataset for high heterogeneity, b 128 and I 6.
Table 4: Training and testing accuracy on Shakespeare dataset.
* FedAvg, I- 67
Training Loss
1.5
STEM, 67
SCAFFOLD, 1
Testing Accuracy
00.55
FodAvg, 536
0.9
SCAFFOLD, I 536
FedAAv,, I 67
STEM, 1 536
0.85
SCAFFOLD I=
STEM, I 67
0.5
FedAvg, $36
0.8
SCAFFOLD. 11I 536
STEM. 1 536
0.75
0.5
1.5
2.5
2.5
Samples Accessed
x10
Samples Accessed
x10
Figure 2: Training loss and the testing accuracy for classification MNIST data against the number samples accessed at each WN for moderate heterogeneity setting with b 8.
related. The goal of our experiments are three-fold: (1) To show that STEM performs on par, better, compared to other algorithms in different heterogeneity settings, (2) there are multiple ways reach the desired solution accuracy, one can either choose large batch size and perform only few local updates or select smaller batch size and perform multiple local updates, and finally, (3) if the local updates and the batch sizes are not chosen appropriately, the WNs might need to perform excessive computations to achieve the desired solution accuracy, thereby slowing down convergence.
Data and Parameter Settings: We compare the algorithms for image classification tasks CIFAR- 10 and MNIST data sets with 100 WNs, and for next character prediction task on Shakespeare dataset 37 with 143 WNs in the network. For both CIFAR-10 and MNIST, each WN implements two-hidden-layer convolutional neural network (CNN) architecture followed by three linear layers for CIFAR-10 and for MNIST. For CIFAR-10 (and MNIST) datset, consider three settings with mild, moderate and high heterogeneity. For all the three settings, the data is partitioned into disjoint sets among the WNs. the mild heterogeneity setting, the WNs have access to partitioned data from all the classes. In the moderate (resp. high) heterogeneity setting the data is partitioned such that each WN can access data from only 5 (resp. 2) out of 10 classes. For CIFAR-10 (resp. MNIST), each WN has access to 490 (resp. 540) samples for training and 90 (resp. 80) samples for testing purposes.
We also compare the performance of algorithms dataset 377 For this task, we adopt the settings from TC and utilize a 2-Layer LSTT network with 100 hidden units and an 8-D embedding layer each WN. Each WN has access 3616 samples on average, and the samples are randomly split into an 80% training set and a 20% testing set. We randomly sample 10 nodes of 143 for the training purpose. All the experiments are implemented on single NVIDIA Quadro RTX 5000 GPU. More details are provided in appendix.
popular FL benchmarking dataset, Shakespeare
For the proposed STEM algorithm, recall that the step-size 7) = R/(w;, + 21)3/3 with momentum parameter defined a, enf The siep-size is used to update the iterates while the momentum

---

parameter is used to construct the stochastic gradient estimate (cf. Algorithm[ aaad Theorem 3.1 For the experiments. we set Wi o² 1 and c 2/R² and tune for [10-1 I0-2 for the CIFAR- 10 for the Shakespeare dataset. For both the datasets we tune for C in the range [1, 10-. For FedProx [10] and FedDDyn (6 we choose the regularization constant to be 0. The momentum parameters for FedGLOMO 18 and MIME 17e are set based on the choices given in the respective papers. Specifically, for Fe we choose the parameter 8g = 0.2 and design the momentum gradient using damping factor given in Appendix A.4 of FedGLOMO [ Moreover, for MIME we choose the momentum parameter as 0.9. For the rest of the algorithms (including FedAvg and SCAFFOLD), the step-size is tuned from the set {10', 100 10 100 210-?).
dataset and for K E {10*, 10 10-1,10
Discussion: We evaluate the training and testing performance of STEM against multiple algorithms for different heterogeneity settings, minibatch sizes, and number of local updates. In Tables 2a 2b and 3 we compare the training and testing accuracy of STEM to that of other algorithms on the CIFAR-10 dataset. Specifie, heterogeneity settings, the choices of minibatches, and number of local updates are stated along with the tables. Note that STEM performs uniformly well under all the conditions. Moreover, note from Table that FedGLOMO diverges once the number of local updates are high. Also, note from Table 3 that FedProx and STEM adapt well to high heterogeneity. Finally, with the next set of experiments
emphasize the importance of choosing and carefully. In Figure 2 we compare the training and testing performance of STEM, FedAvg and SCAFFOLD, against the number of samples accessed each WN for the classification task on MNIST dataset with moderate heterogeneity. We fix 6 = 8 and conduct experiments under two settings, one with / = 67, and the other with / 536 local updates each WN. Note that although large number of local updates might lead to fewer communication rounds but it can make the sample complexity extremely high as is demonstrated by Figure 2 For example, Figure 2 shows that to reach testing accuracy of 96 97% with / = 67, STEM requires approximately 5000 6000 samples, in contrast with I= 536 it requires more than 25000 samples at each WN. Similar behavior can be observed if we fix T> 1 and increase the local batch sizes. This implies not choosing the local updates and the batch sizes judiciously might lead to increased sample complexity. Additional experiments are included in the supplementary material to further evaluate the performance of the proposed algorithms.
Conclusion
In this work, we proposed a novel algorithm STEM, for distributed stochastic non-convex optimiza- tion with applications to FL. We showed that STEM reaches an c-stationary point with Ō(c-30 sample complexity while achieving linear speed-up with the number of WNs. Moreover, the algo- rithm achieves communication complexity of Ó) We established (optimal) trade-off that allows interpolation between varying choices of local updates and the batch sizes at each WN while maintaining (near optimal) sample and communication complexities. We showed that FedAvg ka. LocalSGD) also exhibits similar trade-off while achieving worse complexities. Our results provide guidelines carefully choose the number of local updates, update directions, and minibatch sizes to achieve the best performance. The future directions of this work include developing lower bounds on communication complexity that establishes the tightness of the analysis conducted in this work.
Acknowledgement
We thank the anonymous reviewers for their valuable comments and suggestions. The work of Prashant Khanduri and Mingyi Hong was supported by NSF grant CMMI-1727757, AFOSR grant 19RT0424 and ARO grant N991NNF-1-1-0047 The work of Mingyi Hong was also supported by an IBM Faculty Research award. The work of Jia Liu has been supported in part by NSF grants CAREER CNS-21 CNS-2112471, CNS-2102233, CCF-21 10252, ECCS-2140277 and Google Faculty Research Award.
10

---

References
[1] J. Konečný, H. B. McMahan, D. Ramage, and P. Richtárik, "Federated optimizaiinn Distributed machine learning for on-device intelligence." arXiv preprint arXiv:1610.02527, 2016.
[2] M. Li, D. G. Andersen, A. J. Smola, and K. Yu, "Communication efficient distributed machine learning with the parameter server," in Advances in Neural Information Processing Systems 27, Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger. Eds. Curran Associates, Inc., 2014, pp. 19-27.
[3] J. Dean, G. Corrado, R. Monga, K. Chen, M. Devin, M. Mao, M. Ranzato, A. Senior, P. Tucker, K. Yang al., "Large scale distributed deep networks," in Advances in neural information processing systems, 2012, pp. 1223-1231.
[4] T. Léauté and B. Faltings, "Protecting privacy through distributed computation in multi-agent decision making," Journal of Artificial Intelligence Research, vol. 47, pp. 649-695, 2013.
[5] C. Fang, C. J. Li, Z. Lin, and TT Zhang, "Spider: Near-optimal non-convex optimization via stochastic path-integrated differential estimator," in Advances in Neural Information Processing Systems, 2018, pp. 689-698
[6] D. Zhou, P. Xu, and Q Gu, "Stochastic nested variance reduction for nonconvex optimization," arXiv preprint arXiv:1806.0000 1806.0881, 2018.
[7] A. Cutkosky and F. Orabona, "Momentum-basel variance reduction in non-convex SGD," in Advances in Neural Information Processing Systems 32. Curran Associates, Inc., 2019, pp. 15 236-15 245.
[8] Q. Tran-Dinh, N. H. Pham, D. T. Phan, and L. M. Nguyen, "Hybrid stochastic gradient descent algorithms for stochastic nonconvex optimization," " arXiv preprint urXiv:1905.05920 2019.
[9] X. Zhang, M. Hong, S. Dhople, W. Yin, and Y. Liu, "Fedpd: A federated learning framework with adaptivity to non-iid data," IEEE Transactions on Signal Processing, pp. 1-1, 2021.
[10] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, "Federated optimization heterogeneous networks," arXiv preprint arXiv:1812.0612 2018.
[11] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, "Communiiation--fficceen
learning of deep networks from decentralized data," in Artificial Intelligence and Statistics. PMLR, 2017, pp. 1273-1282.
[12] H. Yu, Yang, and S. Zhu, "Parallel restarted sgd with faster convergence and less communica- tion: Demystifying why model averaging works for deep learning." in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 33, no. 01, 2019, pp. 5693-57000
[13] B. Woodworth, K. Patel, and N. Srebro, "Minibatch vs local sgd for heterogeneous distributed learning." arXiv preprint
2020.
[14] H. Yu, R. Jin, and S. Yang. "On the linear speedup analysis of communication efficient mo- mentum sgd for distributed non-convex optimization," in International Conference on Machine Learning. PMLR, 2019, pp. 7184-7193.
[15] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, "Scaffold: Stochastic controlled averaging for federated learning," " in International Conference on Machine Learning. PMLR, 2020, pp. 5132-5143.
[16] H. Yang, M. Fang, and J. Liu, "Achieving linear speedup with partial worker participation in non-iid federated learning," " arXiv preprint 1rXiv:2101. 1203, 2021.
[17] S. P. Karimireddy, M. Jaggi, Kale, M. Mohri, S. Reddi, U. Stich, and A. T. Suresh, "Mime: Mimicking centralized stochastic algorithms in federated learning," arXiv preprint arXiv:2008. 03600, 2020.
[18] R. Das, A. Hashemi, S. Sanghavi, and S. Dhillon, "Improved convergence rates for non-convex federated learning with compression," arXiv preprint arXiv:2012.04061, 2020.
11

---

[19] B. Woodworth, K. K. Patel, S. U. Stich, Z. Dai, B. Bullins, H. B. McMahan, O. Shamir, and N. Srebro, "Is local sgd better than minibatch sgd?" arXiv preprint arXiv:2002.07839,, 2020.
[20] H. Yu and R. Jin, "On the computation and communication complexity of parallel sgd with dynamic batch sizes for stochastic non-convex optimization," in International Conference Machine Learning. PMLR, 2019, pp. 7174-7183.
[21] J. Wang and G. Joshi, "Cooperative sgd: A unified framework for the design and analysis of local-update sgd algorithms," Journal of Machine Learning Research, vol. 22, no. 213, pp. 1-50, 2021.
[22] A. Khaled, K. Mishchenko, and P. Richtárik, "Better communication complexity for local sgd," arXiv, 2019.
[23] S. U. Stich, "Local sgd converges fast and communicates little," arXiv preprint 2018.
[24] T. Lin, S. U. Stich, K. K. Patel, and M. Jaggi, "Don't t use large mini-batches, use local sgd," in International Conference on Learning Representations, 2020.
[25] Zhou and G. Cong, "On the convergence properties of a k-step averaging stochastic gra- dient descent algorithm for nonconvex optimization," in Proceedings of the Twenty-Seventh International Joint Conference on Artificial Intelligence, IJCAI-18, 7 2018, pp. 3219-3227.
[26] F. Sattler, S. Wiedemann, K.-R. Müller, and W. Samek, "Robust and communication-effice federated learning from non-iid data," IEEE transactions on neural networks and learning systems, vol. 31, no. 9, pp. 3400-3413, 2019.
[27] Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra, "Federated learning with non-iid data," arXiv preprint arXiv: 00582, 2018.
[28] Wang, V. Tantia, N. Ballas, and M. Rabbat, "Slowmo: Improving communication-effice distributed sgd with slow momentum," arXiv preprint 190.006611 2019.
[29] X. Liang, S. Shen, J. Liu, Z. Pan, E. Chen, and Y. Cheng, "Variance reduced local sgd with lower communication complexity," arXiv preprint arXiv: 1912. 12844, 2019.
[30] Sharma, P. Khanduri, S. Bulusu, K. Rajawat, and P. K. Varshney, "Parallel restarted SPI- DER communication efficient distributed nonconvex optimization with optimal computation complexity," arXiv preprint
2019.
[31] Reddi, S. Kale, and S. Kumar, "On the convergence of adam and beyond," arXiv preprint arXiv:1904.099 09237, 2019.
[32] A. Koloskova, N. Loizou, S. Boreiri, M. Jaggi, and Stich, "A unified theory of decentralized sgd with changing topology and local updates," in International Conference on Machine Learning. PMLR, 2020. pp. 5381-5393.
[33] R. Johnson and Zhang, "Accelerating stochastic gradient descent using predictive variance reduction," in Advances in Neural Information Processing Systems 26. Curran Associates, Inc., 2013, pp. 315-323.
[34] L. Bottou, F. E. Curtis, and J. Nocedal, "Optimization methods for large-scale machine learning." " SIAM Review, vol. 60, no. 2, pp. 223-311, 2018.
[35] Y. Drori and O. Shamir, "The complexity of finding stationary points with stochastic gradient descent," in International Conference on Machine Learning. PMLR, 2020, pp. 2658-26677
[36] D. A. E. Acar, Y. Zhao, R. Matas, M. Mattina, P. Whatmough, and V. Saligrama, "Federated learning based on dynamic regularization," in International Conference on Learning Represen- tations, 2020.
[37] S. Caldas, S. M. K. Duddu, P. Wu, T. J. Konečný, H. B. McMahan, V. Smith, and A. Tal- walkar, "Leeff A benchmark for federated settings," 2019.
12


Reviews:
[{"review_id": "VqLnm3S6lRO", "reviewer": "Reviewer_RPhx", "summary": "The authors presented a momentum extension of the classic FedAvg method for non-convex federated learning, along with a unified framework for convergence analysis. The core result shows that the proposed method can achieve near-tight sample and communication complexities under mild conditions, whilst also revealing a trade-off between local update frequencies and local minibatch sizes. The actual performance of the proposed method was evaluated in CIFAR-10 and MNIST image classification tasks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of using momentum-based variance reduction techniques [7, 8] to accelerate the convergence of FedAvg is very interesting, and novel as far as I know about. The main results make sense to me and appear to show the benefit of momentum acceleration in the regime of federated optimization. The analysis is theoretically sound. I am particularly impressed with the clear and clean writing style, and how all steps in the proofs are carefully elaborated in details. That said, I still have some reservations on this manuscript especially with regard to its technical originality and correctness. \n\nMajor comments: \n\n- Originality: A major concern is about the novelty of the technical part, which follows closely the techniques developed in [7] for momentum SGD with slight adaptation to the map-reduce architecture considered by FL. While fancy and the final results look new in the context of FL, the fundamental contribution of this work is relatively weak in my opinion. As a side contribution, the complexity bounds for FedAvg revisit and slightly generalize the corresponding prior results in literature (see, e.g., [12, 15]). \n\n - Clarity and correctness: The update rule of the average vector $\\bar x_t$ and the relevant cost of communication are so far not clearly presented. Since $\\bar x_t$ will be used in the output step for random selection, it has to be maintained at each time instance $t$. However, in Algorithm 1 & 2, $\\bar x_t$ is only updated in the reduce step after every $I$ rounds of iteration; it is not clear how to update it at other time instances. Around Line #492 of the supplement file, it is written that $\\bar x_t$ is updated by averaging the iterates from the workers at each time stamp. If this is the case (which I believe it is), then a round of communication will be needed to do the averaging at each time instance $t$, which will change a lot the story of the communication complexity analysis! The issue must be fixed to warrant correctness of analysis. \n \n- Parameter setting: According to Theorem 3.1, the learning rate of Algorithm 1 needs to be set with the knowledge of variance $\\sigma$ which is typically unavailable in practice. I am wondering if it is possible to extend the analysis to the adaptive learning rate setup as considered in [7]? \n\nMinor comments: \n\n- Experiments: The method has been evaluated in some benchmark image classification tasks. It will be more interesting to see how the method works on some recent FL benchmarks such as Shakespeare and Sent140. \n\n- Problem formulation:  In Equation (1), it is confusing to denote $f^{(k)}$ simultaneously as population risk and loss function. \n\n- Initialization: the initial minibatch-size $B$ in Algorithm 1 & 2 seems never appeared in the convergence analysis and main results as well. \n\n- Step 9 of Algorithm 1: check the notation $\\bar x^{(k)}_{t+1}$.\n\n- In the proof of Lemma B.1, the second line of the equality (a) implicitly uses the independence among the minibatches $B^{(k)}_t$ so that the expectation operator can be moved into the inner product of stochastic gradient differences from each pair of nodes.  It is suggested to clarify such an independence assumption explicitly in the context.   \n\n- The reference part needs to be carefully improved in formatting. \n\n=== Post rebuttal update ===\n\nThe major concerns raised in my initial review have been much clarified in the author response. I thus would like to change my rating one grid up to 6.  ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors have partially addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "9"}, "extra_text": {}}, {"review_id": "6eYUbW-QGMy", "reviewer": "Reviewer_GjiB", "summary": "Authors propose Stochastic Two-Sided Momentum algorithm, that utilizes certain momentum-assisted stochastic gradient directions for both the WNs and SN updates. Authors show that this method achieves near optimal sample complexity $\\tilde{\\mathcal{O}}\\left(\\epsilon^{-3 / 2}\\right)$ and $\\tilde{\\mathcal{O}}\\left(\\epsilon^{-1}\\right)$ communication complexity.\n\nAlso, authors provide interesting observations about trade-off between the minibatch sizes and the local update frequency.  \n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper describes the Stochastic Two-Sided Momentum algorithm. This paper is well-written and it has a good structure. Authors provide better rates for non-convex case in the Federated Learning setting. They improved both sample and communication complexities. They show the trade-off between local computations and mini-batch size. The 3D plot is very illustrative in this case. These results are novel and relevant to the Federated Learning community. \n\nTheoretical analysis is not tight in terms of constants. For example 4096 in theorem 3.1 does not look to be the best constant. The assumptions that are used in the analysis are quite strong. In paper (https://arxiv.org/pdf/2012.04061.pdf) analysis does not need any additional assumptions. Another problem of this analysis is that we need to have a big mini-batch size (\\tilde{O}(\\varepsilon^{1/2}) or a big number of local computations (\\tilde{O}(\\varepsilon^{1/2}), but SCAFFOLD does not require this restriction. Also authors do not compare their result with this result in these papers (https://arxiv.org/pdf/2102.07845.pdf) and (https://arxiv.org/pdf/2007.01154.pdf)\n\nAuthors do not consider compression approach to reduce communication complexity. \n\nThe most problematic part is experiments. It is not clear to me how hyperparameters were tuned. Especially, I do not understand how I and b were picked. Authors provide plots only for b = 8, I = 61 and b = 8, I=61. For the last experiment, we have that I = 536 and I = 67. I do not understand why authors use these parameters and do not try others. Also authors compare only three methods: STEM, SCAFFOLD and FedAvg. They do not compare new method with FedPD, FedProx, MIME, FedGLOMO and MARINA.\n\n==========================================================================================================\nAfter authors response I have increased my score. \n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Authors can add more explanations of experimental settings and process of tuning parameters. Additional experiments with different datasets and level of heterogeneity is desirable. To have fair competition authors can express complexity with mini-batch size parameter. \n\nIf authors can resolve issues with experiments the score can be increased. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "3pFBgRGdE1V", "reviewer": "Reviewer_siNV", "summary": "\nThe paper proposes STEM for federated learning under *sample gradient smoothness condition*. By using momentum-based techniques, the proposed STEM algorithm achieves the optimal sample complexity of $O(\\epsilon^{-3/2})$ and communication complexity of $O(\\epsilon^{-1})$ for non-convex functions. The paper also shows how to balance the batch size and the number of local steps. ", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Main:\n\n## Originality:\n\nThe momentum-based variance reduction techniques are known. Nevertheless, it is new and interesting to see how to combine them into Federated Learning. The paper clearly discussed the differences and comparisons with the previous works. \n\n\n## Quality: \n\nThe paper provides theoretical guarantees for the propsed STEM algorithm. Specifically, the paper gives a clear description of the algorithm and indicates the choices of hyper-parameters (number of local updates $I$ and batch size $b$). Then the paper provides convergence guarantees under *sample gradient smoothness condition* for non-convex functions. Next, the paper shows how to simutaneously achieve optimal sample and communication complexity with the optimal batch sizes and local updates. \n\nFor comparison with the baseline algorithms, the paper also provides theoretical guarantees for FedAvg and gives the batch sizes and local updates trade-off. In the numerical experiment part, the paper shows their proposed STEM algorithm is better than FedAvg and SCAFFOLD.\n\n## Clarity: \n\nThis paper is clearly written, with careful comparisons to previous works, clearly stated assumptions and complexity measures, rigorous expression of the theorems, and detailed remarks on their algorithms. \n\n## Significance: \n\nThe paper is of theoretical significance. An appropriate complexity measure is the key to study of the federated learning. The paper studies both the stochastic gradient complexity (sample complexity) as well as the communication complexity for FedAvg and their proposed STEM. It provides useful message to the community.\n\nThe paper proposes STEM for federated learning under *sample gradient smoothness condition* and provides instructions on hyper-parameter choices. In reality, the objective function may satisfy this stronger smoothness assumption and in this case the proposed STEM is useful, with only a small change to the implementation of the baseline algorithms.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The presented algorithm appears to achieve better communication complexities in theory and has reasonable experimental performance. Reviewers are all in favor of accepting the paper.

Decision: Accept (Poster)