Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper proposes a new method to solve the finite sum problem in federated learning where each local model is an expectation function. The STEM algorithm uses momentum variance-reduced estimator in local updates (similar to STORM estimator in stochastic optimization) and also send these estimators to the server when performing server-side update. Therefore, STEM includes momentum updates for both local and server updates. Convergence analysis reveals that STEM nearly match the best known sample complexity in centralized setting and it also matches the lower bound communication complexity up to a log factor. Numerical experiments on neural network training not only show advantage of STEM compared with FedAvg and SCAFFOLD but also presenting the tradeoff between communication frequency and minibatch size.
main_review: Pros:
- The paper introduces a new framework to incorporate the momentum gradient estimators in federated learing compared to 2 related works which are MIME and FedGLOMO. Given FedGLOMO is rather new, I consider it as concurrent work.
- Convergence analysis provides intuition on the tradeoff between communication and sample complexities (perform more local updates vs communicate more). This effect is further illustrated in the experiments
- Complexity results are also new.
- The analysis of number of local updates and minibatch size

Cons:
- The paper presents FedAvg with full-device participation and compare the results with this version. However, FedAvg or SCAFFOLD do allow partial-participation in their convergence analysis while STEM seems to require full device participation.
- Apart from the local model, STEM also sends "local momentum" which is an estimate of the local gradient. This framework does not exactly match the federated learning framework described in [https://arxiv.org/pdf/1610.05492.pdf, page 2] where local nodes can only send local model back to the server. Moreover, these gradient estimates are susceptible to privacy attack in federated learning [https://proceedings.neurips.cc/paper/2020/file/c4ede56bbd98819ae6112b20ac6bf145-Paper.pdf]. Although this is not exactly the true gradient, but it is still a good estimate (if not then the algorithm is not working well) so it is good to avoid communicating the first-order information. I am aware that the same criticism can be applied to SCAFFOLD, MIME or FedGLOMO.

Minor comments:
- It is better to elaborate the steps from (b) to (c) for the second term at line 504 similar to (9) at line 523 for clarity.
- There should be a $\nabla f^{(k)}(x^{(k)}_t;\xi_t^{(k)})$ at line 561.
limitations_and_societal_impact: Regarding limitations, I have not seen the authors state this clearly in the paper. I don't see any immediate negative societal impact of this work.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper introduced a two-step momentum method for federated learning. The proposed method achieves the best achievable sample and communication complexities for first-order stochastic FL algorithms. The theoretical analysis provides guidance on selecting synchronization interval and batch size for achieving the nearly optimal complexities.
main_review: This paper introduced an interesting idea for achieving near-optimal sample and communication complexities for federated learning. Overall, the paper is well written, but I would suggest authors explicitly state Definitions 2.1 and 2.2 before referencing them. More detailed comments are shown as follows

1. Why do we need to use the consecutive iterates x^k_{t+1} and x^k_t to evaluate the estimator? Can't we just use x^k_{t+1}? How does it affect the complexity bounds?

2. As pointed out by (Khaled, Ahmed, Konstantin Mishchenko, and Peter Richtárik, 2019), a bound on the dissimilarity is not realistic for heterogeneous data setting. Under this assumption, the data is very close to i.i.d.

3. The experiments seems weak in terms of model selection and data scale. Federated EMNIST dataset is a better benchmark data for federated learning (Caldas, Sebastian, et al. 2018).

Khaled, Ahmed, Konstantin Mishchenko, and Peter Richtárik. "First analysis of local gd on heterogeneous data." arXiv preprint arXiv:1909.04715 (2019).
Caldas, Sebastian, et al. "Leaf: A benchmark for federated settings." arXiv preprint arXiv:1812.01097 (2018).
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: This paper studies non-convex stochastic optimization for federated learning under the mean-squared smoothness assumption and shows that better sample convergence and communication complexity are possible with large batch sizes. 
main_review: To me, this is a borderline paper. This paper has its technical merit by establishing stronger sample/communication complexity bounds. However, there seem to be some fundamental issues on the interpretation correctness and presentation clarity.  

The proposed Algorithm 1 is named " Stochastic Two-Sided Momemtum (STEM)", which also appears in the paper title. It is named so because the authors interpret line 9 in Algorithm 1 as "server side momentum step" and argue its novelty and importance for better convergence. However, I don't think it is necessary to let the server do the momentum SGD update in line 9.  We could let each client do local momentum SGD updates in line 8 and then simply average all local x at line 9.  With such an equivalent description,there is no "server side momentum" concept needed. It is really hard for me to appreciate the necessity of server side momentum. In fact, I think line 7-10 is mathematically equivalent to " updating each local node with momentum SGD at each iteration and average all local x at the server when t mod I = 0", which is just the standard "FedAvg with momentum SGD local solver "or "local SGD with momentum".  In this regard, it seems this paper does not propose a new algorithm but just analyzes the existing FedAvg with momentum SGD local solvers under the stronger Assumption 1.  I would also suggest the authors revisit their proof and comment if the proof also follows similar lines as in early works. 

Table 1 lists the complexity of earlier works that can be misleading to general readers who do not have sufficient optimization backgrounds. The current comparison is not apple to apple since the complexities are attained under different technical assumptions. Please add an extra column to list the needed assumption for each row. 

The idea of using large batch sizes to achieve faster convergence may let practitioners frown since large batch sizes may hurt the final model accuracy. In the experiment, the methods are only tested with b=8 or b=64. (These settings are usually considered as O(1) batch size.) More experiments with larger b might be desired.






limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper proposes and analyzes an optimization algorithm (STEM) with two-sided momentum and local steps for (non-convex) Federated Learning. The authors prove its convergence and 
They also claim that it improves over the previous methods in terms of join communication and sample complexity.
main_review: This work considers a problem of non-convex federated optimization, where communication is the bottleneck. A local stochastic gradient method with momentum at the server and workers is suggested. STEM enjoys $\tilde{\mathcal{O}}(\epsilon^{-3/2})$ sample and $\tilde{\mathcal{O}}(\epsilon^{-1})$ communication complexity if mini-batch size and number of local updates are chosen appropriately (e.g. as $\mathcal{O}(\epsilon^{-1/2})$). The authors show a trade-off between the mini-batch sizes and local update frequency.
Experimental comparison to FedAvg and SCAFFOLD on two image classification tasks (with convolutional neural networks) confirmed the benefits of the proposed method.

Related work was mostly adequately cited and compared with.

Overall, the paper is well-organized and leaves an impression of good work, but I think that in its current stage, it is not ready for publication. 
The first reason is the theoretical contributions. The method does not suggest a very novel technique because algorithms with momentum and local steps were previously analyzed in the federated setting. 
The so-called “near-optimal” sample and communication complexities are obtained for mini-batches and number of local updates of the order of $\mathcal{O}(\epsilon^{-1/2})$, which means that for achieving an $\epsilon$-stationary point for $\epsilon = 10^{-10}$ (which corresponds to gradient norm $\||\nabla f(x_t)\|| \leq 10^{-5}$) the method requires very big mini-batches or the number of local updates, which seems not very practical. At the same time, for methods like SCAFFOLD and MIME, it is not required. That is why I think that theoretical comparison is not completely fair from the practical side.

The second issue is the lack of experimental comparison to the methods presented in the comparison Table 1. The chosen benchmarking methods (FedAvg and SCAFFOLD) are the easiest competitors to STEM because the first one is the most basic baseline, while the second is more suited for convex problems in the cross-silo settings. Both of them also have worse theoretical complexities. To show real practical benefits of the proposed method, it is required to compare against FedPD and MIME/FedGLOMO as they have the same sample or communication complexities to STEM. Furthermore, the hyperparameter tunning seems very limited (just two values for $\overline{\kappa}$ and $\overline{c}$). I also recommend multiple runs for every experiment to show mean values and deviations as the methods are stochastic. This is a basic standard to show the consistent superiority of the proposed approach over the baselines.

In addition, real experiments supporting the point on the trade-off between mini-batch size and the number of local updates are quite limited to basically a couple of values (of $b$ and $I$), which does not seem enough to show this effect and provide valuable guidelines for practitioners. If the computational resources are the problem, I suggest the authors using a simpler non-convex model [2]. Showing plots with gradient norm values would be also more consistent with the theory in this case.

To conclude, I think that the current results are not very important nor valuable for the field.

Unfortunately, chosen notation (e.g., for stochastic gradients) complicates the algorithm description and makes convergence proofs very hard to follow and understand. I greatly encourage authors to simplify it. E.g. for local gradient estimates $\nabla f^{kj}$ could be used in case of finite sum losses $f^{(k)} = \sum_{j=1}^J f^{kj}$.

*Minor issues/comments:*

I think that the authors missed an important reference [1].

There is a typo in the Appendix on page 18 in equation (8). In the RHS of the (b) inequality, there should be superscript (j) after the last $\nabla$ symbol.

[1] R. Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020)

[2] L. Zhao, et al. "From Convex to Nonconvex: A Loss Function Analysis for Binary Classification," 2010 IEEE International Conference on Data Mining Workshops, 2010, pp. 1281-1288, doi: 10.1109/ICDMW.2010.57
limitations_and_societal_impact: Limitations of the work were mostly addressed. Negative societal impact concerns are not applicable due to the theoretical nature of the paper.
review_title: null

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