Momentum Benefits Non-iid Federated Learning Simply and Provably

As for experiments, we have added a new experiment with $N=100$ clients with MNIST dataset in the appendix of the revision, see Figure 5 and 6. The results show that FedAvg-M and Scaffold-M outperform FedAvg and Scaffold by an evident margin when $N=100$, which is closer to the pratical FL setups.

Also, since FedAvg-M coincides with FedCM. The superior performance of FedCM reported in FedCM paper is also an evidence that momentum can bring significant benefits, which aligns well with our theory.

Finally, we would like to emphasize that the main contribution of our paper is to clarify the theoretical benefits brought by momentum. We hope the reviewer can find our empirical studies have well justified all our theoretical findings.

Questions

Q1: How to build the relationship between $\nabla f_i(x)$ and $\nabla f(x)$ without the bounded data heterogeneity assumption?

A: We briefly present the high-level idea of FedAvg-M in the full participation setting. Note that when the local learning rate $\eta\to 0$ and global learning rate $\gamma$ is fixed, FedAvg will reduce to minibatch SGD over all the clients, i.e., SGD over $f(x)$ with batch size $NK$. In this case, the heterogeneity assumption $\zeta$ is known to be unnecessary. Intuitively, in FedAvg-M, our strategy is also to show that the local iterates will not drift too far away from the initialization in each communication round (Lemma 9) if $\eta$ is properly small with the help of momentum.

Q2: Lemma 5 seems incorrect.

A: As pointed out and confirmed by reviewer ELLH in the part of "Typos", the second line in Lemma 5 is merely a typo and does not affect the overall validity. The correction should be

$$f(x^{r+1})\leq f(x^r)+\langle \nabla f(x^r),x^{r+1}-x^r\rangle +\frac{L}{2}\\|x^{r+1}-x^r\\|^2 = f(x^r)-\gamma \\|\nabla f(x^r)\\|^2+\gamma \langle \nabla f(x^r),\nabla f(x^r)-g^{r+1}\rangle+\frac{L\gamma^2}{2}\\|g^{r+1}\\|^2.$$

We have addressed this in our revision.

Q3: How to use Young inequality?

A: In page 14, $\gamma \langle \nabla f(x^r),\nabla f(x^r)-g^{r+1}\rangle \leq \frac{\gamma }{2}\\|\nabla f(x^r)\\|^2+\frac{\gamma }{2} \\|\nabla f(x^r)-g^{r+1}\\|^2$.

In page 17 and 22, $\mathbb{E} [\\|\nabla f_i(x^r)\\|^2] \leq (1+q)\mathbb{E}[\\|\nabla f_i(x^{r-1})\\|^2]+(1+q^{-1})\\| \nabla f_i(x^r)-\nabla f_i(x^{r-1})\\|^2 \leq (1+q)\mathbb{E}[\\|\nabla f_i(x^{r-1})\\|^2]+(1+q^{-1})L^2\mathbb{E}[\\|x^r-x^{r-1}\\|^2]$.

In page 25, $\left(1-\frac{S}{N}\right)\frac{1}{N}\sum_{i=1}^N\mathbb{E}[\\|c_i^r-\nabla f_i(x^r)\\|^2]\leq \left(1-\frac{S}{N}\right)\frac{1}{N}\sum_{i=1}^N\mathbb{E}\left[\left(1+\frac{S}{2N}\right)\\|c_i^r-\nabla f_i(x^{r-1})\\|^2+\left(1+\frac{2N}{S}\right)\\| \nabla f_i(x^r)-\nabla f_i(x^{r-1})\\|^2\right]$ and $\\| \nabla f_i(x^r)-\nabla f_i(x^{r-1})\\|^2 \leq L^2\\|x^r-x^{r-1}\\|^2$.

Q4: About the second inequality in page 22 $\mathbb{E}||x^{r} – x^{r−1}||^2 \leq 2\gamma^{2}\mathbb{E}(\mathcal{E}_{r-1} + \mathbb{E}||\nabla f(x^{r-1})||^{2})$.

A: First, we would like to emphasize that we did not use any additional assumption. Note that by the definition of $\mathcal{E}_{r-1}$ in the beginning of Appendix A, we have

$\mathbb{E}\\|x^r-x^{r-1}\\|^2 = \gamma^2\mathbb{E}\\|g^r\\|^2 \leq 2\gamma^2 \mathbb{E}(\\|g^r-\nabla f(x^{r-1})\\|^2 + \\|\nabla f(x^{r-1})\\|^2) = 2\gamma^2 (\mathcal{E}_{r-1}+\mathbb{E}[\\|\nabla f(x^{r-1})\\|^2])$.

And we derive a recursive bound for $\mathcal{E}_{r-1}$ in Lemma 8. Therefore, we do not need the bounded heterogeneity assumption.

Q5: About VRL-SGD.

A: Thanks for pointing this out and we have made additional comments on the work in our revision, as highlighted in the red color in Table 1. For reviewer's convenience, we repeat the comments as follows: The works have not been published in peer-reviewed venues.

We hope the above clarification can relieve the reviewer's concern. We believe our paper makes a novel and important contribution to the community on understanding the theoretical benefits of momentum in federated learning, which has never been clarified before. We are happy to address any further questions or commmens from the reviewer.

[R1] Xu, Jing, et al. ‘Fedcm: Federated Learning with Client-Level Momentum’. arXiv Preprint arXiv:2106. 10874, 2021.

[R2] Karimireddy, Sai Praneeth, Satyen Kale, et al. ‘Scaffold: Stochastic Controlled Averaging for Federated Learning’. International Conference on Machine Learning, PMLR, 2020, pp. 5132–5143.

[R3] Karimireddy, Sai Praneeth, Martin Jaggi, et al. ‘Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning’. arXiv Preprint arXiv:2008. 03606, 2020.

We thank the reviewer for the comments. All questions have been clarified. We are glad to address any further comments or questions.

1. Comparison with FedCM

Thanks for pointing this out and bringing FedCM [R1] to our attention. It is indeed the same as our FedAvg-M. We will remove the statements like "resulting in the new algorithm FedAvg-M" from our paper. Please note that the main novelty of our paper does not lie in the algorithm development, but in clarifying the theoretical improvements brought by momentum. The FedCM paper cannot justify the theoretical benefits by incorporating momentum.

To resolve the reviewer's concern, we added a detailed comparison with FedCM in the revised paper, see the highlighted paragraph in Page 6. For reviewer's convenience, we repeat the highlighted paragraph as follows:

FEDAVG-M coincides with the FEDCM algorithm proposed by Xu et al. (2021b). However, our results outperform that of Xu et al. (2021b) in several aspects. First, our convergence only utilizes the standard smoothness of objectives and gradient stochasticity while Xu et al. (2021b) additionally require bounded data heterogeneity and bounded gradients which are rarely valid in practice, suggesting the limitation of their analysis. Second, the convergence established by Xu et al. (2021b) is significantly weaker than ours and cannot even asymptotically approach the ideal rate $O(1/ \sqrt{NKR})$ in non-convex FL, as demonstrated by the results stated in Table 1.

The rate comparison established in FedCM and ours are as follows:

FedCM: $\left(\frac{L\Delta ({\sigma^2}+NK{G}^2)}{NKR}\right)^{1/2}+\left(\frac{L\Delta ({\sigma}/{\sqrt{K}}+G)}{R}\right)^{2/3}$.

FedAvg-M: $\left(\frac{L\Delta \sigma^2}{NKR}\right)^{1/2}+\frac{L\Delta}{R}$.

It is observed that the rate established by FedCM is much weaker than FedAvg-M.

2. Mistakes in the proof

Please check the detailed responses to questions below.

3. The extra requirement over the communication budget or the local storage

First, it is generally true that downlink bandwidth is substantially larger than uplink bandwidth in most practical systems. This significant asymmetry between downlink and uplink capacities suggests that heavier downlink loads may be less problematic compared to lighter uplink loads.

Second, our FedAvg-M does not incur additional downlink load compared to FedAvg when more memory is available. For practical implementation of FedAvg-M, the server does not need to communicate both $x^{r+1}$ and $g^{r+1}$ simultaneously to clients. If clients can store a model copy $x^r$ locally, by only broadcasting the latest server model $x^{r+1}$ from the server to all clients, clients are able to recover the momentum direction through $g^{r+1}=(x^{r+1}-x^r)/\gamma$. By employing the equivalent implementation, FedAvg-M does not suffer from a double in downlink communication.

Moreover, we would like to remark that the doubling in local storage is customary for advanced FL algorithms beyond FedAvg, see, e.g., Scaffold [R2], Mime [R3]. In fact, as long as the local update rule is not the vanilla SGD, the double local storage can be inevitable to store optimizer states, e.g., momentum state.

To resolve the reviewer's concern, we added comment in Section 3 of the revised paper. The comment is "Notably, no extra downlink commmunication cost is needed if clients store the last iterate model $x^r$ so that momentum $g^{r+1}$ can recovered through $(x^{r+1}-x^r)/\gamma$."

4. The method is not novel and the experiment is too simple

We respectively disagree with the reviewer on this point. The main novelty in our paper lies in proving the theoretical improvements brought by momentum, not in proposing new momentum variants of FL algorithms. While momentum has been widely used in federated learning empirically, it is still unknown in literature how momentum can benefit federated learning theoretically. Our paper focused on this open question and proved two novel results:

• Momentum can remove the influence of data heterogeneity. When all clients participate in FL, momentum enables FedAvg to get rid of the restrictive data heterogeneity assumption without any impractical algorithmic structure.

• Momentum can theoretically speed up FL algorithm. We have shown that FedAvg, Scaffold, and their variance-reduced variants with momentum achieve faster convergence rate than without momentum. In fact, all achieved results are state-of-the-art compared to existing algorithms.

With the above two contributions, we are the first, to our best knowledge, to validate the usage of momentum in federated learning in theory, which we believe is a novel and important contribution to the community.

We thank the reviewer for the follow-up question. We use $\xi$ to denote the random data sampled per iteration within each client and use $\mathcal{E}_r = \mathbb{E}[\|\nabla f(x^r) - g^{r+1}\|^2]$ throughout the manuscript. We believe the reviewer meant to ask for the clarification regarding the latter one.

The reviewer has misunderstood the term $\mathcal{E}_r = \mathbb{E}[\\|\nabla f(x^r) - g^{r+1}\\|^2]$. This term is NOT an assumption. We do not need to assume the boundedness of $\mathbb{E}[\\|\nabla f(x^r) - g^{r+1}\\|^2]$. Instead, its boundedness can be derived rigorously buildng upon the standard objective smoothness (i.e., Assumption 1) and stochastic gradients (i.e., Assumption 3). Specifically, we have shown in the proof of Theorem 11 that

$$\sum_{r=0}^{R-1}\mathcal{E}\_r\leq \frac{9}{7\beta }\mathcal{E}\_{-1} + \frac{2}{7}\mathbb{E} [\sum_{r=-1}^{R-2}\\|\nabla f(x^{r})\\|^2] + \frac{144}{7}(e\beta \eta KL)^2G_0 R+ \frac{36\beta \sigma^2}{7NK}R$$

where $G_0=\sum_{i=1}^N\\|\nabla f_i(x^0)\\|^2/N$ and $\mathcal{E}_{-1} = \\|\nabla f(x^0) - g^{0}\\|^2$ are all constants by definition. By combining this with Lemma 5, we have

\frac{1}{7} \sum_{r=0}^{R-1} \mathbb{E}\left[\left\\|\nabla f\left(x^r\right)\right\\|^2\right] \leq \frac{1}{\gamma} \mathbb{E}\left[f\left(x^0\right)-f^\star\right] +\frac{39}{56 \beta} \mathcal{E}_{-1}+\frac{78}{7}(e \beta \eta K L)^2 G_0 R+\frac{39 \beta \sigma^2}{14 N K} R,

which reveals that $\mathbb{E} [\sum_{r=-1}^{R-2}\\|\nabla f(x^{r})\\|^2]/R$ is bounded provided with properly small $\eta$ and $\beta$. This, in turn, reveals that $\sum_{r=0}^{R-1}\mathcal{E}_r/R$ is bounded is the averaged sense, which is sufficient to attain the ultimate convergence rate. It is guaranteed that $\mathcal{E}_r$ cannot go to $+\infty$, in the averaged sense.

Again, We would also like to emphasize that no additional assumption regarding the data heterogeneity, i.e., $\sup_x \frac{1}{N}\sum_{i=1}^N\\|\nabla f_i(x)-\nabla f(x)\\|^2$, is needed in our paper due to the usage of momentum, which is a non-trival contribution to the community.

We hope this can address your concern. We are glad to have further discussion on this point.

We thank the reviewer for the detailed comments. All questions have been clarified. We are glad to address any further comments or questions.

1. Comparison with FedCM

Thanks for bringing FedCM to our attention. It is indeed the same as our FedAvg-M. We will remove the statements like "resulting in the new algorithm FedAvg-M" from our paper. Please note that the main novelty of our paper does not lie in the algorithm development, but in clarifying the theoretical improvements brought by momentum. The FedCM paper cannot justify the theoretical benefits by incorporating momentum.

To resolve the reviewer's concern, we added a detailed comparison with FedCM in the revised paper, see the highlighted paragraph in Page 6. For reviewer's convenience, we repeat the highlighted paragraph as follows:

FEDAVG-M coincides with the FEDCM algorithm proposed by Xu et al. (2021b). However, our results outperform that of Xu et al. (2021b) in several aspects. First, our convergence only utilizes the standard smoothness of objectives and gradient stochasticity while Xu et al. (2021b) additionally require bounded data heterogeneity and bounded gradients which are rarely valid in practice, suggesting the limitation of their analysis. Second, the convergence established by Xu et al. (2021b) is significantly weaker than ours and cannot even asymptotically approach the ideal rate $O(1/ \sqrt{NKR})$ in non-convex FL, as demonstrated by the results stated in Table 1.

The rate comparison established in FedCM and ours are as follows:

FedCM: $\left(\frac{L\Delta ({\sigma^2}+NK{G}^2)}{NKR}\right)^{1/2}+\left(\frac{L\Delta ({\sigma}/{\sqrt{K}}+G)}{R}\right)^{2/3}$.

FedAvg-M: $\left(\frac{L\Delta \sigma^2}{NKR}\right)^{1/2}+\frac{L\Delta}{R}$.

It is observed that the rate established by FedCM is much weaker than FedAvg-M.

2. VR rates

Thanks for the sharp observation. There is no mismatch between our rate and that in [R1]. The difference lies in the number of data batches used in the initialization stage (i.e., the the number of data batches in the very first iteration). In the paper we set initial number of batches $B$ as $\Theta\left((NKR)^{1/3}(\frac{\sigma^2}{L\Delta})^{4/3}\right)$ for simplicity so we can get rid of the additional $\sigma^2/\epsilon^2$ term. On the other hand, if we set initial batch size as $\Theta(NKR)$ as constrained by the setup in [R1], our rate will match the lower bound stated in [R1].

To resolve the reviewer's concern, we also considered the setting with initial batch size $\Theta(NKR)$ and derived the convergence rate, see Theorem 15 and Theorem 22 in the revised paper. We also list the new convergence rate in Table 1. For reviewer's convenience, we put the comparison in FedAvg-M-VR as follows:

For initial batchsize $B=\Theta\left((NKR)^{1/3}(\frac{\sigma^2}{L\Delta})^{4/3}\right)$, the rate is $\left(\frac{L\Delta \sigma}{NKR}\right)^{2/3}+\frac{L\Delta}{R}$. For initial batchsize $B=\Theta(NKR)$, the rate is $\left(\frac{L\Delta \sigma}{NKR}\right)^{2/3}+ \frac{\sigma^2}{NKR}+\frac{L\Delta}{R}$.

3. More experiments

Thanks for the advice. Due to the hardware resource limitation and tight rebuttal deadline, we cannot afford an experiment with $N=100$ with Cifar-10 dataset at this stage. Instead, we have conducted experiments with $N=100$ clients with MNIST dataset in the appendix of the revision, see Figure 5 and 6. The results show that FedAvg-M and Scaffold-M outperform FedAvg and Scaffold by an evident margin when $N=100$.

Also, since FedAvg-M coincides with FedCM. The superior performance of FedCM reported in FedCM paper is also an evidence that momentum can bring significant benefits, which aligns well with our theory. Finally, we would like to emphasize that the main contribution of our paper is to clarify the theoretical benefits brought by momentum. The empirical studies are mainly to confirm our theoretical findings.

4. Downlink load

First, it is generally true that downlink bandwidth is substantially larger than uplink bandwidth in most practical systems. This significant asymmetry between downlink and uplink capacities suggests that heavier downlink loads may be less problematic compared to lighter uplink loads.

Second, our FedAvg-M does not incur additional downlink load compared to FedAvg when more memory is available. For practical implementation of FedAvg-M, the server does not need to communicate both $x^{r+1}$ and $g^{r+1}$ simultaneously to clients. If clients can store a model copy $x^r$ locally, by only broadcasting the latest server model $x^{r+1}$ from the server to all clients, clients are able to recover the momentum direction through $g^{r+1}=(x^{r+1}-x^r)/\gamma$. By employing the equivalent implementation, FedAvg-M does not suffer from a double in downlink communication.

To resolve the reviewer's concern, we added comment in Section 3 of the revised paper. The comment is "Notably, no extra downlink commmunication cost is needed if clients store the last iterate model $x^r$ so that momentum $g^{r+1}$ can recovered through $(x^{r+1}-x^r)/\gamma$."

5. Link between $R$ and $\beta$

In Theorem 1 we set $\beta=1$ when $R$ is small for simplicity. However, we remark that $\beta$, in principle, can be any $\Theta(1)$ constant that is less than $1$ instead of being exactly $1$. The convergence rates will not affected by this choice, which can be easily derived with very minor modifications to the proof.

To resolve the reviewer's concern, we have revised the choise of $\beta$ in Theorems 1, 11, 15, 19 and 22.

6. Momentum scales with more clients

Thanks for raising this very interesting question. Based on our theories, the incorporation of momentum (once properly tuned) scales well when more clients participate in FL with the linear speedup term $O(1/\sqrt{NKR})$. With our new experiments with $N=100$, we do not observe the collapse of Scaffold-M. Once we have more computational recourse and allowed time, we will conduct more expereiments to examine it.

7. Difference from [R4]

[R4] focuses on communication compression in distributed optimization and does not involve any fundamental characteristic of FL, i.e. local updates and partial participation. Therefore, our analysis is fundamentally different from [R4]. On the other hand, we propose momentum to handle the effect of client drift in local iterations, which is beyond the scope of [R4]. We don't think there is a direct connection between our analysis and [R4] except for the usage of momentum in algorithmic development.

8. References

We thank the reviewer for bringing various relevant references to our attention. We have cited them properly and added corresponding discussions in the revised paper.

[R1] Arjevani, Yossi, Carmon, Yair, Duchi, John C., Foster, Dylan J., Srebro, Nathan and Woodworth, Blake. Lower bounds for non-convex stochastic optimization, Mathematical Programming, 2023.

[R2] Karimireddy, Sai Praneeth, Satyen Kale, et al. ‘Scaffold: Stochastic Controlled Averaging for Federated Learning’. International Conference on Machine Learning, PMLR, 2020, pp. 5132–5143.

[R3] Karimireddy, Sai Praneeth, Martin Jaggi, et al. ‘Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning’. arXiv Preprint arXiv:2008. 03606, 2020.

[R4] Fatkhullin, Ilyas, et al. ‘Momentum Provably Improves Error Feedback!’ arXiv Preprint arXiv:2305. 15155, 2023.

We thank the reviewer again for the careful and valuable comments. We hope these response can clarify the reviewer's questions. We are looking forward to the follow-up discussion, and more than happy to address any further comments or questions.

I thank the authors for their detailed answers.

Given that:

• They recognized the existence of FedCM and added a comparison of convergence rates in Table 1, as well as added a discussion paragraph in their main paper,
• They clarified my concern about their convergence rate of the VR versions, and added those clarifications in the paper,
• They added more experiments to show that FedAVG-M is indeed helping at scale,
• They corrected my statement by remarking that no additional downlink communications are needed for a small memory overhead on the clients' part,
• They clarified my concern about the use of momentum for low training rounds $R$,
• The fact that, upon a second quick check of the proof of their Theorem 1, I did not see any blatant error in their derivations, making me confident about the validity of their result,

and, most importantly, given that the theoretical contributions of this paper on showing the ability of a simple momentum term to provably help federated optimization without relying on the bounded heterogeneity assumption seems to be of high interest for the community,

I subsequently raise my score and support fully the acceptance of the paper.

I would, however, point out that more experiments on CIFAR-10 with ResNets to discuss the scaling properties of the momentum versions of the algorithms as well as the proper tuning of it when scaling in less of a "toy scenario" would strengthen the paper.

We thank the reviewer for the comments. All questions have been clarified. We are glad to address any further comments or questions.

1. Influence of momentum coefficient

The momentum coefficient does not appear in the final convergence rate because we have set up a concrete and proper value for $\beta$ when deriving the convergence rate. For example, when establishing the convergence of FedAvg-M in Theorem 11 in the appendix, we first prove that

$\frac{1}{R}\sum_{r=0}^{R-1}\mathbb{E}\|\nabla f(x^r)\|^2 = O\left(\frac{L\Delta}{R} + \frac{L\Delta}{\beta R} + \frac{\beta \sigma^2}{NK} + (\beta \eta K L)^2 G_0\right)$

which reflects how $\beta$ would influence the convergence rate. When we set $\beta = O(\sqrt{\frac{NKL\Delta}{\sigma^2 R}})$ and let $\eta \beta L$ be sufficiently small (see Theorem 11 for the value of $\eta \beta L$), we can simplify the convergence rate as

$\frac{1}{R}\sum_{r=0}^{R-1}\mathbb{E}\|\nabla f(x^r)\|^2 = O\left(\sqrt{\frac{L \Delta \sigma^2}{NKR}} + \frac{L\Delta}{R}\right)$

which is the SOTA convergence rate. Notably, we remark the above annealing $\beta$ is mainly set for theory purpose.

In practice, we found a proper constant value (e.g., $\beta=0.5, 0.2$) is also capable to attain substantial improvement, as demonstrated by our experiments.

2. FedDyn and our algorithms are not comparable

We respectively disagree with the comments that FedDyn demonstrates faster convergence rate than our proposed algorithms. FedDyn and our algorithm belong to different algorithm families, and they are using different oracles in local updates. Our algorithm use the standard stochastic gradient oracle in local updates, while FedDyn relies on the black-box oracle that gives the exact solutions to local optimization problems (see line 5 in Algorithm 1 of FedDyn). Note that the local optimization problem is generally intractable (particularly for non-convex functions) and thus the oracle required by FedDyn is usually impractical, which outweighs its benefits in convergence if any. If the local optimization problem is not solved exactly, it is not known whether the current rate in FedDyn still hold. For this reason, we believe FedDyn and our algorithms are not comparable. The convergence rate of FedDyn does not hurt our theoretical contribution.

3. Empirical comparison with federated adaptive algorithms

We thank the reviewer for bringing up the federated adaptive algorithms. However, the main contribution in our paper is to clarify the theoretical improvements brought by introducing momentum to FL algorithms. This goal, to our knowledge, has not been achieved by [R1] or any other literature.

Again, we are not aiming to develop adaptive algorithms that can outperform any federated adaptive algorithms in [R1] empirically. All of our experiments are conducted to validate our theoretical findings, i.e., momentum can help remove the influence of data heterogeneity and enhance the convergence rate. Whether our algorithms can outperform [R1] or not does not provide direct evidence to our theoretical findings. We hope the reviewer can understand that the comparison between our algorithms and adaptive algorithms in [R1] is not necessary and beyond the scope of our paper.

4. Experiments on different $\beta$

As illustrated in the descriptions of algorithms, when $\beta=1$, ours will reduce to vanilla FedAvg or Scaffold. And in the experiments, we have shown the difference between $\beta=1$ and other constants, i.e. between the non-momentum variant and the momentum variant. Also, we supplement an ablation study on different choices of $\beta\in (0, 1)$ in the appendix of revision. The results suggest that stronger momentum can lead to better practical performance, see Figures 7 and 8 in Section D.5 in the appendix.

5. Scaffold with low sampling rate

We thank the reviewer for the sharp observation. When the number of sampled clients $S$ is far less than the number of all clients $N$, i.e., $S \ll N$, all federated learning algorithms listed in Table 2 will suffer from slower convergence, see their rates in Table 2. It implies that slow convergence in the scenario with a low client-sampling ratio is a common issue in federated learning.

However, incorporating momentum does bring some benefits. For example, incorporating momentum has improved Scaffold's convergence rate from $O(\sqrt{\frac{L \Delta \sigma^2}{S K R}} + \frac{L \Delta}{R}(N/S)^{2/3})$ to $O(\sqrt{\frac{L \Delta \sigma^2}{S K R}} + \frac{L \Delta}{R}(N^{2/3}/S))$, which quantitively illustrates that momentum indeed accelerates the rate.

We hope the above clarification can relieve the reviewer's concerns. We believe our paper makes a novel and important contribution to the community on understanding the theoretical benefits of momentum in federated learning, which has never been clarified before.

[R1] Reddi, Sashank, et al. ‘Adaptive Federated Optimization’. arXiv Preprint arXiv:2003. 00295, 2020.

We thank the reviewer for the comments. All questions have been clarified. We are glad to address any further comments or questions.

1. Novelty

The main novelty in our paper lies in proving the theoretical improvements brought by momentum, not in proposing new momentum variants of FL algorithms. While momentum has been widely used in federated learning empirically, it is still unknown in literature how momentum can benefit federated learning theoretically. Our paper focused on this open question and proved two novel results:

• Momentum can remove the influence of data heterogeneity. When all clients participate in FL, momentum enables FedAvg to get rid of the restrictive data heterogeneity assumption without any impractical algorithmic structure.

• Momentum can theoretically speed up FL algorithm. We have shown that FedAvg, Scaffold, and their variance-reduced variants with momentum achieve faster convergence rate than without momentum. In fact, all achieved results are state-of-the-art compared to existing algorithms.

With the above two contributions, we are the first, to our best knowledge, to validate the usage of momentum in federated learning in theory, which we believe is a novel and important contribution to the community.

2. Comparison with federated ADAM

Reference [R1] is a great and influential paper. However, it does not provide sufficient theoretical justification for the benefits brought by momentum. [R1] primarily focuses on proving convergence guarantees, but does not clearly explicate the advantages of using momentum and adaptability. Its analysis is limited in that it does not clarify whether and how momentum can relax restrictive assumptions or improve convergence rates. The major novelty in our analysis is to utilize new analytical techniques to clarify that momentum can relieve data heterogeneity assumption and speed up convergence theoretically.

There are also several important technical differences between [R1] and our paper. First, the analysis in [R1] relies on the impractical and restrictive assumptions of bounded gradient and bounded client variance, which can significantly simplify their analysis. In our analysis, we have relaxed all these assumptions. Second, the algorithm in [R1] use vanilla SGD in local updates (see Algorithm 2 in [R1]), while our paper incorporates momentum in local updates, which also cause significant difference in analysis.

We hope the above clarification can relieve the reviewer's concern on the novelty. We believe our paper makes a novel and important contribution to the community on understanding the theoretical benefits of momentum in federated learning, which has never been clarified before.

[R1] Reddi, Sashank, et al. ‘Adaptive Federated Optimization’. arXiv Preprint arXiv:2003. 00295, 2020.