PaperlayerView exact revision on arXiv

Source-linked AI summary

High-Probability Convergence of SGD via Batched Updates

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

arXiv:2609.12765v1math.OCcs.LGcs.SYeess.SY

TL;DR

High-probability guarantees for SGD’s last iterate remain difficult, and prior approaches often rely on restrictive assumptions. This paper uses Batched SGD, averaging online samples before each update, and reports near-optimal high-probability results for convex and non-convex objectives as well as federated learning.

  • Problem

    Prior high-probability analyses often require bounded domains or gradients, assumptions that limit applicability to unconstrained problems and can fail even for simple quadratic objectives.

    Source: §I · p. 1

  • Method

    Batched SGD partitions online samples into batches, holds the parameter fixed while collecting gradients, and updates using a refined direction.

    Source: §I · p. 1, §III · p. 4

  • Results

    Batched SGD achieves near-optimal high-probability concentration bounds for smooth strongly convex and non-convex functions, and its federated extension gives logarithmic communication rounds, linear speedup, and no bias due to heterogeneity.

    Source: §I · p. 2, §VI · p. 8

  • Takeaways & Limitations

    Batching yields simpler high-probability proofs without restrictive assumptions or auxiliary sequences, and the batching idea extends naturally to federated learning.

    Source: §I · p. 2, §VI · p. 8

Abstract

from arXiv · show

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.

I. INTRODUCTION

High-probability guarantees for SGD’s last iterate remain difficult: prior analyses often impose restrictive boundedness assumptions or rely on complex proofs. The paper introduces Batched SGD and reports near-optimal rates for strongly convex and non-convex objectives, plus high-probability guarantees for federated learning.

I. INTRODUCTION

  • The last-iterate high-probability problem remains poorly understood: existing work often assumes bounded domains or gradients, or gives only averaged-iterate guarantees with complex analyses.Some last-iterate guarantees also rely on an auxiliary sequence.Source: §I · p. 1
  • Batched SGD averages samples within equal-sized batches and makes fewer updates using more precise, variance-reduced directions.The batch size is chosen according to properties of the objective.Source: §I · p. 1
  • The reported rates for smooth strongly convex and non-convex objectives hold with probability at least 1 − δ and match their in-expectation counterparts in their dependence on T.The strongly convex guarantee is directly for the last iterate; the dependence on 1/δ is logarithmic.Source: §I · pp. 1–2
  • For smooth, strongly convex federated losses, batching gives high-probability guarantees with logarithmic communication rounds, linear speedup in agent count, and no bias due to local-loss heterogeneity.The authors describe this as the first high-probability result in federated learning.Source: §I · p. 2
  • The stated novelty is not mini-batching itself, but that batching yields near-optimal high-probability rates with simpler proofs.Source: §I · p. 2

II. PROBLEM FORMULATION AND MOTIVATION

The section formulates a high-probability last-iterate error bound for SGD under smoothness and stochastic-gradient assumptions. It motivates the target by contrasting logarithmic dependence on the failure probability with the weaker bounds obtained by directly converting expectation guarantees.

II. PROBLEM FORMULATION AND MOTIVATION

  • Standard SGD queries a stochastic-gradient oracle at each online step and updates the parameter by subtracting the gradient estimate times the step size.The objective is assumed smooth; the section also introduces strong convexity and unbiased, norm-sub-Gaussian gradient noise assumptions.Source: §II · p. 2
  • Given T gradient queries and failure probability δ, the target is an error bound on f(x_T) − f(x*) that holds with probability at least 1 − δ.Source: §II · p. 2
  • A direct Markov- or Chebyshev-based conversion of an expectation bound yields polynomial dependence on 1/δ, whereas the desired high-probability bound scales logarithmically in 1/δ.Source: §II · p. 2
  • The one-step descent bound separates a descent term from a gradient-noise cross term and a squared-noise term, whose combined effect must be controlled.Under strong convexity, the descent term can be related to the function gap and induces contraction toward the optimum.Source: §II · p. 3

A. Challenges in the Analysis

The main obstacle is controlling the stochastic-gradient cross term without taking expectations or assuming bounded iterates or gradients. A direct Young’s-inequality argument creates a step-size-independent bias, while standard martingale concentration needs bounded increments.

A. Challenges in the Analysis

  • The cross term does not vanish in a high-probability analysis because expectations cannot be taken, even though it vanishes under conditional expectation for an in-expectation analysis.Source: A. Challenges in the Analysis · p. 3
  • Young’s inequality introduces a noise-squared term scaling linearly with the step size; unrolling the recursion leaves a step-size-independent bias that can make the bound vacuous.Source: A. Challenges in the Analysis · p. 3
  • Although the cross terms form a martingale-difference sequence, their iterate-dependent increments cannot be uniformly bounded as standard Azuma–Hoeffding concentration requires.Source: A. Challenges in the Analysis · p. 3
  • Prior work often assumes bounded gradients or domains, while approaches avoiding those assumptions commonly guarantee only averaged iterates or use involved proofs; batching is proposed to address these gaps.The authors state that their approach offers a simpler analysis, a direct last-iterate guarantee, and an extension to federated learning.Source: A. Challenges in the Analysis · p. 3

III. BATCHED SGD ALGORITHM

Batched SGD partitions the online samples into batches, holds the parameter fixed while gathering each batch’s gradients, then makes one update using their average. The refined direction lowers variance and supports a simpler high-probability analysis.

III. BATCHED SGD ALGORITHM

  • Batched SGD partitions T samples into K batches of H and makes one parameter update at the end of each batch.The illustration marks update locations with red bars.Source: §III · p. 4, Figure 1 · p. 4
  • During each batch, the parameter stays fixed while H stochastic gradients are collected and averaged into a refined gradient direction.The gradients are evaluated at the parameter value from the start of the batch.Source: §III · p. 4
  • The averaged direction has an H-fold lower variance than a single-sample SGD gradient, and the authors state that batching yields near-optimal high-probability guarantees with a simpler proof.Source: §III · p. 4

IV. MAIN RESULTS AND ANALYSIS

Batched SGD obtains high-probability convergence guarantees by using averaged batch gradients, whose variance reduction supports a simpler analysis. The results cover strongly convex and non-convex objectives, with logarithmically many updates sufficient in the strongly convex case.

IV. MAIN RESULTS AND ANALYSIS

  • Batching scales the statistical noise term down with batch size H; when H = 1, the resulting bound is vacuous, as in standard SGD.The optimization term decays exponentially in K, while the noise term decays linearly with H; balancing them motivates K = O(log T).Source: §IV · p. 4
  • O(log(T) log(dT/δ)/T) is the near-optimal high-probability rate established for Batched SGD in the strongly convex case.The result matches [7] up to a log d factor, with a significantly simpler proof.Source: §IV · p. 4
  • Theorem 1 establishes a high-probability guarantee on SGD’s last iterate; as far as the authors know, [7] is the only work achieving this bound, using a more involved analysis that exploits an auxiliary sequence.Source: §IV · p. 4
  • K = O(log T) parameter updates suffice for near-optimal high-probability bounds in the strongly convex case.The theorem selects K = 4κ log T and α = 1/(2L).Source: §IV · p. 4
  • Averaging independent norm-sub-Gaussian vectors reduces variance, and the analysis applies this property to Batched SGD's gradient estimates.The resulting gradient variance-reduction guarantee holds simultaneously across batches with probability at least 1 − δ.Source: §IV · p. 5
  • The high-probability guarantee also extends to general non-convex objectives, with a rate that incurs only an extra log d factor versus [5], [7].The proof is described as significantly simpler than those prior works.Source: §IV · pp. 5–6

V. EXTENSION TO THE FEDERATED SETTING

The paper extends its batching principle to federated learning by treating local computations between communication rounds as batches. The batch length sets communication frequency, and the approach yields logarithmic communication rounds with a simple high-probability analysis.

V. EXTENSION TO THE FEDERATED SETTING

  • Batched SGD's refined gradients reduce noise before updates, simplifying high-probability analysis and saving update rounds.The same batching principle is carried over to federated learning.Source: §V · p. 6
  • In federated learning, local gradient computations between successive communication rounds act as batches, and batch size H determines communication frequency.The method yields logarithmically many communication rounds and a simple high-probability analysis.Source: §V · p. 6
  • Batched FL translates the batching idea into logarithmic communication rounds for federated optimization.The analysis is high-probability.Source: §V · p. 6

A. Problem Formulation

The section formulates federated learning with a central server and agents that hold local objectives and online samples. It describes intermittent communication rounds and motivates high-probability analysis, under assumptions on gradients, noise, and cross-agent independence.

A. Problem Formulation

  • The federated setup has a central server and M agents, each associated with a local objective function.The standard FL problem is to solve a global optimization objective.Source: A. Problem Formulation · p. 6
  • Each agent's T online samples are split across K communication rounds of H steps, with the server broadcasting a global parameter before local updates and aggregating models afterward.The schedule satisfies T = KH.Source: A. Problem Formulation · p. 6
  • The section identifies a gap in non-trivial high-probability bounds for FL and assumes unbiased gradients, norm-sub-Gaussian noise, and independent agent stochastic processes.The independence assumption is described as standard and needed for collaborative speedups with the number of agents.Source: A. Problem Formulation · p. 6

B. Batched SGD for FL

Batched FL replaces multiple noisy local updates per communication round with one update using a refined direction from that round's samples. This keeps local models synchronized and is presented as a way to avoid heterogeneity bias while retaining collaborative benefits.

B. Batched SGD for FL

  • Multiple unsynchronized local updates using single-sample noisy gradients can cause client drift toward agent-specific optima, making aggregation biased.The bias hinders the benefits of collaboration.Source: B. Batched SGD for FL · p. 6
  • Batched FL constructs a refined gradient from the samples collected during a round, then uses it for one local-model update.This replaces multiple updates using highly noisy stochastic gradients.Source: B. Batched SGD for FL · p. 6
  • The figure depicts each agent partitioning T samples into K communication rounds, with red bars marking parameter updates and server aggregation.The method is called Batched FL and is summarized in Algorithm 2.Source: Figure 2 · p. 7, B. Batched SGD for FL · p. 7
  • Using one precise update with an H-fold lower-variance direction keeps local models synchronized, preventing heterogeneity bias while preserving collaborative benefits.For strongly convex objectives, the mechanism requires only logarithmically many communication rounds.Source: B. Batched SGD for FL · p. 7

C. Results And Analysis

Batched FL provides high-probability guarantees for strongly convex objectives, with a rate of order O(log(MT) log(dT/δ)/(MT)) even when local losses are arbitrarily heterogeneous. The analysis also reports linear speedup and logarithmic communication rounds for strongly convex FL, while non-convex FL retains linear speedup at a higher communication frequency.

Strongly Convex Objective

  • O(log(MT) log(dT/δ)/(MT)) is the reported high-probability rate for strongly convex Batched FL, even with arbitrarily heterogeneous local losses that need not be smooth or convex.The theorem concerns a strongly convex global objective.Source: C. Results And Analysis · pp. 7–8

Linear Speedup with No Heterogeneity Bias

  • A factor-M tighter rate gives linear speedup through agent collaboration without gradient-tracking or control-variate techniques.The comparison is between rates (16) and (36).Source: C. Results And Analysis · p. 8

Logarithmic Communication Rounds

  • K = O(log MT) communication rounds suffice for near-optimal strongly convex FL rates, reducing communication overhead.Batching determines communication frequency in the federated setting.Source: C. Results And Analysis · p. 8

Non-Convex Objective

  • Batched FL also has a high-probability guarantee for general non-convex objectives and an explicit linear speedup in the number of agents.The non-convex result requires a higher communication frequency than the strongly convex setting to achieve that speedup.Source: C. Results And Analysis · p. 8

VI. CONCLUSION

The paper introduces Batched SGD and establishes near-optimal high-probability concentration bounds for smooth strongly convex and non-convex functions. It also extends batching to federated learning and conjectures that the idea may yield simple concentration proofs for other stochastic approximation problems.

Conclusion

  • Batched SGD partitions the online stream into batches, forms low-variance gradient directions, and uses them for parameter updates; it yields near-optimal high-probability bounds for smooth strongly convex and non-convex functions.The proofs are significantly simpler than prior work and require no restrictive assumptions.Source: §VI · p. 8
  • The batching idea extends naturally to federated learning, and its use for simple concentration proofs in other stochastic approximation problems, including reinforcement learning, remains a conjecture.Formalizing that conjecture is ongoing work.Source: §VI · p. 8
Loading 2609.12765v1…