Batched updates improve stochastic gradient descent with high confidence
High-Probability Convergence of SGD via Batched Updates
Machine Learning
Summary
Stochastic gradient descent (SGD) is a key method used to train machine learning models, but it’s hard to guarantee its performance with strong confidence just from the last update. The authors introduce a simple change called Batched SGD, which groups data samples into chunks and updates the model once per chunk using more precise estimates. This approach helps prove that the method works well with high probability, without needing strict assumptions about the problem. They also show that this idea can improve federated learning, where many devices cooperate to train models efficiently and reliably.
What this means in practice
- •For machine learning engineers: Implement batched SGD to achieve reliable model training with proven high-probability guarantees under realistic conditions without restrictive assumptions.
- •For federated learning system builders: Design federated learning systems that reduce communication rounds and handle diverse data distributions while assuring strong convergence confidence.
Authors
Feng Zhu, Robert W. Heath, 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.