math.OCSep 14, 2026

High-Probability Convergence of SGD via Batched Updates

Authors: Feng Zhu, Robert W. Heath Jr., Aritra Mitra

Abstract

Stochastic gradient descent (SGD) is the primary workhorse for large-scale optimization. While the average behavior of its iterates, typically characterized by mean-squared error bounds, is well-understood, obtaining high-probability guarantees for the last iterate remains challenging. Prior approaches to this problem have either imposed restrictive assumptions (such as bounded domains or gradients) or relied on complex proofs involving auxiliary sequences. In this work, we propose Batched SGD, a simple variant that partitions online samples into epochs and performs a single update per epoch using a refined, low-variance gradient estimate. Our main contribution demonstrates that this batching mechanism enables a surprisingly simple high-probability analysis that avoids both restrictive assumptions and auxiliary sequences. Under standard smoothness and norm-sub-Gaussian noise assumptions, we establish near-optimal rates for both strongly convex and non-convex objectives. Furthermore, we show that our batching idea extends naturally to federated learning (FL). We provide the first high-probability guarantees for FL, achieving logarithmic communication complexity, linear speedup in the number of agents, and resilience to data heterogeneity.

Explore similar work

CardsList
May 11, 2026cs.LG

Unveiling High-Probability Generalization in Decentralized SGD

Decentralized stochastic gradient descent (D-SGD) is an efficient method for large-scale distributed learning. Existing generalization studies mainly address expected results, achieving rates limited to O(1δmn)\mathcal{O}\left(\frac{1}{δ\sqrt{mn}}\right), where δδ is the confidence parameter, mm the number of workers, and nn the sample size. When m=1m=1, D-SGD reduces to traditional SGD, whose optimal high-probability generalization bound is O(1nlog⁡(1/δ))\mathcal{O}\left(\frac{1}{\sqrt{n}}\log (1/δ)\right). This discrepancy reveals a gap between high-probability guarantees for SGD and those for D-SGD. To close this, we develop a high-probability learning theory for D-SGD, aiming for the optimal O(1mnlog⁡(1/δ))\mathcal{O}\left(\frac{1}{\sqrt{mn}}\log (1/δ)\right) rate. We refine bounds for D-SGD using pointwise uniform stability in distributed learning-a weaker notion than uniform stability-and analyze them across convex, strongly convex, and non-convex settings. We also provide high-probability results for gradient-based measures in non-convex cases where only local minima exist, and derive optimization error and excess risk bounds. Finally, accounting for communication overhead, we analyze generalization bounds for local models within time-varying frameworks.
Jiahuan Wang, Ping Luo, Ziqing Wen +2
Stochastic Gradient DescentGeneralization BoundsDistributional LearningOptimal Sample Complexity
AskPDF
Jul 16, 2026cs.LG

What's in a Smoothness Constant? Tighter Rates for Local SGD with Bounded Second-order Heterogeneity

Local SGD, also known as Federated Averaging, is a widely used distributed optimization algorithm. Although Local SGD often outperforms alternatives such as Mini-batch SGD in practice, theory still only partially explains when and why local updates help under realistic data heterogeneity. Recent work by [Patel et al., 2025] shows that a bounded second-order heterogeneity assumption captures the efficiency of Local SGD for strongly convex objectives, and conjectures that the same principle extends to the general convex setting. In this paper, we prove this conjecture by establishing an improved convergence guarantee for Local SGD on general convex objectives under bounded second-order heterogeneity. We also improve the best-known lower bounds for Local SGD in this setting, showing that our upper bounds are nearly tight. Together, these results provide a sharper, more fine-grained convergence theory for Local SGD. As a further application of our techniques, we provide a lower bound for serial SGD with replacement, showing how second-order heterogeneity captures the impact of rare high-curvature clients.
Kumar Kshitij Patel, Rustem Islamov, Sebastian U Stich +3
Stochastic Gradient DescentConvex OptimizationOptimization Convergence AnalysisSmoothingFederated AveragingHeterogeneity
AskPDF
May 19, 2026stat.ML

Factor Augmented High-Dimensional SGD

Stochastic gradient descent (SGD) is a fundamental optimization algorithm widely used in modern machine learning. In this paper, we propose Factor-Augmented SGD (FSGD), a new optimization method that leverages latent factor representations in high-dimensional learning tasks. Unlike standard two-stage dimension reduction approaches that rely on offline representation learning and full data storage, a key novelty of FSGD is that it operates purely on streaming data, making it scalable to large-scale and high-dimensional problems. Furthermore, we establish the first theoretical framework that explicitly incorporates latent factor estimation error into the analysis of SGD, and provide moment convergence in ℓs\ell^s norm under decaying step sizes and mini-batch updates. Our results provide a new foundation for employing SGD reliably and scalably in high-dimensional machine learning systems.
Shubo Li, Yuefeng Han, Xiufan Yu
Stochastic Gradient DescentDp-SgdHigh-DimensionalLatent Factor RepresentationsSingular Value DecompositionFactor
AskPDF