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 sarrd00950mm....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. 1angg55922bbuukyymalll.. osu. edu
Ketan Rajawat Indian Institute of Technology Kanpur ketan@iitk ac. in
Pramod K. Varshney
Syracuse University
varshneydsyr 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
min ER
f(x) = 1 Σ K
(z)
= K Σ B (f(0) (x: (x; (0))
(I)
where f(A) ; _d RR denotes the smooth (possibly non-convex) objective function and g(b) ~ D(k) 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
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 ' 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= ((T
(3)
where u satisfies E [0, 11. 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, O 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/ /) 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- 1
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 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 т) = 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 O) 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-699
[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": "_l3KnJHq0FW", "reviewer": "Reviewer_pfmU", "summary": "This paper introduced a two-step momentum method for federated learning. The proposed method achieves the best achievable sample and communication complexities for first-order stochastic FL algorithms. The theoretical analysis provides guidance on selecting synchronization interval and batch size for achieving the nearly optimal complexities.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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 introduced an interesting idea for achieving near-optimal sample and communication complexities for federated learning. Overall, the paper is well written, but I would suggest authors explicitly state Definitions 2.1 and 2.2 before referencing them. More detailed comments are shown as follows\n\n1. Why do we need to use the consecutive iterates x^k_{t+1} and x^k_t to evaluate the estimator? Can't we just use x^k_{t+1}? How does it affect the complexity bounds?\n\n2. As pointed out by (Khaled, Ahmed, Konstantin Mishchenko, and Peter Richtárik, 2019), a bound on the dissimilarity is not realistic for heterogeneous data setting. Under this assumption, the data is very close to i.i.d.\n\n3. The experiments seems weak in terms of model selection and data scale. Federated EMNIST dataset is a better benchmark data for federated learning (Caldas, Sebastian, et al. 2018).\n\nKhaled, Ahmed, Konstantin Mishchenko, and Peter Richtárik. \"First analysis of local gd on heterogeneous data.\" arXiv preprint arXiv:1909.04715 (2019).\nCaldas, Sebastian, et al. \"Leaf: A benchmark for federated settings.\" arXiv preprint arXiv:1812.01097 (2018).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "No", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "_d5ubtMhk-B", "reviewer": "Reviewer_2QZ9", "summary": "This paper studies non-convex stochastic optimization for federated learning under the mean-squared smoothness assumption and shows that better sample convergence and communication complexity are possible with large batch sizes. ", "questions": "", "limitations": "", "rating": 5, "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": "To me, this is a borderline paper. This paper has its technical merit by establishing stronger sample/communication complexity bounds. However, there seem to be some fundamental issues on the interpretation correctness and presentation clarity.  \n\nThe proposed Algorithm 1 is named \" Stochastic Two-Sided Momemtum (STEM)\", which also appears in the paper title. It is named so because the authors interpret line 9 in Algorithm 1 as \"server side momentum step\" and argue its novelty and importance for better convergence. However, I don't think it is necessary to let the server do the momentum SGD update in line 9.  We could let each client do local momentum SGD updates in line 8 and then simply average all local x at line 9.  With such an equivalent description,there is no \"server side momentum\" concept needed. It is really hard for me to appreciate the necessity of server side momentum. In fact, I think line 7-10 is mathematically equivalent to \" updating each local node with momentum SGD at each iteration and average all local x at the server when t mod I = 0\", which is just the standard \"FedAvg with momentum SGD local solver \"or \"local SGD with momentum\".  In this regard, it seems this paper does not propose a new algorithm but just analyzes the existing FedAvg with momentum SGD local solvers under the stronger Assumption 1.  I would also suggest the authors revisit their proof and comment if the proof also follows similar lines as in early works. \n\nTable 1 lists the complexity of earlier works that can be misleading to general readers who do not have sufficient optimization backgrounds. The current comparison is not apple to apple since the complexities are attained under different technical assumptions. Please add an extra column to list the needed assumption for each row. \n\nThe idea of using large batch sizes to achieve faster convergence may let practitioners frown since large batch sizes may hurt the final model accuracy. In the experiment, the methods are only tested with b=8 or b=64. (These settings are usually considered as O(1) batch size.) More experiments with larger b might be desired.\n\n\n\n\n\n", "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": "3"}, "extra_text": {}}, {"review_id": "LyeNkLIrksD", "reviewer": "Reviewer_4cGF", "summary": "The paper proposes a new method to solve the finite sum problem in federated learning where each local model is an expectation function. The STEM algorithm uses momentum variance-reduced estimator in local updates (similar to STORM estimator in stochastic optimization) and also send these estimators to the server when performing server-side update. Therefore, STEM includes momentum updates for both local and server updates. Convergence analysis reveals that STEM nearly match the best known sample complexity in centralized setting and it also matches the lower bound communication complexity up to a log factor. Numerical experiments on neural network training not only show advantage of STEM compared with FedAvg and SCAFFOLD but also presenting the tradeoff between communication frequency and minibatch size.", "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": "Pros:\n- The paper introduces a new framework to incorporate the momentum gradient estimators in federated learing compared to 2 related works which are MIME and FedGLOMO. Given FedGLOMO is rather new, I consider it as concurrent work.\n- Convergence analysis provides intuition on the tradeoff between communication and sample complexities (perform more local updates vs communicate more). This effect is further illustrated in the experiments\n- Complexity results are also new.\n- The analysis of number of local updates and minibatch size\n\nCons:\n- The paper presents FedAvg with full-device participation and compare the results with this version. However, FedAvg or SCAFFOLD do allow partial-participation in their convergence analysis while STEM seems to require full device participation.\n- Apart from the local model, STEM also sends \"local momentum\" which is an estimate of the local gradient. This framework does not exactly match the federated learning framework described in [https://arxiv.org/pdf/1610.05492.pdf, page 2] where local nodes can only send local model back to the server. Moreover, these gradient estimates are susceptible to privacy attack in federated learning [https://proceedings.neurips.cc/paper/2020/file/c4ede56bbd98819ae6112b20ac6bf145-Paper.pdf]. Although this is not exactly the true gradient, but it is still a good estimate (if not then the algorithm is not working well) so it is good to avoid communicating the first-order information. I am aware that the same criticism can be applied to SCAFFOLD, MIME or FedGLOMO.\n\nMinor comments:\n- It is better to elaborate the steps from (b) to (c) for the second term at line 504 similar to (9) at line 523 for clarity.\n- There should be a $\\nabla f^{(k)}(x^{(k)}_t;\\xi_t^{(k)})$ at line 561.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Regarding limitations, I have not seen the authors state this clearly in the paper. I don't see any immediate negative societal impact of this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "0iAs4eCS3la", "reviewer": "Reviewer_9RXX", "summary": "This paper proposes and analyzes an optimization algorithm (STEM) with two-sided momentum and local steps for (non-convex) Federated Learning. The authors prove its convergence and \nThey also claim that it improves over the previous methods in terms of join communication and sample complexity.", "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 work considers a problem of non-convex federated optimization, where communication is the bottleneck. A local stochastic gradient method with momentum at the server and workers is suggested. STEM enjoys $\\tilde{\\mathcal{O}}(\\epsilon^{-3/2})$ sample and $\\tilde{\\mathcal{O}}(\\epsilon^{-1})$ communication complexity if mini-batch size and number of local updates are chosen appropriately (e.g. as $\\mathcal{O}(\\epsilon^{-1/2})$). The authors show a trade-off between the mini-batch sizes and local update frequency.\nExperimental comparison to FedAvg and SCAFFOLD on two image classification tasks (with convolutional neural networks) confirmed the benefits of the proposed method.\n\nRelated work was mostly adequately cited and compared with.\n\nOverall, the paper is well-organized and leaves an impression of good work, but I think that in its current stage, it is not ready for publication. \nThe first reason is the theoretical contributions. The method does not suggest a very novel technique because algorithms with momentum and local steps were previously analyzed in the federated setting. \nThe so-called “near-optimal” sample and communication complexities are obtained for mini-batches and number of local updates of the order of $\\mathcal{O}(\\epsilon^{-1/2})$, which means that for achieving an $\\epsilon$-stationary point for $\\epsilon = 10^{-10}$ (which corresponds to gradient norm $\\||\\nabla f(x_t)\\|| \\leq 10^{-5}$) the method requires very big mini-batches or the number of local updates, which seems not very practical. At the same time, for methods like SCAFFOLD and MIME, it is not required. That is why I think that theoretical comparison is not completely fair from the practical side.\n\nThe second issue is the lack of experimental comparison to the methods presented in the comparison Table 1. The chosen benchmarking methods (FedAvg and SCAFFOLD) are the easiest competitors to STEM because the first one is the most basic baseline, while the second is more suited for convex problems in the cross-silo settings. Both of them also have worse theoretical complexities. To show real practical benefits of the proposed method, it is required to compare against FedPD and MIME/FedGLOMO as they have the same sample or communication complexities to STEM. Furthermore, the hyperparameter tunning seems very limited (just two values for $\\overline{\\kappa}$ and $\\overline{c}$). I also recommend multiple runs for every experiment to show mean values and deviations as the methods are stochastic. This is a basic standard to show the consistent superiority of the proposed approach over the baselines.\n\nIn addition, real experiments supporting the point on the trade-off between mini-batch size and the number of local updates are quite limited to basically a couple of values (of $b$ and $I$), which does not seem enough to show this effect and provide valuable guidelines for practitioners. If the computational resources are the problem, I suggest the authors using a simpler non-convex model [2]. Showing plots with gradient norm values would be also more consistent with the theory in this case.\n\nTo conclude, I think that the current results are not very important nor valuable for the field.\n\nUnfortunately, chosen notation (e.g., for stochastic gradients) complicates the algorithm description and makes convergence proofs very hard to follow and understand. I greatly encourage authors to simplify it. E.g. for local gradient estimates $\\nabla f^{kj}$ could be used in case of finite sum losses $f^{(k)} = \\sum_{j=1}^J f^{kj}$.\n\n*Minor issues/comments:*\n\nI think that the authors missed an important reference [1].\n\nThere is a typo in the Appendix on page 18 in equation (8). In the RHS of the (b) inequality, there should be superscript (j) after the last $\\nabla$ symbol.\n\n[1] R. Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020)\n\n[2] L. Zhao, et al. \"From Convex to Nonconvex: A Loss Function Analysis for Binary Classification,\" 2010 IEEE International Conference on Data Mining Workshops, 2010, pp. 1281-1288, doi: 10.1109/ICDMW.2010.57", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations of the work were mostly addressed. Negative societal impact concerns are not applicable due to the theoretical nature of the paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a borderline paper. The reviewers agree that the communication complexity and sample complexity results are new in the federated learning setting, but there are questions on the novelty of the algorithm as the main ideas and techniques (including the complxities) are adapted from existing variance-reduced estimators for stochastic optimization. The are also concerns on additional vectors to communicate, the large local batch-size or number of iterations, and the requirement of full participation. After discussion between the reviewers and assessment of two Area Chairs and senior AC, we think the results are not significant enough and have to recommend rejection. 
Decision: Accept (Poster)