Stochastic Optimization and SGD — Mathematics Knowledge System

Overview

Stochastic optimization minimizes an objective function given as an expectation: min_x E_xi[f(x; xi)], where xi is a random variable. The stochastic gradient descent (SGD) algorithm replaces the exact gradient with a noisy estimate (stochastic gradient): x_{t+1} = x_t - alpha_t * g_t where g_t is an unbiased estimator of the gradient. SGD is the workhorse of machine learning. Its convergence analysis reveals the trade-off between step size (learning rate) and variance: too large => instability, too small => slow convergence. Variance reduction methods (SVRG, SAG, ADAM) improve practical convergence.

Intuition

In a machine learning problem, the loss is an average over n training examples: L(w) = (1/n) sum_{i=1}^n ell(w, x_i, y_i). Computing the full gradient costs O(n) per step -- too expensive for large n. SGD approximates the gradient by sampling a mini-batch of examples: g_t = (1/|B|) sum_{i in B_t} nabla ell(w_t, x_i, y_i). This is an unbiased estimate of nabla L(w_t). Although noisier, SGD can make many more steps per computation, often converging faster in practice. The noise also acts as implicit regularization.

Formal Definition

Definition

For convex f: R^d -> R with E[f] minimized at x*, SGD with step sizes alpha_t satisfying sum alpha_t = inf and sum alpha_t^2 < inf (Robbins-Monro conditions) converges to x* almost surely. For non-convex f (as in deep learning): E[||nabla f(x_t)||^2] -> 0 at rate 1/sqrt(T) under mild conditions. The update is x_{t+1} = x_t - alpha_t * g_t, where E[g_t | x_t] = nabla f(x_t) and Var(g_t) <= sigma^2. Mini-batch SGD with batch size B reduces variance by B: Var(g_t^B) = sigma^2/B.

x_{t+1} = x_t - \alpha_t g_t,\quad \mathbb{E}[g_t \mid x_t] = \nabla f(x_t)

SGD update rule

\sum_{t=1}^\infty \alpha_t = \infty,\quad \sum_{t=1}^\infty \alpha_t^2 < \infty

Robbins-Monro step size conditions

\mathbb{E}\left[f(x_T) - f(x^*)\right] \le \frac{\|x_1 - x^*\|^2 + \sigma^2 \sum_{t=1}^T \alpha_t^2}{2\sum_{t=1}^T \alpha_t}

SGD convergence bound (convex case)

\frac{1}{T}\sum_{t=1}^T \mathbb{E}\|\nabla f(x_t)\|^2 \le \frac{2(f(x_1)-f^*)}{\alpha T} + \alpha \sigma^2 L

SGD convergence (non-convex, constant step)

Notation

Notation	Meaning
$g_t$	Stochastic gradient at step t
$\alpha_t$	Step size (learning rate) at step t
$\sigma^2$	Variance of the stochastic gradient
$B$	Mini-batch size

Theorems

Theorem 1: SGD Convergence (Convex Case)

For a convex L-smooth function f with minimum value f*, SGD with constant step alpha satisfies: (1/T) sum_{t=1}^T E[f(x_t)] - f* <= L*||x_1 - x*||^2/(2T*alpha) + alpha*sigma^2/2. Setting alpha = ||x_1-x*||/(sigma*sqrt(T)) gives E[f(x_bar_T)] - f* <= ||x_1-x*||*sigma/sqrt(T) + O(1/T), i.e., O(1/sqrt(T)) convergence rate.

Theorem 2: SGD Convergence (Non-Convex)

For an L-smooth (possibly non-convex) function f, SGD with step size alpha satisfies (1/T)*sum_{t=0}^{T-1} E[||nabla f(x_t)||^2] <= 2*(f(x_0) - f*)/(alpha*T) + alpha*L*sigma^2. Setting alpha = 1/(sigma*sqrt(LT)) gives O(sigma/sqrt(T)) convergence to a stationary point. This is the standard guarantee for finding approximate critical points of non-convex losses.

Theorem 3: SVRG Variance Reduction

Stochastic Variance Reduced Gradient (SVRG): periodically compute the full gradient mu = nabla f(x_s) at a snapshot x_s. Then use the corrected gradient v_t = nabla f_i(x_t) - nabla f_i(x_s) + mu (reduced variance). SVRG achieves linear convergence for strongly convex functions: E[f(x_T) - f*] <= rho^T*(f(x_0) - f*) for some rho < 1. This improves over SGD's O(1/T) rate for strongly convex objectives.

Worked Examples

1
Full gradient: nabla L(w) = (2/n) sum_{i=1}^n x_i*(w^T x_i - y_i) -- costs O(nd) to compute.
2
SGD: sample a random index i ~ Uniform({1,...,n}). Stochastic gradient: g = 2 x_i*(w^T x_i - y_i).
$g = 2x_i(w^\top x_i - y_i)$
3
Update: w_{t+1} = w_t - alpha * g = w_t - 2*alpha*x_i*(w_t^T x_i - y_i).
$w_{t+1} = w_t - 2\alpha x_i(w_t^\top x_i - y_i)$
4
E[g | w_t] = nabla L(w_t): the stochastic gradient is unbiased. Cost per step: O(d) instead of O(nd).

✓ Answer

SGD for linear regression: sample one data point (x_i, y_i), compute gradient g = 2*x_i*(w^T*x_i - y_i), update w <- w - alpha*g. Cost O(d) per step vs O(nd) for full gradient.

Practice Problems

Mediumfree response

Explain the bias-variance tradeoff in choosing the mini-batch size B for SGD.

Common Mistakes

Common Mistake

Thinking a smaller learning rate always leads to faster or better convergence.

Too small a learning rate means the algorithm makes negligible progress per step, so convergence is slow in wall-clock time even if it's theoretically stable. The optimal step size balances the bias term (decreasing in alpha) and the variance term (increasing in alpha): for convex objectives, alpha ~ 1/sqrt(T) is optimal, not alpha -> 0 too aggressively. In practice, learning rate warmup and scheduling (cosine, linear decay) are crucial for good performance.

Quiz

The Robbins-Monro conditions for SGD convergence are:

alpha_t = 1/t and alpha_t constantsum alpha_t = inf and sum alpha_t^2 < infalpha_t -> 0 exponentiallyalpha_t = 1/sqrt(t) only

Historical Background

Robbins and Monro (1951) introduced stochastic approximation for root-finding. Kiefer and Wolfowitz (1952) extended to optimization. The modern ML usage traces to Rumelhart, Hinton, and Williams (1986) for training neural networks. The convergence theory was developed by Polyak (1987), Ruppert (1988), and later Bottou (1998), Shalev-Shwartz and Zhang (2013) for variance reduction. Adam (Kingma-Ba, 2014) is the most popular adaptive SGD variant in practice.

1951
Robbins-Monro stochastic approximation algorithm
Herbert Robbins, Sutton Monro
1986
Rumelhart-Hinton-Williams apply SGD to backpropagation in neural networks
David Rumelhart, Geoffrey Hinton, Ronald Williams
1998
Bottou analyzes SGD for large-scale learning
Leon Bottou
2014
Kingma-Ba introduce Adam optimizer
Diederik Kingma, Jimmy Ba

Summary

SGD: x_{t+1} = x_t - alpha_t * g_t with E[g_t | x_t] = nabla f(x_t); cost O(d) per step.
Convergence: O(1/sqrt(T)) for convex f; finds approximate stationary points for non-convex f.
Mini-batch size B: reduces variance by B but costs B times more per step.
Variance reduction (SVRG, SAG): achieves linear convergence for strongly convex objectives.

References

BookBottou, L., Curtis, F.E., and Nocedal, J. Optimization Methods for Large-Scale Machine Learning. SIAM Review, 2018.
BookShalev-Shwartz, S. and Ben-David, S. Understanding Machine Learning. Cambridge, 2014.
WebsiteWikipedia -- Stochastic gradient descent