Gradient Descent

The Optimization Problem

Machine learning reduces to empirical risk minimization (ERM): finding parameters $\theta$ that minimize a loss function averaged over a training dataset $\mathcal{D} = \{(x_i, y_i)\}_{i=1}^N$:

$\theta^* = \arg\min_\theta \; \mathcal{L}(\theta) = \arg\min_\theta \; \frac{1}{N} \sum_{i=1}^{N} \ell(f_\theta(x_i), y_i)$

For most models of interest (neural networks), $\mathcal{L}(\theta)$ is non-convex and high-dimensional (billions of parameters), ruling out closed-form solutions. Gradient descent solves this iteratively by moving in the direction of steepest descent.

Vanilla Gradient Descent

Starting from an initial $\theta_0$, the update rule computes the gradient over the full dataset and takes a step in the negative gradient direction:

$\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t) = \theta_t - \frac{\eta}{N} \sum_{i=1}^{N} \nabla_\theta \ell(f_{\theta_t}(x_i), y_i)$

where $\eta > 0$ is the learning rate (step size).

**Geometric interpretation.** At each point $\theta_t$, the gradient $\nabla \mathcal{L}(\theta_t)$ is a vector pointing in the direction of steepest *increase* of $\mathcal{L}$. Moving in the direction $-\nabla \mathcal{L}$ is the locally optimal descent direction. However, "locally optimal" can be misleading: the gradient only gives a good direction for infinitesimally small steps. For finite step sizes, the curvature of $\mathcal{L}$ matters.

Convergence Analysis

A differentiable function $\mathcal{L}$ has **$L$-Lipschitz continuous gradients** (is $L$-smooth) if:

$\|\nabla \mathcal{L}(\theta) - \nabla \mathcal{L}(\theta')\| \leq L \|\theta - \theta'\| \quad \forall \, \theta, \theta'$

Equivalently, the loss is bounded by a quadratic around any point:

$\mathcal{L}(\theta') \leq \mathcal{L}(\theta) + \nabla \mathcal{L}(\theta)^\top (\theta' - \theta) + \frac{L}{2}\|\theta' - \theta\|^2$

The constant $L$ is the largest eigenvalue of the Hessian $\nabla^2 \mathcal{L}$.

For an $L$-smooth convex function, gradient descent with step size $\eta = 1/L$ satisfies:

$\mathcal{L}(\theta_T) - \mathcal{L}(\theta^*) \leq \frac{L \|\theta_0 - \theta^*\|^2}{2T}$

This is an $O(1/T)$ rate: to achieve $\epsilon$-accuracy, we need $T = O(L \|\theta_0 - \theta^*\|^2 / \epsilon)$ iterations.

If $\mathcal{L}$ is additionally **$\mu$-strongly convex** ($\nabla^2 \mathcal{L} \succeq \mu I$), gradient descent with $\eta = 1/L$ converges linearly:

$\mathcal{L}(\theta_T) - \mathcal{L}(\theta^*) \leq \left(1 - \frac{\mu}{L}\right)^T \left(\mathcal{L}(\theta_0) - \mathcal{L}(\theta^*)\right)$

The condition number $\kappa = L/\mu$ determines convergence speed. Poorly conditioned problems ($\kappa \gg 1$) converge slowly because the loss surface is elongated: the gradient oscillates across the narrow direction while making slow progress along the wide direction.

Stochastic Gradient Descent (SGD)

SGD approximates the full gradient with a mini-batch $\mathcal{B} \subset \mathcal{D}$ of size $B = |\mathcal{B}|$, sampled uniformly:

$\theta_{t+1} = \theta_t - \eta_t \cdot g_t, \quad g_t = \frac{1}{B} \sum_{i \in \mathcal{B}_t} \nabla_\theta \ell(f_{\theta_t}(x_i), y_i)$

Properties of the stochastic gradient $g_t$:

1. Unbiased: $\mathbb{E}_{\mathcal{B}}[g_t] = \nabla \mathcal{L}(\theta_t)$
2. Bounded variance: $\mathbb{E}\|g_t - \nabla \mathcal{L}(\theta_t)\|^2 \leq \sigma^2 / B$, where $\sigma^2$ is the per-sample gradient variance
3. Computational cost: $O(B)$ per step instead of $O(N)$

For an $L$-smooth function with stochastic gradients having variance bounded by $\sigma^2$, SGD with step size $\eta = c / \sqrt{T}$ satisfies:

$\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\|\nabla \mathcal{L}(\theta_t)\|^2 \leq O\!\left(\frac{\sigma}{\sqrt{T}} + \frac{1}{T}\right)$

This is an $O(1/\sqrt{T})$ rate to a stationary point (where $\nabla \mathcal{L} \approx 0$). Note: for non-convex functions, a stationary point may be a saddle point or local minimum, not necessarily a global minimum.

**SGD noise as implicit regularization.** The stochastic noise in mini-batch gradients serves as implicit regularization. Empirical and theoretical evidence suggests:

1. SGD noise helps escape sharp minima (high curvature) in favor of flat minima (low curvature), which tend to generalize better ([?keskar2017large]).
2. The noise scale is proportional to $\eta / B$ (learning rate divided by batch size), which is why the linear scaling rule ([?goyal2017accurate]) prescribes scaling $\eta$ proportionally to $B$.
3. Large-batch training (small noise) can converge to sharp minima that generalize poorly, unless carefully tuned with warmup and learning rate schedules.

Momentum

SGD oscillates in directions with high curvature (large eigenvalues of the Hessian) while making slow progress along directions with low curvature. Momentum fixes this by accumulating a velocity vector:

$v_{t+1} = \mu \, v_t + g_t, \qquad \theta_{t+1} = \theta_t - \eta \, v_{t+1}$

where $\mu \in [0, 1)$ is the momentum coefficient (typically $\mu = 0.9$). The velocity $v_t$ is an exponential moving average of past gradients with effective window size $\sim 1/(1-\mu)$.

**Physical analogy.** Think of a ball rolling down the loss surface. Without momentum, the ball stops immediately when the slope changes direction (oscillation). With momentum, the ball has inertia: it accelerates along consistent downhill directions and dampens oscillations. On a quadratic with condition number $\kappa$, optimal momentum $\mu^* = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1)$ improves convergence from $O(\kappa \log 1/\epsilon)$ to $O(\sqrt{\kappa} \log 1/\epsilon)$ iterations.

Nesterov accelerated gradient evaluates the gradient at the "look-ahead" position $\theta_t - \eta \mu v_t$:

$v_{t+1} = \mu \, v_t + \nabla \mathcal{L}(\theta_t - \eta \mu v_t), \qquad \theta_{t+1} = \theta_t - \eta \, v_{t+1}$

Nesterov momentum achieves the optimal $O(1/T^2)$ convergence rate for smooth convex functions, compared to $O(1/T)$ for vanilla GD ([?nesterov1983method]). In practice, the improvement over classical momentum is often modest for deep learning.

Adam (Adaptive Moment Estimation)

**Adam** [@kingma2015adam] combines momentum with per-parameter adaptive learning rates by tracking exponential moving averages of the first and second moments of the gradient:

$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t \qquad \text{(first moment / mean)}$

$v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \qquad \text{(second moment / uncentered variance)}$

Bias-corrected estimates (since $m_0 = v_0 = 0$):

$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$

Parameter update:

$\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$

Default hyperparameters: $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$.

**Why Adam works.** The effective learning rate for parameter $\theta_i$ at step $t$ is $\eta / (\sqrt{\hat{v}_{t,i}} + \epsilon)$:

• Parameters with large gradient variance (noisy signal) have large $\hat{v}_{t,i}$, so their effective learning rate is smaller -- Adam is cautious where the signal is noisy.
• Parameters with small gradient variance (consistent signal) have small $\hat{v}_{t,i}$, so their effective learning rate is larger -- Adam makes confident updates.
• The bias correction is essential in the first few steps: without it, $m_t$ and $v_t$ are biased toward zero because they are initialized at zero and the exponential average has not converged.

**AdamW (decoupled weight decay)** [@loshchilov2019adamw]. In standard Adam, L2 regularization adds $\lambda \theta$ to the gradient, which then gets divided by $\sqrt{\hat{v}_t}$. This means the effective regularization strength varies per parameter -- defeating the purpose of uniform weight decay. AdamW fixes this by applying weight decay *directly* to the parameters, outside the adaptive learning rate:

$\theta_{t+1} = (1 - \eta \lambda) \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$

AdamW is the standard optimizer for transformer training. Typical hyperparameters: $\eta \in [10^{-5}, 10^{-3}]$, $\beta_1 = 0.9$, $\beta_2 \in [0.95, 0.999]$, weight decay $\lambda \in [0.01, 0.1]$.

Other Optimizers

Optimizer	Update Rule (Simplified)	Memory	Key Property
SGD	$\theta \leftarrow \theta - \eta g$	$O(P)$	Simple, good generalization
SGD + Momentum	$v \leftarrow \mu v + g$; $\theta \leftarrow \theta - \eta v$	$O(2P)$	Dampens oscillations
Adam	Uses $m_t, v_t$ as above	$O(3P)$	Adaptive per-parameter rates
AdamW	Adam + decoupled weight decay	$O(3P)$	Standard for transformers
Adafactor	Factored second moments	$O(P + \text{rows} + \text{cols})$	Memory-efficient for large matrices
LAMB	Adam + layerwise LR scaling	$O(3P)$	Large-batch training
Lion	Sign-based momentum update	$O(2P)$	Memory-efficient, empirically strong
Muon	Orthogonalized momentum	$O(2P)$	Emerging; strong empirical results

Learning Rate Schedules

The learning rate $\eta_t$ is rarely constant. A schedule adapts it over training:

Schedule	Formula	Typical Use
Constant	$\eta_t = \eta_0$	Baselines, short fine-tuning
Step decay	$\eta_t = \eta_0 \cdot \gamma^{\lfloor t/s \rfloor}$	CNNs (ResNet); $\gamma = 0.1$, $s$ at 30/60/90 epochs
Cosine annealing	$\eta_t = \eta_{\min} + \frac{1}{2}(\eta_0 - \eta_{\min})(1 + \cos(\pi t / T))$	Transformer pretraining ([?loshchilov2017sgdr])
Linear warmup + cosine	$\eta_t = \eta_0 \cdot \min(t/w, \; \frac{1}{2}(1+\cos(\pi(t-w)/(T-w))))$	Standard for LLM pretraining
Warmup + inverse sqrt	$\eta_t = \eta_0 \cdot \min(t/w, \; \sqrt{w/t})$	Original Transformer (Vaswani et al., 2017)
WSD (Warmup-Stable-Decay)	Warmup $\to$ constant $\to$ cosine decay	Practical for unknown training length

**Why warmup?** In the first few steps, Adam's second-moment estimates $\hat{v}_t$ are based on very few samples and are unreliable. Large learning rates during this phase cause wild updates. Linear warmup (ramping $\eta$ from 0 to $\eta_0$ over $w$ steps) gives the adaptive estimates time to stabilize. Empirically, $w$ = 1--5% of total training steps works well.

Gradient Clipping

To prevent exploding gradients, clip the gradient norm before applying the update:

$g_t \leftarrow g_t \cdot \min\!\left(1, \; \frac{c}{\|g_t\|}\right)$

where $c$ is the maximum allowed gradient norm (typically $c = 1.0$). This preserves the gradient direction but bounds its magnitude.

Gradient clipping is essential for training transformers and RNNs. Without it, a single batch with an unusually large loss can produce a gradient spike that destabilizes all the optimizer state (Adam's running averages get corrupted). A typical setting is `max_grad_norm=1.0` in both PyTorch and JAX training loops.

Convergence Rate Summary

Setting	Algorithm	Rate	Metric
$L$-smooth convex	GD ($\eta = 1/L$)	$O(1/T)$	$\mathcal{L}(\theta_T) - \mathcal{L}^*$
$L$-smooth, $\mu$-strongly convex	GD ($\eta = 1/L$)	$O((1-\mu/L)^T)$	$\mathcal{L}(\theta_T) - \mathcal{L}^*$
$L$-smooth convex	Nesterov	$O(1/T^2)$	$\mathcal{L}(\theta_T) - \mathcal{L}^*$
$L$-smooth convex	SGD ($\eta \propto 1/\sqrt{T}$)	$O(1/\sqrt{T})$	$\mathbb{E}[\mathcal{L}(\theta_T) - \mathcal{L}^*]$
$L$-smooth non-convex	SGD ($\eta \propto 1/\sqrt{T}$)	$O(1/\sqrt{T})$	$\frac{1}{T}\sum \mathbb{E}\|\nabla \mathcal{L}\|^2$

These rates are worst-case. In practice, deep learning loss surfaces have favorable structure (saddle points are easily escaped, local minima are often near-global) that makes convergence faster than theoretical bounds suggest. The rates are most useful for comparing algorithms and understanding scaling behavior (e.g., doubling the number of SGD steps gives $\sqrt{2}\times$ improvement).

Notation Summary

Symbol	Meaning
$\theta$	Model parameters
$\eta, \eta_t$	Learning rate (possibly time-varying)
$\mathcal{L}$	Loss function (empirical risk)
$\ell$	Per-sample loss
$g_t$	Gradient (or stochastic gradient) at step $t$
$v_t$	Velocity (momentum) or second moment (Adam)
$m_t$	First moment estimate (Adam)
$\hat{m}_t, \hat{v}_t$	Bias-corrected moment estimates
$\beta_1, \beta_2$	Exponential decay rates (Adam)
$\mu$	Momentum coefficient
$L$	Lipschitz constant of the gradient ($L$-smoothness)
$\mu$ (strong convexity)	Strong convexity parameter
$\kappa = L/\mu$	Condition number
$B =	\mathcal{B}
$\sigma^2$	Per-sample gradient variance
$c$	Gradient clipping threshold

---

References

• Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin (2017). Attention Is All You Need. NeurIPS.