Stochastic Gradient Descent Convergence

Why This Matters

Gradient descent computes the full gradient $\nabla F(w) = \frac{1}{n}\sum_{i=1}^n \nabla f_i(w)$ at each step, costing $O(n)$ per iteration. When $n$ is millions or billions (as in modern ML), this is too expensive. SGD replaces the full gradient with a stochastic estimate: the gradient of a single sample (or small mini-batch). The cost drops to $O(1)$ per iteration.

The trade-off: SGD's gradient estimate is noisy, which slows convergence relative to full gradient descent. But this noise also provides implicit regularization, helping SGD find solutions that generalize better. Understanding SGD convergence theory explains learning rate scheduling, batch size and learning dynamics, and why SGD remains the default optimizer for deep learning.

Setup

We minimize $F(w) = \mathbb{E}_{\xi}[f(w; \xi)]$ where $\xi$ is a random data sample. The theory relies on properties from convex optimization. The SGD update is:

$w_{t+1} = w_t - \eta_t g_t$

where $g_t = \nabla f(w_t; \xi_t)$ is a stochastic gradient satisfying $\mathbb{E}[g_t | w_t] = \nabla F(w_t)$ (unbiasedness).

Definition

Stochastic Gradient Oracle

A stochastic gradient oracle returns $g_t$ such that:

1. $\mathbb{E}[g_t | w_t] = \nabla F(w_t)$ (unbiased)
2. $\mathbb{E}[\|g_t - \nabla F(w_t)\|^2 | w_t] \leq \sigma^2$ (bounded variance)

The variance $\sigma^2$ measures the noise level. For a single sample from a finite dataset, $\sigma^2$ depends on the heterogeneity of the individual gradients $\nabla f_i(w)$.

Main Theorems

Theorem

SGD Convergence for Convex Functions

Statement

For convex, $L$-smooth $F$ with bounded gradient variance $\sigma^2$, SGD with learning rate $\eta = \frac{1}{L\sqrt{T}}$ satisfies:

$\mathbb{E}[F(\bar{w}_T)] - F(w^*) \leq \frac{L\|w_0 - w^*\|^2}{2\sqrt{T}} + \frac{\sigma^2}{2L\sqrt{T}}$

where $\bar{w}_T = \frac{1}{T}\sum_{t=0}^{T-1} w_t$ is the average iterate. The convergence rate is $O(1/\sqrt{T})$.

Intuition

The first term is the deterministic convergence rate (how fast GD would converge if there were no noise). The second term is the noise penalty. Both decrease as $O(1/\sqrt{T})$. Compared to full GD, which converges as $O(1/T)$ for smooth convex functions, SGD is slower by a factor of $\sqrt{T}$. This is the price of using cheap, noisy gradients.

Proof Sketch

Start with the smoothness inequality: $F(w_{t+1}) \leq F(w_t) + \langle \nabla F(w_t), w_{t+1} - w_t \rangle + \frac{L}{2}\|w_{t+1} - w_t\|^2$. Substitute $w_{t+1} - w_t = -\eta g_t$ and take expectations. The cross term $\mathbb{E}[\langle \nabla F(w_t), -\eta g_t \rangle] = -\eta \|\nabla F(w_t)\|^2$ (by unbiasedness). The quadratic term gives $\frac{L\eta^2}{2}(\|\nabla F(w_t)\|^2 + \sigma^2)$. Combine, use convexity to relate $\|\nabla F(w_t)\|^2$ to $F(w_t) - F(w^*)$, telescope, and choose $\eta$ to balance terms.

Why It Matters

This $O(1/\sqrt{T})$ rate is optimal for first-order stochastic methods on convex functions. No algorithm using only stochastic gradient queries can do better in the worst case (Nemirovsky and Yudin, 1983). The rate determines how many epochs (passes over the data) you need for a given accuracy.

Failure Mode

The bound requires a constant learning rate tuned to $T$ (the total number of iterations), which must be known in advance. In practice, this is handled by decaying the learning rate. If the learning rate does not decay, SGD oscillates around the optimum with radius proportional to $\eta \sigma$.

Theorem

SGD Convergence for Strongly Convex Functions

Statement

For $\mu$-strongly convex, $L$-smooth $F$ with bounded gradient variance $\sigma^2$, SGD with learning rate $\eta_t = \frac{2}{\mu(t+1)}$ satisfies:

$\mathbb{E}[F(w_T)] - F(w^*) \leq \frac{2\sigma^2}{\mu T} + \frac{2L\|w_0 - w^*\|^2}{T^2}$

The dominant term is $O(\sigma^2 / (\mu T))$, giving a convergence rate of $O(1/T)$.

Intuition

Strong convexity provides curvature that pulls iterates toward $w^*$ more aggressively. This improves the rate from $O(1/\sqrt{T})$ to $O(1/T)$. The condition number $\kappa = L/\mu$ does not appear in the leading term, but $\mu$ does: stronger curvature means faster convergence.

Proof Sketch

Use the strong convexity inequality: $F(w^*) \geq F(w_t) + \langle \nabla F(w_t), w^* - w_t \rangle + \frac{\mu}{2}\|w^* - w_t\|^2$. Combined with the SGD update and taking expectations, this gives a recurrence on $\mathbb{E}[\|w_t - w^*\|^2]$. With the decaying learning rate $\eta_t = 2/(\mu(t+1))$, the recurrence solves to $O(\sigma^2/(\mu T))$.

Why It Matters

This $O(1/T)$ rate matches the minimax-optimal rate for strongly convex stochastic optimization. For the special case of least squares regression with $n$ data points, this means $O(n/T)$ suboptimality per epoch, so a constant number of passes over the data suffices for a fixed accuracy level.

Failure Mode

The rate requires knowing $\mu$ to set the learning rate. If $\mu$ is overestimated, the learning rate decays too fast and convergence stalls. If $\mu$ is underestimated, the learning rate stays too large and the iterates oscillate. Adaptive methods (AdaGrad, Adam) avoid this by adjusting rates per-coordinate.

Mini-Batch SGD

Using a mini-batch of size $B$ reduces variance: if $g_t^{(B)} = \frac{1}{B}\sum_{j=1}^B \nabla f(w_t; \xi_t^{(j)})$, then:

$\text{Var}(g_t^{(B)}) = \frac{\sigma^2}{B}$

This reduces the noise term in the convergence bound by a factor of $B$. But each iteration now costs $B$ times as much compute. The total work to achieve error $\epsilon$ is:

$\text{Work} = B \times T \propto B \times \frac{\sigma^2}{B\epsilon^2} = \frac{\sigma^2}{\epsilon^2}$

for the convex case. The total work is independent of $B$. Mini-batches do not reduce total computation for convex problems; they trade fewer iterations for more work per iteration. The real benefit is parallelism: $B$ gradient computations can run simultaneously on a GPU.

Learning Rate Schedules

Definition

Robbins-Monro Conditions

A learning rate sequence $\{\eta_t\}$ satisfies the Robbins-Monro conditions if:

$\sum_{t=1}^{\infty} \eta_t = \infty \quad \text{and} \quad \sum_{t=1}^{\infty} \eta_t^2 < \infty$

The first condition ensures the iterates can reach any point. The second ensures the noise is eventually damped. Examples: $\eta_t = c/t$ satisfies both; $\eta_t = c/\sqrt{t}$ satisfies the first but not the second.

Common schedules in practice:

• Constant then decay. Train with a constant $\eta$ for most of training, then decay (linear or cosine) in the final phase.
• Cosine annealing. $\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})(1 + \cos(\pi t / T))$. Smooth decay that spends more time at low learning rates.
• Linear warmup. Start with a small $\eta$ and increase linearly for the first few thousand steps, then follow one of the above schedules. This stabilizes early training when parameter initialization is far from the eventual scale.

Stationary Points in Nonconvex SGD

For nonconvex objectives, SGD convergence is typically stated in terms of stationarity rather than global optimality. The relevant distinction:

Definition

First-Order Stationary Point (FOSP)

A first-order stationary point satisfies $\nabla F(w) = 0$. An $\epsilon$-FOSP satisfies $\|\nabla F(w)\| \leq \epsilon$. Local minima, local maxima, and saddle points are all FOSPs.

Definition

Second-Order Stationary Point (SOSP)

A second-order stationary point satisfies $\nabla F(w) = 0$ and $\nabla^2 F(w) \succeq 0$ (no negative Hessian eigenvalue). An $\epsilon$-SOSP adds $\lambda_{\min}(\nabla^2 F(w)) \geq -\sqrt{\rho \epsilon}$ for Hessian-Lipschitz constant $\rho$. Saddle points are FOSPs but not SOSPs.

Standard nonconvex SGD analyses guarantee only an $\epsilon$-FOSP in $O(1/\epsilon^4)$ stochastic gradient queries. Jin, Ge, Netrapalli, Kakade, Jordan (2017), "How to Escape Saddle Points Efficiently", showed that perturbed gradient descent finds an $\epsilon$-SOSP in $\tilde{O}(1/\epsilon^2)$ iterations (deterministic setting). In the stochastic setting, analogous results (Jin et al., 2021) give polynomial rates for SGD to find approximate SOSPs under Hessian-Lipschitz assumptions.

Noise as Implicit Regularization

SGD noise is not purely harmful. Several observed benefits:

Escaping saddle points. Gradient noise helps SGD escape strict saddle points (where the Hessian has a negative eigenvalue, so they are FOSPs but not SOSPs). Full GD can get stuck at saddle points; SGD almost surely does not.

Flat minima preference. Empirical evidence suggests SGD converges to flatter minima (minima with smaller Hessian eigenvalues), which generalize better. Larger learning rates and smaller batch sizes increase noise, pushing SGD toward flatter regions.

Implicit bias of SGD. For linear models, SGD with small initialization converges to the minimum-norm solution. For matrix factorization problems, it converges to low-rank solutions. These implicit biases help generalization without explicit regularization.

Common Confusions

Watch Out

SGD is not the same as mini-batch GD with B equals n

When $B = n$ (full batch), you recover deterministic GD, not SGD. The convergence rates and implicit regularization properties are qualitatively different. Full-batch GD converges as $O(1/T)$ for smooth convex functions vs $O(1/\sqrt{T})$ for SGD, but SGD's noise provides regularization benefits that full-batch GD lacks.

Watch Out

Adam is not SGD

Adam uses adaptive per-coordinate learning rates and momentum. Its convergence theory is different from SGD. Adam can diverge on simple convex problems (Reddi et al., 2018). The AMSGrad fix addresses this, but in practice, Adam often works well despite the theoretical gap.

Exercises

ExerciseCore

Problem

You are training a model with SGD on a $\mu$-strongly convex loss with $\mu = 0.01$, $\sigma^2 = 1$, and you want $\mathbb{E}[F(w_T) - F(w^*)] \leq 0.001$. Using the strongly convex rate $O(\sigma^2 / (\mu T))$, approximately how many iterations $T$ do you need?

Hint

Reveal Solution

ExerciseAdvanced

Problem

Prove that for SGD with constant learning rate $\eta$ on a $\mu$-strongly convex, $L$-smooth function, the iterates do not converge to $w^*$ but instead oscillate in a ball of radius $O(\eta \sigma^2 / \mu)$ around $w^*$.

Hint

Reveal Solution

References

Canonical:

• Nemirovsky & Yudin, Problem Complexity and Method Efficiency in Optimization (1983)
• Robbins & Monro, "A Stochastic Approximation Method" (1951)

Current:

• Bottou, Curtis, Nocedal, "Optimization Methods for Large-Scale Machine Learning" (SIAM Review, 2018)

• Shalev-Shwartz & Ben-David, Understanding Machine Learning (2014), Chapter 14

• Jin, Ge, Netrapalli, Kakade, Jordan, "How to Escape Saddle Points Efficiently" (2017), ICML; arXiv:1703.00887. Perturbed GD finds an $\epsilon$-SOSP in $\tilde{O}(1/\epsilon^2)$ iterations.

• Boyd & Vandenberghe, Convex Optimization (2004), Chapters 2-5

• Nesterov, Introductory Lectures on Convex Optimization (2004), Chapters 1-3