mlsanityData to 14 Oct 2026

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

arXiv 2106.10435 ↗First posted 19 Jun 2021cs.LG math.OC stat.MLcited by 43

Not on the map. Why?

Abstract (arXiv)

Federated Learning (FL) refers to the paradigm where multiple worker nodes (WNs) build a joint model by using local data. Despite extensive research, for a generic non-convex FL problem, it is not clear, how to choose the WNs' and the server's update directions, the minibatch sizes, and the local update frequency, so that the WNs use the minimum number of samples and communication rounds to achieve the desired solution. This work addresses the above question and considers a class of stochastic algorithms where the WNs perform a few local updates before communication. We show that when both the WN's and the server's directions are chosen based on a stochastic momentum estimator, the algorithm requires tildeO(ε-3/2) samples and tildeO(ε-1) communication rounds to compute an ε-stationary solution. To the best of our knowledge, this is the first FL algorithm that achieves such it near-optimal sample and communication complexities simultaneously. Further, we show that there is a trade-off curve between local update frequencies and local minibatch sizes, on which the above sample and communication complexities can be maintained. Finally, we show that for the classical FedAvg (a.k.a. Local SGD, which is a momentum-less special case of the STEM), a similar trade-off curve exists, albeit with worse sample and communication complexities. Our insights on this trade-off provides guidelines for choosing the four important design elements for FL algorithms, the update frequency, directions, and minibatch sizes to achieve the best performance.

Key facts. Cited by 43 papers in the map (data through 14 Oct 2026). Builds on 30 of its 41 references that are in the map. First posted on arXiv 19 Jun 2021.

Builds on

30 of 41 references.

References · 30

  1. 2016Communication-Efficient Learning of Deep Networks from Decentralized Datacited by 7,669
  2. 2018Federated Optimization in Heterogeneous Networkscited by 2,987
  3. 2016Optimization Methods for Large-Scale Machine Learningcited by 1,150
  4. 2018Federated Learning with Non-IID Datacited by 1,108
  5. 2019On the Convergence of Adam and Beyondcited by 841
  6. 2018LEAF: A Benchmark for Federated Settingscited by 792
  7. 2018Local SGD Converges Fast and Communicates Littlecited by 605
  8. 2016Federated Optimization: Distributed Machine Learning for On-Device Intelligencecited by 564
  9. 2019Robust and Communication-Efficient Federated Learning from Non-IID Datacited by 419
  10. 2020A Unified Theory of Decentralized SGD with Changing Topology and Local Updatescited by 367
  11. 2019Momentum-Based Variance Reduction in Non-Convex SGDcited by 306
  12. 2018SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimatorcited by 287
18 more
  1. 2018Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learningcited by 278
  2. 2018Don't Use Large Mini-Batches, Use Local SGDcited by 202
  3. 2021Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learningcited by 193
  4. 2019On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimizationcited by 180
  5. 2020Mime: Mimicking Centralized Stochastic Algorithms in Federated Learningcited by 154
  6. 2020Is Local SGD Better than Minibatch SGD?cited by 151
  7. 2020Minibatch vs Local SGD for Heterogeneous Distributed Learningcited by 147
  8. 2020FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Datacited by 144
  9. 2019SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentumcited by 121
  10. 2018Cooperative SGD: A unified Framework for the Design and Analysis of Communication-Efficient SGD Algorithmscited by 106
  11. 2017On the convergence properties of a K-step averaging stochastic gradient descent algorithm for nonconvex optimizationcited by 92
  12. 2019Variance Reduced Local SGD with Lower Communication Complexitycited by 80
  13. 2018Stochastic Nested Variance Reduction for Nonconvex Optimizationcited by 63
  14. 2019The Complexity of Finding Stationary Points with Stochastic Gradient Descentcited by 35
  15. 2019Hybrid Stochastic Gradient Descent Algorithms for Stochastic Nonconvex Optimizationcited by 32
  16. 2019On the Computation and Communication Complexity of Parallel SGD with Dynamic Batch Sizes for Stochastic Non-Convex Optimizationcited by 15
  17. 2019Parallel Restarted SPIDER -- Communication Efficient Distributed Nonconvex Optimization with Optimal Computation Complexitycited by 8
  18. 2014Protecting Privacy through Distributed Computation in Multi-agent Decision Makingcited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 15; 2023: 11; 2024: 8; 2025: 6; 2026: 3.

43 papers.

Sort:

By area:

  1. 2022Generalized Federated Learning via Sharpness Aware Minimizationcited by 144
  2. 2023Stochastic Controlled Averaging for Federated Learning with Communication Compressioncited by 101
  3. 2022FedVARP: Tackling the Variance Due to Partial Client Participation in Federated Learningcited by 64
  4. 2022On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyondcited by 38
  5. 2023Momentum Benefits Non-IID Federated Learning Simply and Provablycited by 38
  6. 2022Faster Adaptive Federated Learningcited by 35
  7. 2022Communication-Efficient Federated Learning with Accelerated Client Gradientcited by 24
  8. 2022Accelerated Federated Learning with Decoupled Adaptive Optimizationcited by 23
  9. 2022DASHA: Distributed Nonconvex Optimization with Communication Compression, Optimal Oracle Complexity, and No Client Synchronizationcited by 16
  10. 2023On the Role of Server Momentum in Federated Learningcited by 16
  11. 2024The Limits and Potentials of Local SGD for Distributed Heterogeneous Learning with Intermittent Communicationcited by 12
  12. 2022Taming Fat-Tailed ("Heavier-Tailed'' with Potentially Infinite Variance) Noise in Federated Learningcited by 10
31 more
  1. 2023Solving a Class of Non-Convex Minimax Optimization in Federated Learningcited by 10
  2. 2023Communication-Efficient Federated Bilevel Optimization with Local and Global Lower Level Problemscited by 8
  3. 2023Federated Learning with Manifold Regularization and Normalized Update Reaggregationcited by 8
  4. 2022Fast Adaptive Federated Bilevel Optimizationcited by 7
  5. 2022DIAMOND: Taming Sample and Communication Complexities in Decentralized Bilevel Optimizationcited by 7
  6. 2023When Decentralized Optimization Meets Federated Learningcited by 7
  7. 2025On Provable Benefits of Muon in Federated Learningcited by 7
  8. 2026FedBCD:Communication-Efficient Accelerated Block Coordinate Gradient Descent for Federated Learningcited by 7
  9. 2024FedLion: Faster Adaptive Federated Optimization with Fewer Communicationcited by 6
  10. 2022Adaptive Federated Minimax Optimization with Lower Complexitiescited by 5
  11. 2023Federated Conditional Stochastic Optimizationcited by 4
  12. 2022FedSSO: A Federated Server-Side Second-Order Optimization Algorithmcited by 3
  13. 2024Momentum Approximation in Asynchronous Private Federated Learningcited by 3
  14. 2024Understanding Server-Assisted Federated Learning in the Presence of Incomplete Client Participationcited by 3
  15. 2025Federated Learning Survey: A Multi-Level Taxonomy of Aggregation Techniques, Experimental Insights, and Future Frontierscited by 3
  16. 2022Faster Adaptive Momentum-Based Federated Methods for Distributed Composition Optimizationcited by 2
  17. 2023FedDA: Faster Framework of Local Adaptive Gradient Methods via Restarted Dual Averagingcited by 2
  18. 2023Oracle Complexity Reduction for Model-free LQR: A Stochastic Variance-Reduced Policy Gradient Approachcited by 2
  19. 2024Non-Convex Optimization in Federated Learning via Variance Reduction and Adaptive Learningcited by 2
  20. 2022On the Convergence of Momentum-Based Algorithms for Federated Bilevel Optimization Problemscited by 1
  21. 2023FedDRO: Federated Compositional Optimization for Distributionally Robust Learningcited by 1
  22. 2024Momentum-Based Federated Reinforcement Learning with Interaction and Communication Efficiencycited by 1
  23. 2025Federated Learning on Riemannian Manifolds: A Gradient-Free Projection-Based Approachcited by 1
  24. 2022SAGDA: Achieving O(ε-2) Communication Complexity in Federated Min-Max Learningcited by 0
  25. 2024FedEGG: Federated Learning with Explicit Global Guidancecited by 0
  26. 2024Convergent Differential Privacy Analysis for General Federated Learningcited by 0
  27. 2025TACO: Tackling Over-correction in Federated Learning with Tailored Adaptive Correctioncited by 0
  28. 2025Enhancing Privacy in Decentralized Min-Max Optimization: A Differentially Private Approachcited by 0
  29. 2025Adaptive Federated Learning via Dynamical System Modelcited by 0
  30. 2026Federated Compositional Muon Optimizer for Matrix-Wise Modelscited by 0
  31. 2026A Momentum-Based Variance-Reduced Algorithm for Federated Multiobjective Optimizationcited by 0