STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning
Year: 2021

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 number of local updates, 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 Õ(✏ 3/2) samples and Õ(✏ 1) communication rounds to compute an ✏-stationary solution. To the best of our knowledge, this is the first FL algorithm that achieves such near-optimal sample and communication complexities simultaneously. Further, we show that there is a trade-off curve between the number of local updates and the 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 number of local updates, WNs’ and server’s update directions, and minibatch sizes to achieve the best performance.

heading: 1 Introduction
text: In Federated Learning (FL), multiple worker nodes (WNs) collaborate with the goal of learning a joint model, by only using local data. Therefore it has become popular for machine learning problems where datasets are massively distributed [1]. In FL, the data is often collected at or off-loaded to multiple WNs which in collaboration with a server node (SN) jointly aim to learn a centralized model
⇤Corresponding Author: Mingyi Hong.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
[2, 3]. The local WNs share the computational load and since the data is local to each WN, FL also provides some level of data privacy [4]. A classical distributed optimization problem that K WNs aim to solve:
min x2Rd
⇢ f(x) := 1
K
KX
k=1
f (k)(x) := 1
K
KX
k=1
E⇠(k)⇠D(k) [f (k)(x; ⇠(k))] . (1)
where f (k) : Rd ! R denotes the smooth (possibly non-convex) objective function and ⇠(k) ⇠ D(k) represents the sample/s drawn from distribution D(k) at the kth WN with k 2 [K]. When the distributions D(k) 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 to achieve an ✏-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 to 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 significant 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-assisted 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 Õ(✏ 3/2)2 samples and Õ(✏ 1) communication rounds to reach an ✏-stationary solution; see Figure 1 for an illustration. These complexity results are the best achievable for first-order stochastic FL algorithms (under certain assumptions, cf. Assumption 1); see [5–8] and [9, 10], as well as Remark 1 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.k.a. the Local SGD) [11–13]. 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 O(✏ 2) samples and O(✏ 3/2) communication rounds to achieve an ✏-stationary solution.
2The notation Õ(·) hides the logarithmic factors.
Collectively, our insights on the trade-offs provide practical guidelines for choosing different design elements for FL algorithms.
Related Works. FL algorithms were first proposed in the form of FedAvg [11], 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 [19–25], while many recent studies have focused on designing new algorithms to deal with heterogeneous data settings, as well as problems where the local loss functions are non-convex [9, 10, 12–16, 18, 26–32]. In [12], the authors showed that Parallel Restarted SGD (Local SGD or FedAvg [11]) achieves linear speed up while requiring O(✏ 2) samples and O(✏ 3/2) rounds of communication to reach an ✏-stationary solution. In [14], a Momentum SGD was 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(✏ 1) while maintaining the same sample complexity. The works in [15, 16] conducted tighter analysis for FedAvg with partial WN participation with O(1) local updates and batch sizes. Their analysis showed that FedAvg’s sample and communication complexities are both O(✏ 2). 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 [29] also utilized variance reduction and showed improved communication complexity of O(✏ 1), 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 [10] used a penalty based method to improve the communication complexity of FedAvg (i.e., the Parallel Restarted and Momentum SGD [14, 12]) to O(✏ 1). FedProx used a gradient similarity assumption to model data heterogeneity which can be stringent for many practical applications. This assumption was relaxed by FedPD proposed in [9].
Recently, the works [17, 18] proposed to utilize hybrid momentum gradient estimators [7, 8]. The MIME algorithm [17] matched the optimal sample complexity (under certain smoothness assumptions) of O(✏ 3/2) of the centralized non-convex stochastic optimization algorithms [5–8]. Similarly, Fed-GLOMO [18] achieved the same sample complexity while employing compression to further reduce communication. Both MIME and Fed-GLOMO required O(✏ 3/2) communication rounds to achieve an ✏-stationary solution. Please see Table 1 for a summary of the above discussion.
The comparison of Local SGD (FedAvg) to Minibatch SGD for convex and strongly convex problems with homogeneous data setting was first conducted in [19] and later extended to heterogeneous setting in [13]. It was shown that Minibatch SGD almost always dominates the Local SGD. In contrast, it was shown in [24] that Local SGD dominates Minibatch SGD in terms of generalization performance. Although existing FL results are rich, but they are somehow ad hoc and there is a lack of principled
understanding of the algorithms. We note that the proposed STEM algorithmic framework provides a 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 E[X] and its expectation conditioned on an Event A is denoted as E[X|Event A]. We denote by R (and Rd) the real line (and the d-dimensional Euclidean space). The set of natural numbers is denoted by N. Given a positive integer K 2 N, we denote [K] , {1, 2, . . . ,K}. Notation k · k denotes the `2-norm and h·, ·i the Euclidean inner product. For a discrete set B, |B| denotes the cardinality of the set. Uniform distribution over a discrete set {1, . . . , T} is denoted as U{1, . . . , T}.

heading: 2 Preliminaries
text: Before we proceed to the algorithms, we make the following assumptions about problem (1). Assumption 1 (Sample Gradient Lipschitz Smoothness). The stochastic functions f (k)(·, ⇠(k)) with ⇠(k) ⇠ D(k) for all k 2 [K], satisfy the mean squared smoothness property, i.e, we have
Ekrf (k)(x; ⇠(k)) rf (k)(y; ⇠(k))k2  L2kx yk2 for all x, y 2 Rd. Assumption 2 (Unbiased gradient and Variance Bounds). (i) Unbiased Gradient. The stochastic gradients computed at each WN are unbiased
E[rf (k)(x; ⇠(k))] = rf (k)(x), 8 ⇠(k) ⇠ D(k), 8 k 2 [K]. (ii) Intra- and inter- node Variance Bound. The following bounds hold:
Ekrf (k)(x; ⇠(k)) rf (k)(x)k2  2, krf (k)(x) rf (`)(x)k2  ⇣2, 8 ⇠(k) ⇠ D(k), 8k, ` 2 [K].
Note that Assumption 1 is stronger than directly assuming f (k)’s are Lipschitz smooth (which we will refer to as the averaged gradient Lipschitz smooth condition), but it 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 [6], STORM [7] (and many others) as well as in FL algorithms such as MIME [17] and Fed-GLOMO [18]. The second relation in Assumption 2-(ii) quantifies the data heterogeneity, and we call ⇣ > 0 as the heterogeneity parameter. This is a typical assumption required to evaluate the performance of FL algorithms. If data distributions across individual WNs are identical, i.e., D(k) = D(`) for all k, ` 2 [K] then we have ⇣ = 0.
Next, we define the ✏-stationary solution for non-convex optimization problems, as well as quantify the computation and communication complexities to achieve an ✏-stationary point. Definition 2.1 (✏-Stationary Point). A point x is called ✏-stationary if krf(x)k2  ✏. Moreover, a stochastic algorithm is said to achieve an ✏-stationary point in t iterations if E[krf(xt)k2]  ✏, 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) framework [34], where, given a sample ⇠(k) ⇠ D(k) at the kth node and iterate x, the oracle returns (f (k)(x; ⇠(k)),rf (k)(x; ⇠(k))). Each access to the oracle is counted as a 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 ✏-stationary point given in Definition 2.1. Definition 2.3 (Communication complexity). We define a communication round as a one back-andforth sharing of 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 ✏-stationary point given in Definition 2.1.
3 The STEM algorithm and the trade-off analysis
In this section, we 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. 1, 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 1. In Step 10, each node locally updates its model parameters using the local direction dkt , computed by using b stochastic gradients at two
Algorithm 1 The Stochastic Two-Sided Momemtum (STEM) Algorithm 1: Input: Parameters: c > 0, the number of local updates I , batch size b, stepsizes {⌘t}. 2: Initialize: Iterate x(k)1 = x̄1 = 1K PK k=1 x (k) 1 , descent direction d (k) 1 = d̄1 = 1 K PK k=1 d (k) 1
with d(k)1 = 1 B P ⇠(k)1 2B (k) 1 rf (k)(x(k)1 ; ⇠ (k) 1 ) and |B (k) 1 | = B for k 2 [K].
3: Perform: x(k)2 = xk1 ⌘1d (k) 1 , 8 k 2 [K] 4: for t = 1 to T do 5: for k = 1 to K do #at the WN 6: d(k)t+1 = 1
b
X
⇠(k)t+12B (k) t+1
rf (k)(x(k)t+1; ⇠ (k) t+1)+(1 at+1) ⇣ d(k)t 1
b
X
⇠(k)t+12B (k) t+1
rf (k)(x(k)t ; ⇠ (k) t+1)
⌘
where we choose |B(k)t+1| = b, and at+1 = c · ⌘2t ; 7: if t mod I = 0 then #at the SN 8: d(k)t+1 = d̄t+1 := 1 K PK k=1 d (k) t+1 9: x(k)t+2 := x̄t+1 ⌘t+1d̄t+1 = 1 K PK k=1 x (k) t+1 ⌘t+1d̄t+1 #server-side momentum
10: else x(k)t+2 = x (k) t+1 ⌘t+1d (k) t+1 #worker-side momentum 11: end if 12: end for 13: end for 14: Return: x̄a where a ⇠ U{1, ..., T}.
consecutive iterates x(k)t+1 and x (k) t . After every I local steps, the WNs share their current local models {x(k)t+1} K k=1 and directions {d (k) t+1} K k=1 with the SN. The SN aggregates these quantities, and performs a server-side momentum step, before returning x̄t+1 and d̄t+1 to all the WNs. Because both the WNs and the SN perform momentum based updates, we call the algorithm a stochastic two-sided momentum algorithm. The key parameters are: b the minibatch size, I the local update steps between two communication rounds, ⌘t the stepsizes, and at the momentum parameters.
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-based gradient estimator itself is not new and has been used in the literature before (see e.g., in [7, 8] and [17, 18] to improve the sample complexities of centralized and decentralized stochastic optimization problems, respectively), it 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 to 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 b, and the number of local updates I . Overall, it is the judicious choice of all these design elements that results in the optimal sample and communication complexities.
Next, we present the convergence guarantees of the STEM algorithm.
3.1 Main results: convergence guarantees for STEM
In this section, we analyze the performance of STEM. We first present our main result, and then provide discussions about a few parameter choices. In the next subsection, we discuss a special case of STEM related to the classical FedAvg and minibatch SGD algorithms. Theorem 3.1. Under the Assumptions 1 and 2, suppose the stepsize sequence is chosen as:
⌘t = ̄
(wt + 2t)1/3 , (2)
where we define :
̄ = (bK)2/3 2/3
L , wt = max
⇢ 2 2, 4096L3I3̄3 2t, c3̄3
4096L3I3
.
Further, let us set c = 64L 2 bK + 2 24̄3LI = L 2 ✓ 64 bK + 1 24(bK)2I ◆ , and set the initial batch size as
B = bI; set the local updates I and minibatch size b as follows:
I = O (T/K2)⌫/3 , b = O (T/K2)1/2 ⌫/2
(3)
where ⌫ satisfies ⌫ 2 [0, 1]. Then for STEM the following holds:
(i) For x̄a chosen according to Algorithm 1, we have:
Ekrf(x̄a)k2 = O ✓ f(x̄1) f⇤
K2⌫/3T 1 ⌫/3
◆ + Õ ✓ 2
K2⌫/3T 1 ⌫/3
◆ + Õ ✓ ⇣2
K2⌫/3T 1 ⌫/3
◆ . (4)
(ii) For any ⌫ 2 [0, 1], we have Sample Complexity: The sample complexity of STEM is Õ(✏ 3/2). This implies that each WN requires at most Õ(K 1✏ 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 Õ(✏ 1).
The proof of this result is relegated to the Supplemental Material. A few remarks are in order. Remark 1 (Near-Optimal sample and communication complexities). Theorem 3.1 suggests that when I and b are selected appropriately, then STEM achieves Õ(✏ 3/2) and Õ(✏ 1) 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 1. We note that the O(✏ 3/2) 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 O(✏ 1) 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(✏ 1) 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 in Assumption 2-(ii)). We leave the precise characterization of the communication lower bound with inter-node variance as future work. 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 a drawback, however, we note that it is in fact an 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 = I = O(1) (i.e., b and I are chosen as constants), STEM achieves the same (optimal) sample and communication complexities as achieved by FedGLOMO [18] and MIME [17]. We further note that to the best of our knowledge the algorithms that achieve the communication complexity of O(✏ 1) either require the number of local updates or the batch-sizes that depend on the solution accuracy ✏. For example, FedProx [10], FedPD [9], and FedDyn [36] rely on solving the “local problems" to achieve an ✏-accuracy, which implies that the number of local updates (or the batch sizes) implicitly depends on the desired solution accuracy ✏, as is the case for STEM. Similarly, as shown in [12] and [14] the communication complexity of FedAvg and its momentum version can be improved from O(✏ 2) to O(✏ 3/2) when the number of local updates (or batch size) is chosen as O(✏ 1/2) (cf. Section 3.2 for a more detailed discussion). Remark 3 (The Optimal Batch Sizes and Local Updates Trade-off). The parameter ⌫ 2 [0, 1] is used to balance the local minibatch sizes b, and the number of local updates I . Eqs. in (3) suggest that when ⌫ increases from 0 to 1, b decreases and I increases. Specifically, if ⌫ = 1, then b is a constant but I = O(T 1/3/K2/3). In this case, each WN chooses a 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 = O(T 1/2/K) but I is a 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 a trade-off can be seen in Fig. 1b. Due to space limitation, these two special cases will be precisely stated in the supplementary materials as corollaries of Theorem 3.1.
Algorithm 2 The FedAvg Algorithm 1: Input: {⌘t}Tt=0; I , the # of local updates per communication round; b, the minibatch sizes. 2: for t = 1 to T do 3: for k = 1 to K do 4: d(k)t = 1 b P ⇠(k)t 2B (k) t rf (k)(x(k)t ; ⇠ (k) t ) with |B (k) t | = b
5: x(k)t+1 = x (k) t ⌘td (k) t 6: if t mod I = 0 then 7: x(k)t+1 = x̄t+1 = 1 K PK k=1 x (k) t+1 8: end if 9: end for
10: end for 11: Return: x̄a where a ⇠ U{1, ..., T}.
Remark 4 (The Sub-Optimal Batch Sizes and Local Updates Trade-off). From our proof (Theorem C.10 included in the supplemental material), we can see that STEM requires Õ max (b ·
I)✏ 1,K 1✏ 3/2 samples and Õ max ✏ 1, (b · I) 1K 1✏ 3/2 and communication rounds.
According to the above expressions, if b · I increases beyond O(K 1✏ 1/2), then the sample complexity will increase from the optimal Õ(✏ 3/2); otherwise, the optimal sample complexity Õ(✏ 3/2) is maintained. On the other hand, if b · I decreases beyond O(K 1✏ 1/2), the communication complexity increases from Õ(✏ 1). For instance, if we choose b = O(1) and I = O(1) the communication complexity becomes Õ(✏ 3/2) while the optimal sample complexity Õ(✏ 3/2) is maintained. This trade-off is illustrated in Figure 1a, where we maintain the optimal sample complexity, while changing b and I to generate the trade-off surface. Remark 5 (Data Heterogeneity). The term Õ ⇣ ⇣2
K2⌫/3T 1 ⌫/3
⌘ in the gradient bound (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 2-(ii). Highly heterogeneous data with large ⇣2 can adversely impact the performance of STEM. Note that such a 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 = 1, i.e., the minibatch SGD counterpart of STEM where only a single local iteration is performed between two communication rounds. We have the following corollary.
Corollary 1 (Minibatch STEM). Under Assumptions 1 and 2 , and choose the algorithm parameters as in Theorem 3.1. At each WN, choose I = 1, b = (T/K2)1/2, and the initial batch size B = b · I . Then STEM satisfies:
(i) For x̄a chosen according to Algorithm 1, we have
Ekrf(x̄a)k2 = O ⇣f(x̄1) f⇤
T
⌘ + Õ ⇣ 2
T
⌘ .
(ii) Minibatch STEM achieves Õ(✏ 3/2) sample and Õ(✏ 1) communication complexity.
Next, we show that FedAvg also exhibits a trade-off similar to that of STEM but with worse sample and communication complexities.

heading: 3.2 Special cases: The FedAvg algorithm
text: We briefly discuss another interesting special case of STEM, where the local momentum update is replaced by the conventional SGD (i.e., at = 1, 8 t), while the server does not perform the momentum update (i.e., d̄t = 0, 8 t). This is essentially the classical FedAvg algorithm, just that it balances the number of local updates I and the minibatch size b. We show that this algorithm also exhibits a trade-off between b and I and on the trade-off curve it achieves O(✏ 2) sample complexity and O(✏ 3/2) communication complexity.
Theorem 3.2 (The FedAvg Algorithm). Under Assumptions 1 and 2, suppose the stepsize is chosen as: ⌘ = q bK T ; Let us set:
I = O (T/K3)⌫/4 , b = O (T/K3)1/3 ⌫/3
(5)
where ⌫ 2 [0, 1] is a constant. Then for FedAvg with T 81L2I2bK, the following holds
(i) For x̄a chosen according to Algorithm 2, we have
Ekrf(x̄a)k2 = O ✓ f(x̄1) f⇤
K⌫/2T 2/3 ⌫/6
◆ +O ✓ 2
K⌫/2T 2/3 ⌫/6
◆ +O ✓ ⇣2
K⌫/2T 2/3 ⌫/6
◆ .
(ii) For any choice of ⌫ 2 [0, 1] we have: Sample Complexity: The sample complexity of FedAvg is O(✏ 2). This implies that each WN requires at most O(K 1✏ 2) gradient computations, thereby achieving linear speedup with the number of WNs in the network. Communication Complexity: The communication complexity of FedAvg is O(✏ 3/2).
Note that the requirement on T being lower bounded is only relevant for theoretical purposes, a similar requirement was also imposed in [14] to prove convergence. Again, the parameter ⌫ 2 [0, 1] in the statement of Theorem 3.2 balances I and b at each WN while maintaining state-of-the-art sample and communication complexities; please see Table 1 for a comparison of those bounds with existing FedAvg bounds. For ⌫ = 1, FedAvg (cf. Theorem 3.2) reduces to FedAvg proposed in [12, 14] and for ⌫ = 0, the algorithm can be viewed as a large batch FedAvg with constant local updates [15, 16]. Note that similar to STEM, it is known that for I = 1, the Minibatch SGD’s performance is independent of the heterogeneity parameter, ⇣ [13]. We also point out that if Algorithm 1 uses Nesterov’s or Polyak’s momentum [14] at 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 of 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 to best select b and I 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.

heading: 4 Numerical results
text: In this section, we validate the proposed STEM algorithm and compare its performance with the de facto standard FedAvg [11], and the algorithms stated in Table 1. Note that instead of FedPD we include the performance comparison with FedDyn [36] since they are known to be very closely
Algorithm Training Acc. Testing Acc. FedAvg 40.1 39.2 FedProx 43.5 43.2 FedDyn 43.7 43.2 SCAFFOLD 40.3 41.3 MIME 32.1 32.1 FedGLOMO 40.3 40.1 STEM 44.5 43.8
Table 4: Training and testing accuracy on Shakespeare dataset.
related. The goal of our experiments are three-fold: (1) To show that STEM performs on par, if not better, compared to other algorithms in different heterogeneity settings, (2) there are multiple ways to reach the desired solution accuracy, one can either choose a large batch size and perform only a few local updates or select a 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 on CIFAR10 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 a two-hidden-layer convolutional neural network (CNN) architecture followed by three linear layers for CIFAR-10 and two for MNIST. For CIFAR-10 (and MNIST) datset, we consider three settings with mild, moderate and high heterogeneity. For all the three settings, the data is partitioned into disjoint sets among the WNs. In 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 on a popular FL benchmarking dataset, Shakespeare dataset [37]. For this task, we adopt the settings from [10] and utilize a 2-Layer LSTM network with 100 hidden units and an 8-D embedding layer at each WN. Each WN has access to 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 out of 143 for the training purpose. All the experiments are implemented on a single NVIDIA Quadro RTX 5000 GPU. More details are provided in appendix.
For the proposed STEM algorithm, recall that the step-size is ⌘t = ̄/(wt + 2t)1/3 with momentum parameter defined as at = c⌘2t . The step-size is used to update the iterates while the momentum
parameter is used to construct the stochastic gradient estimate (cf. Algorithm 1 and Theorem 3.1). For the experiments, we set wt = 2 = 1 and c = c̄/̄2 and tune for ̄ 2 [10 1, 10 2] for the CIFAR-10 dataset and for ̄ 2 {101, 100, 10 1, 10 2} for the Shakespeare dataset. For both the datasets we tune for c̄ in the range [1, 10]. For FedProx [10] and FedDyn [36] we choose the regularization constant to be 0.1. The momentum parameters for FedGLOMO [18] and MIME [17] are set based on the choices given in the respective papers. Specifically, for FedGLOMO we choose the parameter k = 0.2 and design the momentum gradient using a damping factor given in Appendix A.4 of FedGLOMO [18]. 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 {101, 100, 10 1, 10 2}.
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. Specific, 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 2b 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 we emphasize the importance of choosing b and I carefully. In Figure 2, we compare the training and testing performance of STEM, FedAvg and SCAFFOLD, against the number of samples accessed at each WN for the classification task on MNIST dataset with moderate heterogeneity. We fix b = 8 and conduct experiments under two settings, one with I = 67, and the other with I = 536 local updates at each WN. Note that although a 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 I = 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 I > 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 optimization with applications to FL. We showed that STEM reaches an ✏-stationary point with Õ(✏ 3/2) sample complexity while achieving linear speed-up with the number of WNs. Moreover, the algorithm achieves a communication complexity of Õ(✏ 1). We established a (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 (a.k.a LocalSGD) also exhibits a similar trade-off while achieving worse complexities. Our results provide guidelines to 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.

heading: Acknowledgement
text: 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 W911NF-19-1-0247. 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-2110259, CNS-2112471, CNS-2102233, CCF-2110252, ECCS-2140277, and a Google Faculty Research Award.

Review summaries:
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.

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.

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.

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.  Also, authors provide interesting observations about trade-off between the minibatch sizes and the local update frequency.    

 The 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.