# Sharper Rates and Flexible Framework for Nonconvex SGD with Client and Data Sampling

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 27 references that are in the map.
- First posted on arXiv 5 Jun 2022.

## Builds on (most cited first, 17 of 17)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Learning: Strategies for Improving Communication Efficiency](https://mlsanity.com/paper/1610.05492/) (2016)
- [SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives](https://mlsanity.com/paper/1407.0202/) (2014)
- [Lower Bounds for Non-Convex Stochastic Optimization](https://mlsanity.com/paper/1912.02365/) (2019)
- [SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimator](https://mlsanity.com/paper/1807.01695/) (2018)
- [SARAH: A Novel Method for Machine Learning Problems Using Stochastic Recursive Gradient](https://mlsanity.com/paper/1703.00102/) (2017)
- [PAGE: A Simple and Optimal Probabilistic Gradient Estimator for Nonconvex Optimization](https://mlsanity.com/paper/2008.10898/) (2020)
- [Stochastic Nested Variance Reduction for Nonconvex Optimization](https://mlsanity.com/paper/1806.07811/) (2018)
- [Non-convex Finite-Sum Optimization Via SCSG Methods](https://mlsanity.com/paper/1706.09156/) (2017)
- [Permutation Compressors for Provably Faster Distributed Nonconvex Optimization](https://mlsanity.com/paper/2110.03300/) (2021)
- [Accelerating Minibatch Stochastic Gradient Descent using Stratified Sampling](https://mlsanity.com/paper/1405.3080/) (2014)
- [Adaptivity of Stochastic Gradient Methods for Nonconvex Optimization](https://mlsanity.com/paper/2002.05359/) (2020)
- [Nonconvex Variance Reduced Optimization with Arbitrary Sampling](https://mlsanity.com/paper/1809.04146/) (2018)
- [L-SVRG and L-Katyusha with Arbitrary Sampling](https://mlsanity.com/paper/1906.01481/) (2019)
- [FL_PyTorch: optimization research simulator for federated learning](https://mlsanity.com/paper/2202.03099/) (2022)

## Built on by (most cited first, 10 of 10)

- [Can 5th Generation Local Training Methods Support Client Sampling? Yes!](https://mlsanity.com/paper/2212.14370/) (2022)
- [Improving Accelerated Federated Learning with Compression and Importance Sampling](https://mlsanity.com/paper/2306.03240/) (2023)
- [Federated Learning under Periodic Client Participation and Heterogeneous Data: A New Communication-Efficient Algorithm and Analysis](https://mlsanity.com/paper/2410.23131/) (2024)
- [Freya PAGE: First Optimal Time Complexity for Large-Scale Nonconvex Finite-Sum Optimization with Heterogeneous Asynchronous Computations](https://mlsanity.com/paper/2405.15545/) (2024)
- [Correlated Quantization for Faster Nonconvex Distributed Optimization](https://mlsanity.com/paper/2401.05518/) (2024)
- [SPAM: Stochastic Proximal Point Method with Momentum Variance Reduction for Non-convex Cross-Device Federated Learning](https://mlsanity.com/paper/2405.20127/) (2024)
- [Tighter Performance Theory of FedExProx](https://mlsanity.com/paper/2410.15368/) (2024)
- [BurTorch: Revisiting Training from First Principles by Coupling Autodiff, Math Optimization, and Systems](https://mlsanity.com/paper/2503.13795/) (2025)
- [Convergence Analysis of the ProbAbilistic Gradient Estimator Algorithm for Weakly Convex Finite-Sum Optimization](https://mlsanity.com/paper/2509.00737/) (2025)
- [Theoretical Foundations of Communication-Efficient, Robust, and Practical Distributed and Federated Optimization](https://mlsanity.com/paper/2608.06563/) (2026)

## Abstract (arXiv)

We revisit the classical problem of finding an approximately stationary point of the average of n smooth and possibly nonconvex functions. The optimal complexity of stochastic first-order methods in terms of the number of gradient evaluations of individual functions is Oleft(n + n1/2varepsilon-1right), attained by the optimal SGD methods smallsfcolorgreenSPIDER(arXiv:1807.01695) and smallsfcolorgreenPAGE(arXiv:2008.10898), for example, where varepsilon is the error tolerance. However, i) the big-O notation hides crucial dependencies on the smoothness constants associated with the functions, and ii) the rates and theory in these methods assume simplistic sampling mechanisms that do not offer any flexibility. In this work we remedy the situation. First, we generalize the smallsfcolorgreenPAGE algorithm so that it can provably work with virtually any (unbiased) sampling mechanism. This is particularly useful in federated learning, as it allows us to construct and better understand the impact of various combinations of client and data sampling strategies. Second, our analysis is sharper as we make explicit use of certain novel inequalities that capture the intricate interplay between the smoothness constants and the sampling procedure. Indeed, our analysis is better even for the simple sampling procedure analyzed in the smallsfcolorgreenPAGE paper. However, this already improved bound can be further sharpened by a different sampling scheme which we propose. In summary, we provide the most general and most accurate analysis of optimal SGD in the smooth nonconvex regime. Finally, our theoretical findings are supposed with carefully designed experiments.