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

arXiv: https://arxiv.org/abs/2106.10435 · HTML: https://mlsanity.com/paper/2106.10435/

## 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 (most cited first, 30 of 30)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [LEAF: A Benchmark for Federated Settings](https://mlsanity.com/paper/1812.01097/) (2018)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Robust and Communication-Efficient Federated Learning from Non-IID Data](https://mlsanity.com/paper/1903.02891/) (2019)
- [A Unified Theory of Decentralized SGD with Changing Topology and Local Updates](https://mlsanity.com/paper/2003.10422/) (2020)
- [Momentum-Based Variance Reduction in Non-Convex SGD](https://mlsanity.com/paper/1905.10018/) (2019)
- [SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimator](https://mlsanity.com/paper/1807.01695/) (2018)
- [Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning](https://mlsanity.com/paper/1807.06629/) (2018)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learning](https://mlsanity.com/paper/2101.11203/) (2021)
- [On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization](https://mlsanity.com/paper/1905.03817/) (2019)
- [Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning](https://mlsanity.com/paper/2008.03606/) (2020)
- [Is Local SGD Better than Minibatch SGD?](https://mlsanity.com/paper/2002.07839/) (2020)
- [Minibatch vs Local SGD for Heterogeneous Distributed Learning](https://mlsanity.com/paper/2006.04735/) (2020)
- [FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data](https://mlsanity.com/paper/2005.11418/) (2020)
- [SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum](https://mlsanity.com/paper/1910.00643/) (2019)
- [Cooperative SGD: A unified Framework for the Design and Analysis of Communication-Efficient SGD Algorithms](https://mlsanity.com/paper/1808.07576/) (2018)
- [On the convergence properties of a K-step averaging stochastic gradient descent algorithm for nonconvex optimization](https://mlsanity.com/paper/1708.01012/) (2017)
- [Variance Reduced Local SGD with Lower Communication Complexity](https://mlsanity.com/paper/1912.12844/) (2019)
- [Stochastic Nested Variance Reduction for Nonconvex Optimization](https://mlsanity.com/paper/1806.07811/) (2018)
- [The Complexity of Finding Stationary Points with Stochastic Gradient Descent](https://mlsanity.com/paper/1910.01845/) (2019)
- [Hybrid Stochastic Gradient Descent Algorithms for Stochastic Nonconvex Optimization](https://mlsanity.com/paper/1905.05920/) (2019)
- [On the Computation and Communication Complexity of Parallel SGD with Dynamic Batch Sizes for Stochastic Non-Convex Optimization](https://mlsanity.com/paper/1905.04346/) (2019)
- [Parallel Restarted SPIDER -- Communication Efficient Distributed Nonconvex Optimization with Optimal Computation Complexity](https://mlsanity.com/paper/1912.06036/) (2019)
- [Protecting Privacy through Distributed Computation in Multi-agent Decision Making](https://mlsanity.com/paper/1402.0589/) (2014)

## Built on by (most cited first, 30 of 43)

- [Generalized Federated Learning via Sharpness Aware Minimization](https://mlsanity.com/paper/2206.02618/) (2022)
- [Stochastic Controlled Averaging for Federated Learning with Communication Compression](https://mlsanity.com/paper/2308.08165/) (2023)
- [FedVARP: Tackling the Variance Due to Partial Client Participation in Federated Learning](https://mlsanity.com/paper/2207.14130/) (2022)
- [On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyond](https://mlsanity.com/paper/2206.05187/) (2022)
- [Momentum Benefits Non-IID Federated Learning Simply and Provably](https://mlsanity.com/paper/2306.16504/) (2023)
- [Faster Adaptive Federated Learning](https://mlsanity.com/paper/2212.00974/) (2022)
- [Communication-Efficient Federated Learning with Accelerated Client Gradient](https://mlsanity.com/paper/2201.03172/) (2022)
- [Accelerated Federated Learning with Decoupled Adaptive Optimization](https://mlsanity.com/paper/2207.07223/) (2022)
- [DASHA: Distributed Nonconvex Optimization with Communication Compression, Optimal Oracle Complexity, and No Client Synchronization](https://mlsanity.com/paper/2202.01268/) (2022)
- [On the Role of Server Momentum in Federated Learning](https://mlsanity.com/paper/2312.12670/) (2023)
- [The Limits and Potentials of Local SGD for Distributed Heterogeneous Learning with Intermittent Communication](https://mlsanity.com/paper/2405.11667/) (2024)
- [Taming Fat-Tailed ("Heavier-Tailed'' with Potentially Infinite Variance) Noise in Federated Learning](https://mlsanity.com/paper/2210.00690/) (2022)
- [Solving a Class of Non-Convex Minimax Optimization in Federated Learning](https://mlsanity.com/paper/2310.03613/) (2023)
- [Communication-Efficient Federated Bilevel Optimization with Local and Global Lower Level Problems](https://mlsanity.com/paper/2302.06701/) (2023)
- [Federated Learning with Manifold Regularization and Normalized Update Reaggregation](https://mlsanity.com/paper/2311.05924/) (2023)
- [Fast Adaptive Federated Bilevel Optimization](https://mlsanity.com/paper/2211.01122/) (2022)
- [DIAMOND: Taming Sample and Communication Complexities in Decentralized Bilevel Optimization](https://mlsanity.com/paper/2212.02376/) (2022)
- [When Decentralized Optimization Meets Federated Learning](https://mlsanity.com/paper/2306.02570/) (2023)
- [On Provable Benefits of Muon in Federated Learning](https://mlsanity.com/paper/2510.03866/) (2025)
- [FedBCD:Communication-Efficient Accelerated Block Coordinate Gradient Descent for Federated Learning](https://mlsanity.com/paper/2603.05116/) (2026)
- [FedLion: Faster Adaptive Federated Optimization with Fewer Communication](https://mlsanity.com/paper/2402.09941/) (2024)
- [Adaptive Federated Minimax Optimization with Lower Complexities](https://mlsanity.com/paper/2211.07303/) (2022)
- [Federated Conditional Stochastic Optimization](https://mlsanity.com/paper/2310.02524/) (2023)
- [FedSSO: A Federated Server-Side Second-Order Optimization Algorithm](https://mlsanity.com/paper/2206.09576/) (2022)
- [Momentum Approximation in Asynchronous Private Federated Learning](https://mlsanity.com/paper/2402.09247/) (2024)
- [Understanding Server-Assisted Federated Learning in the Presence of Incomplete Client Participation](https://mlsanity.com/paper/2405.02745/) (2024)
- [Federated Learning Survey: A Multi-Level Taxonomy of Aggregation Techniques, Experimental Insights, and Future Frontiers](https://mlsanity.com/paper/2511.22616/) (2025)
- [Faster Adaptive Momentum-Based Federated Methods for Distributed Composition Optimization](https://mlsanity.com/paper/2211.01883/) (2022)
- [FedDA: Faster Framework of Local Adaptive Gradient Methods via Restarted Dual Averaging](https://mlsanity.com/paper/2302.06103/) (2023)
- [Oracle Complexity Reduction for Model-free LQR: A Stochastic Variance-Reduced Policy Gradient Approach](https://mlsanity.com/paper/2309.10679/) (2023)

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