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

Abstract: Prashant Khanduri, PRANAY SHARMA, Haibo Yang, Mingyi Hong, Jia Liu, Ketan Rajawat, Pramod Varshney

Reviews and author responses:

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.  The idea of using momentum-based variance reduction techniques [7, 8] to accelerate the convergence of FedAvg is very interesting, and novel as far as I know about. The main results make sense to me and appear to show the benefit of momentum acceleration in the regime of federated optimization. The analysis is theoretically sound. I am particularly impressed with the clear and clean writing style, and how all steps in the proofs are carefully elaborated in details. That said, I still have some reservations on this manuscript especially with regard to its technical originality and correctness. 

Major comments: 

- Originality: A major concern is about the novelty of the technical part, which follows closely the techniques developed in [7] for momentum SGD with slight adaptation to the map-reduce architecture considered by FL. While fancy and the final results look new in the context of FL, the fundamental contribution of this work is relatively weak in my opinion. As a side contribution, the complexity bounds for FedAvg revisit and slightly generalize the corresponding prior results in literature (see, e.g., [12, 15]). 

 - Clarity and correctness: The update rule of the average vector $\bar x_t$ and the relevant cost of communication are so far not clearly presented. Since $\bar x_t$ will be used in the output step for random selection, it has to be maintained at each time instance $t$. However, in Algorithm 1 & 2, $\bar x_t$ is only updated in the reduce step after every $I$ rounds of iteration; it is not clear how to update it at other time instances. Around Line #492 of the supplement file, it is written that $\bar x_t$ is updated by averaging the iterates from the workers at each time stamp. If this is the case (which I believe it is), then a round of communication will be needed to do the averaging at each time instance $t$, which will change a lot the story of the communication complexity analysis! The issue must be fixed to warrant correctness of analysis. 
 
- Parameter setting: According to Theorem 3.1, the learning rate of Algorithm 1 needs to be set with the knowledge of variance $\sigma$ which is typically unavailable in practice. I am wondering if it is possible to extend the analysis to the adaptive learning rate setup as considered in [7]? 

Minor comments: 

- Experiments: The method has been evaluated in some benchmark image classification tasks. It will be more interesting to see how the method works on some recent FL benchmarks such as Shakespeare and Sent140. 

- Problem formulation:  In Equation (1), it is confusing to denote $f^{(k)}$ simultaneously as population risk and loss function. 

- Initialization: the initial minibatch-size $B$ in Algorithm 1 & 2 seems never appeared in the convergence analysis and main results as well. 

- Step 9 of Algorithm 1: check the notation $\bar x^{(k)}_{t+1}$.

- In the proof of Lemma B.1, the second line of the equality (a) implicitly uses the independence among the minibatches $B^{(k)}_t$ so that the expectation operator can be moved into the inner product of stochastic gradient differences from each pair of nodes.  It is suggested to clarify such an independence assumption explicitly in the context.   

- The reference part needs to be carefully improved in formatting. 

=== Post rebuttal update ===

The major concerns raised in my initial review have been much clarified in the author response. I thus would like to change my rating one grid up to 6.   Yes, the authors have partially addressed the limitations and potential negative societal impact of their work.

 We thank the reviewer for his/her insightful comments and positive evaluation of our work. 

We agree with the reviewer that the listed assumptions are more general than the **standard variance assumption**, $\mathbb{E} \\| \nabla f^{(k)}(x ; \xi^{(k)}) - \nabla f^{(k)} (x) \\| \leq \sigma^2$, in the paper. These asumptions certainly present interesting future directions for generalizing the results obtained for STEM.





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.  

  This paper describes the Stochastic Two-Sided Momentum algorithm. This paper is well-written and it has a good structure. Authors provide better rates for non-convex case in the Federated Learning setting. They improved both sample and communication complexities. They show the trade-off between local computations and mini-batch size. The 3D plot is very illustrative in this case. These results are novel and relevant to the Federated Learning community. 

Theoretical analysis is not tight in terms of constants. For example 4096 in theorem 3.1 does not look to be the best constant. The assumptions that are used in the analysis are quite strong. In paper (https://arxiv.org/pdf/2012.04061.pdf) analysis does not need any additional assumptions. Another problem of this analysis is that we need to have a big mini-batch size (\tilde{O}(\varepsilon^{1/2}) or a big number of local computations (\tilde{O}(\varepsilon^{1/2}), but SCAFFOLD does not require this restriction. Also authors do not compare their result with this result in these papers (https://arxiv.org/pdf/2102.07845.pdf) and (https://arxiv.org/pdf/2007.01154.pdf)

Authors do not consider compression approach to reduce communication complexity. 

The most problematic part is experiments. It is not clear to me how hyperparameters were tuned. Especially, I do not understand how I and b were picked. Authors provide plots only for b = 8, I = 61 and b = 8, I=61. For the last experiment, we have that I = 536 and I = 67. I do not understand why authors use these parameters and do not try others. Also authors compare only three methods: STEM, SCAFFOLD and FedAvg. They do not compare new method with FedPD, FedProx, MIME, FedGLOMO and MARINA.

==========================================================================================================
After authors response I have increased my score. 

 Authors can add more explanations of experimental settings and process of tuning parameters. Additional experiments with different datasets and level of heterogeneity is desirable. To have fair competition authors can express complexity with mini-batch size parameter. 

If authors can resolve issues with experiments the score can be increased. 

 Thank you for clarifications! 

I think that there is some space for improvement. For example one can generalize obtain results using this assumption: 
$$\mathbb{E}\left\Vert\nabla f^{(k)}\left(x ; \xi^{(k)}\right)-\nabla f^{(k)}(x)\right\Vert^{2} \leq 2 A_k\left(f^{(k)}(x)-f^{(k)}_{\star}\right)+B_k^{2}. $$ 
This assumption appears in https://arxiv.org/pdf/2006.05988.pdf
 
or even more general assumption:

$$\mathbb{E}\left\Vert\nabla f^{(k)}\left(x ; \xi^{(k)}\right)-\nabla f^{(k)}(x)\right\Vert^{2} \leq 2 A_k\left(f^{(k)}(x)-f^{(k)}_{\star}\right)\+ B_k \cdot\Vert\nabla f^{(k)}(x)\Vert^{2}+C.$$

This assumption was introduced in https://arxiv.org/pdf/2002.03329.pdf

However, I think this paper deserves positive score, so I have increased my score. 





 We thank the reviewer for the valuable suggestions. 

- In the revised version of the paper, we will clearly define $a(T)$ and $\bar{x}_{a(T)}$ to remove any confusion. 
- We thank the reviewer for pointing this out. We will update the discussion following Theorem 3.2 to make this clear.  

 Thank you for providing the detailed clarifications. My major concerns have been much clarified in the response. I thus would upgrade my rating. Still, to make the output of Algorithm 1 well-defined without introducing additional communication cost, it is suggested to at least spell out the expression of $\bar x_{a(T)}$ at the random time instance $a(T)$ of termination. Also, it might be good to mention around Theorem 3.2 that $T = \Omega(K^3)$ is needed at least to make sure that $I,b\ge 1$.

 We sincerely thank the reviewer for his/her efforts in reviewing the paper and providing valuable feedback. Below, we address the comments. 

> **Your comment:** 1. Thank you for your clarifications. Many things are clear now. But there are some aspects that I want to discuss. I want to understand how much this assumption 2 is strongly related to the heterogeneity of the data and whether it can be verified in practice. I understand that this assumption allows to get better complexity, but I wonder how strongly this assumption affects the analysis, and it is interesting whether it is possible to obtain such complexity without this condition or with a weaker condition.

**Our response:** We thank the reviewer for the positive feedback. Below we discuss the inter-node variance bound (bounded gradient dissimilarity) stated as
$$\\| \nabla f^{(k)}(x) - \nabla f^{(\ell)}(x)  \\|^2 \leq \zeta^2 \qquad \forall~  k, \ell \in [K],$$
The above assumption can equivalently be stated as 
$\\| \nabla f^{(k)}(x) - \nabla f(x)  \\|^2 \leq \zeta^2 ~ \forall k \in [K],$ where $\nabla f(x)$ denotes the gradient of the overall objective function evaluated at $x$. This assumption captures the distance of the gradients at individual nodes from the full gradient and can be seen as a bounded variance assumption, but between the individual node's function and the overall objective function. Note that if the data across the nodes are i.i.d., then we have $\nabla f^{(k)}(x) = \nabla f^{(\ell)}(x)$ for all $k, \ell \in [K]$ which implies $\zeta = 0$. However, if the data across the nodes are non-i.i.d., we expect $\zeta > 0$. Intuitively, $\zeta$ captures the heterogeneity of data across nodes and is a widely adopted model used for describing data heterogeneity for non-convex FL algorithms; for example, please see [1-5].

The inter-node variance assumption is utilized by SGD-based algorithms to bound the error accumulation caused by the drift in local nodes' iterates as the local number of iterations increases. Note that this drift could be severe in FL problems due to the data heterogeneity and is challenging to bound even with the inter-node variance bound (please see Lemma B5 and Lemma C8 in the paper for more details). We also note that some works rely on assumptions like (i) Bounded gradients, please see [6,7]; (ii) Bounded Hessian dissimilarity [9]; (iii) Bounded dissimilarity [10], for analyzing FL algorithms, and they are much stronger compared to the inter-node variance assumption utilized by STEM.

To provide the reviewer an intuitive understanding as to why this assumption is mild, below we give two common examples of learning models, where the heterogeneity assumption is easily satisfied (for more examples please see Appendix A in [8])

- **Logistic regression:** Consider a logistic regression problem with each node having $n$ samples. Specifically, we represent the samples at the $k$th node as $ \\{a_i^{(k)}, b_i^{(k)} \\} $ for $i = 1$ to $n$. Note that $a_i^{(k)} \in \mathbb{R}^d$ are the features, whereas $b_i^{(k)} \in \mathbb{R}$ are the corresponding labels. Then the loss function at the $k$th node is:  
$$ f^{(k)} (x) = \frac{1}{n} \sum_{i=1}^n \frac{1}{1 + \exp \big(b_i^{(k)} - \big[a_i^{(k)}\big]^T x \big)}.$$
It can be shown that the inter-node gradient difference between two nodes' gradients is:
$$ \\| \nabla f^{(k)} (x)  - \nabla f^{(\ell)} (x)\\| \leq \frac{1}{\sqrt{n}} \Big[ \\|A^{(k)}\\| + \\|A^{(\ell)} \\| \Big],$$
    where $A^{(k)} = \big[ a_1^{(k)}, a_2^{(k)} , \ldots, a_n^{(k)} \big] \in \mathbb{R}^{d \times n}$. Therefore, $\zeta$ in the inter-node variance assumption can be written as $ \zeta = \sup_{k , \ell \in [K]} \Big\\{\frac{1}{\sqrt{n}} \Big[ \\|A^{(k)}\\| + \\|A^{(\ell)} \\| \Big] \Big\\} $.
- **Single Layer Neural network:** We can also show similar inter-node variance bound for a function defined by a single layer neural network. For ease of presentation, let us define a simple network of width $m$, with each node holding a single data sample $a^{(k)} \in \mathbb{R}^d$ as:
$$ f^{(k)}(x) = \frac{1}{m} \sum_{i = 1}^m \sigma(x_i^{\top}a^{(k)}), $$
    where $\sigma(\cdot)$ represents the activation function and $x = [x_1, x_2, \ldots, x_m] \in \mathbb{R}^{d m}$ are the weights of the activation functions. We use an non-linear activation function that is 1-Lipschitz and smooth. Then, the gradient $\nabla f^{(k)}(x)$ is:
$$ \nabla f^{(k)}(x)  =   \frac{1}{m} \Big[  \sigma'(x_1^{\top}a^{(k)}) a^{(k)\top}, \sigma'(x_2^{\top}a^{(k)}) a^{(k)\top}, \ldots, \sigma'(x_m^{\top}a^{(k)}) a^{(k)\top} \Big]^{\top}. $$
    Therefore, we have for any $k, \ell \in [K]$,
$$ \\| \nabla f^{(k)}(x) -  \nabla f^{(\ell)}(x) \\|  = \frac{1}{m} \Bigg\\|  \Big[  \sigma'(x_1^{\top}a^{(k)\top}) a^{(k)\top}, \ldots, \sigma'(x_m^{\top}a^{(k)}) a^{(k)\top} \Big]^{\top}  -   \Big[  \sigma'(x_1^{\top}a^{(\ell)}) a^{(\ell)\top}, \ldots, \sigma'(x_m^{\top}a^{(\ell)}) a^{(\ell)\top} \Big]^{\top} \Bigg\\|  \leq \zeta
$$
where $\zeta$ can be defined as $\zeta = \sup_{k , \ell \in [K]} \Big\\{   \\| a^{(k)}\\| + \\| a^{(\ell)}\\|    \Big\\}$. Therefore, a learning model defined by the above single-layer $m$-width neural network satisfies the inter-node variance bound. Then it can be easily verified that for any Lipschitz and smooth loss defined using the above network as, $L(f^{(k)}(x), b^{(k)})$, also satisfies the inter-node variance bound. 

In summary, the inter-node variance assumption can be theoretically verified for many popular problem classes.

$[1]$ Yu et al, On the linear speedup analysis of communication efficient momentum SGD for distributed non-convex optimization

$[2]$ Das et al, Improved convergence rates for non-convex federated learning with compression

$[3]$ Woodworth et al, Minibatch vs Local SGD for Heterogeneous Distributed Learning 

$[4]$ Li et al, Communication-Efficient Local Decentralized SGD Methods 

$[5]$ Reddi et al, Adaptive Federated Optimization

$[6]$ Li et al, On the Convergence of FedAvg on Non-IID Data

$[7]$ Yu et al, “Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning

$[8]$ Zhang et al, FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data

$[9]$ Li, et al, Federated Optimization in Heterogeneous Networks 

$[10]$ Karimireddi et al, Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning

> **Your comment:** 2. Thank you for providing additional experiments. I think it is not possible to show plots using this website. If you find a way to do it, I will be happy to see new plots.
>
> Authors clarified a lot of aspects so I am going to increase my score.

**Our response:** We thank the reviewer for the appreciation. We believe the reviewer is correct that it is not possible to show plots on OpenReview. However, we will include the plots of all additional experiments in the revised version of the paper. 



 Thank you for your clarifications. Many things are clear now. But there are some aspects that I want to discuss. I want to understand how much this assumption 2 is strongly related to the heterogeneity of the data and whether it can be verified in practice. I understand that this assumption allows to get better complexity, but I wonder how strongly this assumption affects the analysis, and it is interesting whether it is possible to obtain such a complexity without this condition or with a weaker condition. 

Thank you for the providing additional experiments. I think it is not possible to show plots using this website. If you find a way to do it, I will be happy to see new plots. 

Authors clarified a lot of aspects so I am going to increase my score. 

 > **Your Comment:** 1. This paper describes the Stochastic Two-Sided Momentum algorithm. This paper is well-written and it has a good structure. Authors provide better rates for non-convex case in the Federated Learning setting. They improved both sample and communication complexities. They show the trade-off between local computations and mini-batch size. The 3D plot is very illustrative in this case. These results are novel and relevant to the Federated Learning community.

**Our Response:** We thank the reviewer for this positive comment. Based on the reviewer's suggestions we have conducted additional experiments to thoroughly evaluate the performance of STEM against other algorithms. A summary of the additional experiments is included in the **General Response** section at the beginning of the response. In the revised version of the paper, we will include these results, together with detailed discussion about the parameter selection. In the following, we address the reviewer's comments.



>**Your Comment:** 2. Theoretical analysis is not tight in terms of constants. For example, $4096$ in Theorem 3.1 does not look to be the best constant.

**Our Response:** We would also like to point out that the constant $4096$ arises from the fact that we need to choose the step-size $\eta_t \leq 1/16 L$. Taking a further look into the proof, it seems unlikely that the constants can be improved by too much. 
Moreover, to keep the constant as an exponent of $2$, we have kept this choice. However, it is important to note that our focus is on obtaining algorithms that have optimal convergence rates, measured by the *order* of $\epsilon$. Therefore the constant $4096$  does not affect such an order of convergence rate guarantees.

>**Your Comment:** 3. The assumptions that are used in the analysis are quite strong. In paper (https://arxiv.org/ pdf/2012.04061.pdf) analysis does not need any additional assumptions.

**Our Response:** First, let us discuss and compare the assumptions made for FedGLOMO (and FedLOMO) in ``Faster Non-Convex Federated Learning via Global and
Local Momentum (https://arxiv.org/ pdf/2012.04061.pdf)" to that of STEM. What we observed is that the assumptions made for FedGLOMO are actually **stronger** compared to that for our STEM algorithm.  Below we discuss each assumption in detail:

- **Lipschitz smoothness:** FedGLOMO (and FedLOMO) requires per sample Lipschitz smoothness assumption, i.e., each stochastic sample gradient: $\nabla f^{(k)}(x, \zeta^{(k)})$ satisfies: $$\|| \nabla f^{(k)}(x, \zeta^{(k)}) - \nabla f^{(k)}(y, \zeta^{(k)}) \||^2 \leq L^2 \|| x - y\||^2$$
which is a much stronger assumption compared to the mean-square Lipschitzness assumption required by STEM, stated below:
$$\mathbb{E}\|| \nabla f^{(k)}(x, \zeta^{(k)}) - \nabla f^{(k)}(y, \zeta^{(k)}) \||^2 \leq L^2 \|| x - y\||^2.$$ Please see Assumption 1 in STEM compared to Assumption 1 in FedGLOMO.

- **Intra-Gradient variance bound:** Both FedGLOMO and STEM require the intra-gradient variance bound and achieve the same computational complexities. However, STEM achieves better communication complexity of $\mathcal{O}(\epsilon^{-1})$ compared to that of $\mathcal{O}(\epsilon^{-3/2})$ achieved by FedGLOMO.
Note that Intra-Gradient variance bound is referred to as bounded gradient dissimilarity in FedGLOMO. 

- **Full batch gradient computation:** In contrast to FedGLOMO, FedLOMO does not require Intra-Gradient variance bound but achieves worse (worst case) communication and computation complexities of $\mathcal{O}(\epsilon^{-2})$. FedLOMO avoids the "intra-gradient variance bound" assumption by computing full batch gradients at the start of each round and by assuming per sample Lipschitz smoothness (see the previous point). In contrast, STEM does not require computation of full gradients and achieves better communication and computation complexities of $\mathcal{O}(\epsilon^{{-1}})$ and $\mathcal{O}(\epsilon^{{-3/2}})$, respectively.

In summary, FedGLOMO relies on a much stronger per-sample Lipschitz smoothness assumption compared to the mean-square Lipschitness assumption required by STEM. Both algorithms do rely on intra-gradient variance bound.

>**Your Comment:** 4. Another problem of this analysis is that we need to have a big mini-batch size $\tilde{O}(\epsilon^{1/2})$ or a big number of local computations $\tilde{O}(\epsilon^{1/2})$, but SCAFFOLD does not require this restriction.

**Our Response:** We thank the reviewer for raising this concern, which we would like to further clarify in here. Usually, the goal of FL algorithms is to minimize the total communication rounds required to achieve a desired convergence accuracy, which is referred to as "communication complexity'' in the literature. To achieve an improved communication complexity, it is often desired that nodes perform multiple local computations before parameter sharing. In fact, it is an advantage of STEM that it allows the nodes to perform a larger number of local updates (or larger batch gradients, or trade-off batch gradients with number of local updates) to achieve a near-optimal communication complexity. Importantly, STEM achieves near-optimal communication complexity while attaining optimal **overall** sample complexity. More specifically, even if nodes are performing multiple local updates within each communication round, the total samples required throughout the entire convergence process is $\mathcal{O}(\epsilon^{-3/2})$, which is optimal for stochastic non-convex problems. Moreover, note that even if we choose a **constant** number of local updates with a constant batch size (as suggested by the reviewer), our algorithm still achieves better computation and communication complexities of $\mathcal{O}(\epsilon^{-3/2})$ compared to those of SCAFFOLD. Please note that this point has already been mentioned in the original paper, Remark 3, where we commented that if both the batch size and iteration number are chosen as $\mathcal{O}(1)$, the total number of communication rounds remains to be $\mathcal{O}(\epsilon^{-3/2})$.

>**Your Comment:** 5. Also authors do not compare their result with this result in these papers (https://arxiv.org/pdf/ 2102.07845.pdf) and (https://arxiv.org/pdf/ 2007.01154.pdf).
Authors do not consider compression approach to reduce communication complexity.

**Our Response:** In the introduction of the revised version of the paper, we will add discussions of FedCOM and MARINA mentioned by the reviewer. Here, we will provide a comparison between our work and FedCOM/MARINA.

First, we note that MARINA is not a typical FL algorithm as it **does not** utilize multiple number of local updates (and the problems considered therein are deterministic). 
Second, FedCOM and MARINA are compression-based algorithms, which are not directly comparable to our work. We emphasize that the goal of our current work is to design an algorithm that (i) achieves the optimal sample complexity for FL and, (ii) achieves (near) optimal communication complexity. We believe that it is important to understand the best possible performance for uncompressed version of FL algorithms. After that,  we can combine our new algorithm with existing compression schemes to further reduce communication costs. Due to space limitations, we cannot focus on both topics in a single paper, therefore, we decide to exclusively study the uncompressed version, and try to study it thoroughly.

On the other hand, we agree with the reviewer that compression is an important direction to improve the practical communication performance of algorithms. We plan to study how state-of-the-art techniques can be applied in the STEM algorithm in our future works.

>**Your Comment:** 6. The most problematic part is experiments. It is not clear to me how hyper-parameters were tuned. Especially, I do not understand how $I$ and $b$ were picked. Authors provide plots only for $b = 8$, $I = 61$ and $b = 64$, $I=7$. For the last experiment, we have that $I = 536$ and $I = 67$. I do not understand why authors use these parameters and do not try others.

**Our response:** We would like to clarify how we chose parameters $b$ and $I$. Our parameters choices for $b$ and $I$ are guided by our theoretical results in Theorem 3.1 and Remark 3.
Specifically, according to Theorem 3.1 and Remark 3, a good choice of $b$ and $I$ is that $b\cdot I \approx \mathcal{O}(\epsilon^{-1/2} )$, i.e., when $b\cdot I$ is fixed to a large constant. For the initial set of experiments, we have set $b\cdot I \approx 490$ (same as the number of local data samples at each node). This implies that with $b = 64$ we choose $I = \lfloor490/64\rfloor = 7$ and with $b = 8$ we choose $I = \lfloor 490/8\rfloor = 61$. We choose $I = 536$ and $I = 67$ using the same calculation for the MNIST dataset. Note that this strategy is adopted for choosing $b$ and $I$ because it allows us to efficiently trade-off $b$ and $I$ according to Theorem 3.1. In the revised version of the paper, we will add these explanations together with additional experiments with different choices of hyper-parameters.

>**Your Comment:** 7. Also, the authors compare only three methods: STEM, SCAFFOLD, and FedAvg. They do not compare the new method with FedPD, FedProx, MIME, FedGLOMO, and MARINA.

**Our Response:** We thank the reviewer for pointers to these related works.
Based on the reviewer's suggestion, in the rebuttal period, we have conducted additional experiments to further evaluate the performance of STEM. Specifically, given time and computational resource constraints in this rebuttal period, we have been able to compare the proposed algorithm with FedProx, FedDyn, MIME, FedGLOMO, SCAFFOLD, and FedAvg. We have included a summary of the results at the beginning of the response in the **General Response** section.

In the revised version of the paper, we will add the complete experiments with sufficient details.

 We thank the reviewers for taking their valuable time to provide insightful comments and suggestions for the paper. We believe that the thoughtful reviews and the recommendations made by the reviewers have substantially improved the quality of the paper. Based on reviewers' suggestions we have conducted additional numerical experiments to compare the performance of STEM to other algorithms under different settings. In this response, we include a summary of the results obtained. We will provide the complete details of parameter selection and algorithm implementation in the revised manuscript. 

We evaluate the performance of STEM against multiple algorithms for (i) different heterogeneity settings, (ii) multiple batch size choices and, (iii) two different datasets, specifically, in addition to CIFAR-10 dataset, we also include additional results for the Shakespeare dataset [1]. Since we have additionally included FedDyn [2] in our results, we choose not to include FedPD since they are known to be very closely related. 

### Additional algorithms and heterogeneity setting

In Tables R1, R2, and R3 we present the results for different algorithms for CIFAR-10 dataset under different heterogeneity and batch-sizes choices (note, in Table R1 and R2, newly included algorithms are highlighted using bold letters; Table R3 is completely new since it represents a new high heterogeneous setting). Note that STEM performs uniformly well under all the conditions. Moreover, note from Table R2 that FedGLOMO diverges for the case when the number of local updates is high. Also, note from Table R3 that when the heterogeneity is high, FedProx performs well while SCAFFOLD's performance degrades significantly.  

> **Table R1**  CIFAR-10 data is randomly split among nodes (moderate heterogeneity setting). We choose $b = 64$ and $I = 7$ after $500$ rounds of communication. Highlighted rows are new.\
> **Algorithm**  $\quad$   **Training Accuracy(%)** $\quad$    **Test Accuracy(%)** 
>---
>FedAvg    $~\qquad \qquad \quad$       78.2     $\qquad \qquad \quad \qquad$                74.1
>
> **FedProx** $\qquad ~~~\quad \quad$   79.2 $\qquad \qquad \quad \qquad$  74.8   
>
> **FedDyn** $~~\qquad ~~\quad \quad$  68.9 $\qquad \qquad \quad \qquad$  66.0
>
>SCAFFOLD $\quad \quad \qquad$  71.9 $\qquad \qquad \quad \qquad$  74.0
>
> **MIME** $~~~\qquad \qquad \quad $  82.6 $\qquad \qquad \quad \qquad$  76.8
>
> **FedGLOMO** $~~\quad \quad \quad $  76.1 $\qquad \qquad \quad \qquad$  72.8
>
>STEM $~~~~\qquad  \quad \qquad$  80.1$\qquad \qquad ~\quad \qquad$  78.8


> **Table R2**  CIFAR-10 data is split among nodes such that each node can only access $5$ classes (high heterogeneity setting). We choose $b = 8$ and $I = 61$ after $500$ rounds of communication. Highlighted rows are new.\
> **Algorithm**  $\quad$   **Training Accuracy(%)** $\quad$    **Test Accuracy(%)** 
>---
>FedAvg    $~\qquad \qquad \quad$       73.6     $\qquad \qquad \quad \qquad$                75.4 
>
> **FedProx** $~~~\quad \qquad \quad$   80.0 $\qquad \qquad \quad \qquad$  75.2   
>
> **FedDyn** $~~\qquad ~~\quad \quad$  76.1 $\qquad \qquad \quad \qquad$  71.3 
>
> SCAFFOLD $\quad \quad \qquad$  72.5 $\qquad \qquad \quad \qquad$  73.7 
>
> **MIME** $~~~\qquad \qquad \quad $  61.5 $\qquad \qquad \quad \qquad$  58.6  
>
> **FedGLOMO** $~~\quad \quad \quad $  10.0 $\qquad \qquad \quad \qquad$  10.0  
>
>STEM $~~~~\qquad  \quad \qquad$  81.1 $\qquad \qquad \quad \qquad$  78.5 


> **Table R3**  CIFAR-10 data is split among nodes such that each node can only access $2$ classes (extreme heterogeneity setting). Training and testing accuracy for $b = 128$ and $I = 6$ after $300$ rounds of communication. All the results in this table are new.\
> **Algorithm**  $\quad$   **Training Accuracy(%)** $\quad$    **Test Accuracy(%)** 
>---
>FedAvg    $~\qquad \qquad \quad$       57.6     $\qquad \qquad \quad \qquad$               57.1
>
>FedProx $\qquad \qquad \quad$   59.1  $\qquad \qquad \quad \qquad$  58.5
>
>FedDyn $~~~\qquad ~~\quad \quad$  51.2 $\qquad \qquad \quad \qquad$  51.3
>
>SCAFFOLD $\;\quad \quad ~~~\quad$  40.3 $\qquad \qquad \quad \qquad$  41.3 
>
> MIME $~~~~\qquad \qquad \quad $  56.1 $\qquad \qquad \quad \qquad$  55.1
>
>FedGLOMO $~~~\quad \quad \quad $  56.8 $\qquad \qquad \quad \qquad$  56.1
>
>STEM $~~~~\qquad  \quad \qquad$  58.5  $\qquad \qquad \quad \qquad$  57.4 

### Federated learning Benchmarking Dataset

In Table R4, we evaluate the performance of STEM against other algorithms for the task of next-character prediction on the Shakespeare dataset [1]. Note that the Shakespeare dataset is a popular benchmarking dataset for FL algorithms. We choose the number of local updates as the number of updates required to complete a single pass over each node's data. We present the training and testing accuracy in Table R4. Note from Table R4 that STEM outperforms the rest of the algorithms. We also note that because of time and computation constraints we were unable to average the experiments over multiple runs. In the revised version of the manuscript, we will average the performance of all the algorithms over multiple runs for the Shakespeare dataset and report the performance with error bars. Moreover, we will include the details of all the parameter tuning and experiment settings in the updated manuscript.


> **Table R4**  Performance comparison for Shakespeare dataset with $b = 128$. We report the performance after $50$ rounds of communication. All the results in this table are new.\
> **Algorithm**  $\quad$   **Training Accuracy(%)** $\quad$    **Test Accuracy(%)** 
>---
>FedAvg    $~\qquad \qquad \quad$       40.1    $\qquad \qquad \quad \qquad$                39.2
>
>FedProx $\qquad ~~~\quad \quad$   43.5  $\qquad \qquad \quad \qquad$  43.2
>
>FedDyn $~~\qquad ~~\quad \quad$  43.7 $\qquad \qquad \quad \qquad$  43.2 
>
>SCAFFOLD $\quad \quad \qquad$  40.3 $\qquad \qquad \quad \qquad$  41.3 
>
> MIME $~~~\qquad \qquad \quad $  32.1 $\qquad \qquad \quad \qquad$  32.1
>
>FedGLOMO $~~\quad \quad \quad $  40.3 $\qquad \qquad \quad \qquad$  40.1
>
>STEM $~~~~\qquad  \quad \qquad$  44.5 $\qquad \qquad \quad \qquad$  43.8

$[1]$ Caldas et al, LEAF: A Benchmark for Federated Settings 

$[2]$ Acar et al, Federated Learning Based on Dynamic Regularization



 > **Your Comment:** 1. The idea of using momentum-based variance reduction techniques [7, 8] to accelerate the convergence of FedAvg is very interesting, and novel as far as I know about. ... The analysis is theoretically sound. I am particularly impressed with the clear and clean writing style, and how all steps in the proofs are carefully elaborated in details. That said, I still have some reservations on this manuscript especially with regard to its technical originality and correctness.

**Our Response:** We thank the reviewer for the positive comment. Below we address other comments in detail. 

> **Your Comment:** 2. Originality: A major concern is about the novelty of the technical part, which follows closely the techniques developed in [7] for momentum SGD with slight adaptation to the map-reduce architecture considered by FL. While fancy and the final results look new in the context of FL, the fundamental contribution of this work is relatively weak in my opinion. As a side contribution, the complexity bounds for FedAvg revisit and slightly generalize the corresponding prior results in literature (see, e.g., [12, 15]). 

**Our Response:** We first note that FL and Map-Reduce are quite different. The traditional Map-Reduce architecture is used for distributed processing of computing jobs that are in a *one-way* traffic flow, while distributed learning problems in general (and FL in particular) are usually iterative in nature with *back-and-forth* communication between the nodes and the server. Moreover, for FL problems the completion of jobs highly depends on the convergence process. In other words, FL only uses a *map-reduce-like* operation in each iteration. 
Also, we are not entirely sure if the reviewer meant map-reduce or *all-reduce* operation for parameter sharing. If the reviewer meant *all-reduce*, we would like to point out that FL is more general than "all-reduce," which is just a method for aggregation without requiring a dedicated parameter server. In other words, all-reduce is just an option in the implementation of FL. Our paper does not assume whether the underlying network uses all-reduce or not.

Moreover, we respectfully disagree that FedAvg and/or our STEM are trivial extensions of SGD or Momentum SGD to Map-Reduce (or All-Reduce) architecture. Note that FedAvg and/or STEM will be an easy extension *only* if the local nodes share their parameters immediately after each local update and/or the data is *homogeneous* across nodes. 
However, a key distinguishing design element of FL algorithms is the use of local updates that allows each node to compute its parameters locally (for multiple iterations) without any intervention from the server. With these multiple local updates and heterogeneous data across worker nodes, many open problems arise, and there have been a lot of recent works devoted to developing and analyzing FL algorithms. In this work, we mainly focused on designing FL algorithms that can achieve:

- Near-optimal communication rounds to achieve a desired solution.
- Optimal computation complexity.
- Linear speed-up with the number of nodes. 

Note that achieving all the three objectives above is difficult when the data across the nodes are heterogeneous. The reason is that, with heterogeneous data, the local iterates at each worker node tend to drift further from the global optimal solution with a larger number of local updates. Thus, it is important to note that, to our knowledge, the STEM algorithm is the *first* FL algorithm to achieve all three goals simultaneously, while allowing heterogeneous data distribution across worker nodes. Moreover, the analysis presented in the paper reveals a fundamental trade-off between the batch sizes and the number of local updates required to achieve these optimal guarantees, and these tradeoffs have not been discovered before. More importantly, STEM uniformly improves the guarantees provided by algorithms proposed in [12] and [15].

Finally, we would like to comment that despite the fact that the momentum-based gradient estimator in STEM is not completely new, but making it work in the FL setting has not been done, and it is non-trivial because of the difficulties mentioned above. One key technical innovation of our work is to identify the most suitable way to incorporate momentum-based directions in FL algorithms (that is, the so-called two-sided momentum updates), so that the desired properties mentioned above can be achieved; please see lines 137 to 147 and Remark 1 in the original paper for more discussions.

> **Your Comment:** 3. Clarity and correctness: The update rule of the average vector $\bar{x}_t$ and the relevant cost of communication are so far not clearly presented. Since $\bar{x}_t$ will be used in the output step for random selection, it has to be maintained at each time instance $t$. However, in Algorithm 1 \& 2, $\bar{x}_t$ is only updated in the reduce step after every rounds of $I$ iteration; it is not clear how to update it at other time instances. Around Line number 492 of the supplement file, it is written that $\bar{x}_t$ is updated by averaging the iterates from the workers at each time stamp. If this is the case (which I believe it is), then a round of communication will be needed to do the averaging at each time instance $t$, which will change a lot the story of the communication complexity analysis! The issue must be fixed to warrant correctness of analysis. 

**Our Response:** There appears to be some misunderstanding in our STEM algorithm. We would like to emphasize that STEM **does not** require averaging the iterates at each time instant. Instead, the averaging is performed after every $I$ local updates at each worker node 8 [see Step 9 in Algorithm 1]. More specifically, the variable $\bar{x}_t$ mentioned by the reviewer, which appears in Line 492 of the appendix, is a **virtual** variable constructed only for analytical purposes. It represents the average of the local parameters at time $t$, and it is evaluated to characterize the distance between each node's parameters from the average because with multiple local updates the local parameters tend to drift away from the average.
 
About the question on random selection, please note that the random selection is used for theoretical purposes to show that at a randomly selected time the average parameter will be an $\epsilon$-stationary solution in expectation. We note that it is common to use random selection to facilitate theoretical analysis for non-convex stochastic algorithms  (see, e.g., [Remark 1 in [1], Remark 2 in [2], [3]). In this paper, we simply follow the convention in the literature. In practice, this can be easily achieved by generating a random iteration count $a(T)$ uniformly from $\{ 1, \ldots, T\}$ and terminating the algorithm at iteration $a(T)$ and output $\bar{x}_{a(T)}$ as the solution (please see [4]). 

In the revised version of the manuscript, we will add all the above clarifications.

$[1]$ Yu et al, “Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning,

$[2]$ Yu et al “On the linear speedup analysis of communication efficient momentum sgd for distributed non-convex optimization

$[3]$ Fang et al, “Spider: Near-optimal non-convex optimization via stochastic path-integrated differential estimator 

$[4]$ Ghadim et al, Stochastic first-and zeroth-order methods for nonconvex stochastic programming

> **Your Comment:** 4. Parameter setting: According to Theorem 3.1, the learning rate of Algorithm 1 needs to be set with the knowledge of variance $\sigma$ which is typically unavailable in practice. I am wondering if it is possible to extend the analysis to the adaptive learning rate setup as considered in [7]? 

**Our Response:** In practice, $\sigma^2$ can be estimated by computing the empirical variance using a few representative samples. Meanwhile, the step-size (learning rate) of reference [7] in the paper depends on $G$, the bound on the stochastic gradients which is also difficult to determine in practice. This means that even for this kind of adaptive gradient method, there is still some quantities that need to be estimated from data.  

In general, for FL algorithms it is usually difficult to extend the analysis to algorithms with adaptive learning rates, particularly when the worker nodes are allowed to perform multiple local updates. This problem arises because with adaptive learning rates, each node employs a different step-size strategy based on its local data. This makes it challenging to uniformly bound the parameter drift in each round. However, as shown in [1] it might be possible to use adaptive step-sizes at the server. Designing adaptive algorithms when each node uses an adaptive step-size is an interesting research problem in its own right and deserves an independent study in our future work of STEM. We thank the reviewer for bringing up this topic.

$[1]$ Reddi et al., Adaptive Federated Optimization. 

> Minor comments: 

- Thanks for your suggestion. In this rebuttal period, we have been able to conduct additional experiments on additional algorithms to further evaluate the performance of STEM. A summary of the experimental results is provided at the beginning of the response in the **General Response** section.
-In the revised version of the manuscript, we will denote the loss function by $F^{(k)}(\cdot)$ and the population risk by $f^{(k)}(\cdot)$.
-The initial batch size $B$ is defined in line 156 in the statement of Theorem 3.1.
- Thanks for catching this typo. We will correct the typo in the revised version of the manuscript.
-Thanks for your suggestions. We will update the proof of Lemma B.1 with clear explanations of the independence assumptions. We will also update the formatting of the references. 




 We greatly thank the reviewer for positive evaluation of our work.


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.   Main:

## Originality:

The momentum-based variance reduction techniques are known. Nevertheless, it is new and interesting to see how to combine them into Federated Learning. The paper clearly discussed the differences and comparisons with the previous works. 


## Quality: 

The paper provides theoretical guarantees for the propsed STEM algorithm. Specifically, the paper gives a clear description of the algorithm and indicates the choices of hyper-parameters (number of local updates $I$ and batch size $b$). Then the paper provides convergence guarantees under *sample gradient smoothness condition* for non-convex functions. Next, the paper shows how to simutaneously achieve optimal sample and communication complexity with the optimal batch sizes and local updates. 

For comparison with the baseline algorithms, the paper also provides theoretical guarantees for FedAvg and gives the batch sizes and local updates trade-off. In the numerical experiment part, the paper shows their proposed STEM algorithm is better than FedAvg and SCAFFOLD.

## Clarity: 

This paper is clearly written, with careful comparisons to previous works, clearly stated assumptions and complexity measures, rigorous expression of the theorems, and detailed remarks on their algorithms. 

## Significance: 

The paper is of theoretical significance. An appropriate complexity measure is the key to study of the federated learning. The paper studies both the stochastic gradient complexity (sample complexity) as well as the communication complexity for FedAvg and their proposed STEM. It provides useful message to the community.

The paper proposes STEM for federated learning under *sample gradient smoothness condition* and provides instructions on hyper-parameter choices. In reality, the objective function may satisfy this stronger smoothness assumption and in this case the proposed STEM is useful, with only a small change to the implementation of the baseline algorithms. Yes

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


Decision: accept