Gradient Descent §

The learning rate: the step size $\eta$ controls the tradeoff between speed and stability.

• $\eta$ too large: the iteration overshoots the minimum and diverges (violates the contraction condition $\eta < 2/L$).
• $\eta$ too small: convergence is slow (the contraction factor $\rho \approx 1 - \eta\mu$ is close to 1).
• $\eta = 1/L$: the optimal choice for strongly convex problems; gives the fastest linear convergence.

The epoch: in practice, data is processed in passes. One epoch is one complete pass over all $n$ training examples. The total number of gradient steps is $\text{epochs} \times \lceil n/b \rceil$ where $b$ is the batch size.

Example: minimizing $f(x) = 5x^2$ from $x_0 = 1$ with $\eta = 0.1$ (the optimal $1/L$ where $L = 10$): $x_1 = 1 - 0.1 \times 10 = 0$. Converged in one step.

Example problem:

Why noise helps: in high-dimensional non-convex landscapes, saddle points dominate local minima (Dauphin et al., 2014). The stochastic noise in SGD provides the perturbation needed to escape saddle points, which deterministic GD cannot do (it gets stuck at any critical point).

Batch size tradeoff:

• $b = 1$ (pure SGD): maximum noise, fastest per-step, but high variance in the gradient estimate.
• $b = n$ (batch GD): zero noise, slowest per-step, but a precise gradient.
• $b \in [32, 256]$ (typical): a practical compromise. The noise is small enough for stable training but large enough to provide regularization.

Learning rate schedule: in practice, $\eta$ is decayed over training (e.g., step decay, cosine annealing). A large initial $\eta$ explores the landscape; a small final $\eta$ fine-tunes the solution. This is analogous to simulated annealing in physics.

Connection to Regression: the train/validation split and overfitting diagnostics from regression apply directly to SGD training. The stochastic noise in SGD acts as an implicit regularizer, often reducing the need for explicit weight decay.

Example problem:

Intuition: imagine a ball rolling in a landscape. Without momentum, the ball stops at every small bump. With momentum, the ball builds up speed in consistent directions and plows through small bumps. In a narrow valley (high curvature in one direction, low in another), the oscillations in the high-curvature direction cancel out (the velocity averages to zero), while the progress in the low-curvature direction accumulates.

The effective learning rate: with momentum, the effective step size in a consistent direction is $\eta/(1-\beta)$. For $\beta = 0.9$, the effective learning rate is $10\eta$. This allows a larger nominal $\eta$ without instability, because the oscillatory components are damped.

Nesterov momentum: a refinement that looks ahead: $v_t = \beta v_{t-1} + (1-\beta)\nabla f(\theta_t - \eta v_{t-1})$. The gradient is evaluated at the predicted next position, not the current one. This gives a small but consistent improvement in convergence rate.

Example problem:

Why adaptive step sizes help: in neural networks, different parameters have very different gradient scales. Embedding table entries for rare tokens have sparse, large gradients; bias terms have small, dense gradients. Adam normalizes each coordinate by its historical gradient magnitude, so all parameters receive comparable step sizes.

Typical hyperparameters: $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\varepsilon = 10^{-8}$, $\eta = 10^{-3}$ to $10^{-4}$. These defaults work well across most problems, which is why Adam became the default optimizer for deep learning.

Convergence: Adam does not converge for all problems (Reddi et al., 2018 showed a counterexample). In practice, it works well for non-convex deep learning objectives. For convex problems, SGD with a proper learning rate schedule is often preferable.

Connection to Neural Networks: Adam is the default optimizer for training neural networks. The compute section of that page covers how the gradients are computed (backpropagation) and how they are applied in parallel on GPUs.

Example problem:

Convex vs. non-convex: for convex $f$, every local minimum is global, and gradient descent with $\eta = 1/L$ converges linearly. For non-convex $f$ (neural networks), the landscape is complex:

• Saddle points: points where the gradient is zero but the Hessian has both positive and negative eigenvalues. In high dimensions, these are far more common than local minima.
• Plateaus: regions where the gradient is very small but the point is not a minimum. Adam’s adaptive step size helps traverse plateaus.
• Sharp vs. flat minima: sharp minima (high Hessian eigenvalues) generalize poorly; flat minima (low Hessian eigenvalues) generalize well. SGD’s noise implicitly prefers flat minima.

The valley structure: in ill-conditioned problems (large $L/\mu$ ratio), the loss landscape has narrow valleys. Gradient descent oscillates across the valley while slowly descending. Momentum and adaptive methods (Adam) address this by damping the oscillatory component.

Example problem:

Parameters:

• Layer 1: $W_1 \in \mathbb{R}^{128 \times 784}$ → $128 \times 784 = 100{,}352$ weights + $128$ biases = $100{,}480$.
• Layer 2: $W_2 \in \mathbb{R}^{10 \times 128}$ → $10 \times 128 = 1{,}280$ weights + $10$ biases = $1{,}290$.
• Total: $101{,}770$ parameters.

FLOPs (forward):

• Layer 1: $2 \times 128 \times 784 = 200{,}704$ (multiply-accumulate counts as 2 FLOPs).
• Layer 2: $2 \times 10 \times 128 = 2{,}560$.
• Total: $\approx 203{,}264$ FLOPs per example.

Backpropagation costs 2× the forward FLOPs: 1× forward for activation gradients (error signals) + 1× forward for weight gradients. Total training step = 3× forward ≈ $610{,}000$ FLOPs.