📝 Gradient Descent

Description

< What is it? >

• [The idea]

  • Gradient descent is the iterative optimizer behind almost all model training: start from random parameters, measure which direction increases the loss, and step the other way. Repeat until the loss stops improving (convergence).

  • The gradient ∇J(θ) (∇: Nabla = differential operator, here is the gradient) is the vector of partial derivatives of the loss with respect to every parameter. It points in the direction of steepest increase, so −∇J(θ) is the steepest descent direction — hence the minus sign:

    θ  :=  θ  −  α · ∇J(θ)
    │         │      │
    new       │      gradient of the loss w.r.t. every parameter
    params    learning rate α (step size)

  • Nothing about this is specific to neural networks. It is the same rule for linear regression, logistic regression and a 100-billion-parameter transformer — only the way the gradient is computed differs (a closed-form expression for the simple cases, backpropagation for deep nets).

• [What it requires]
  • A loss that is differentiable with respect to the parameters. This is why classification models optimize cross-entropy rather than accuracy: accuracy is a step function with zero gradient almost everywhere.
  • Models whose parameters are not fit by a continuous objective — decision trees and their ensembles — do not use it. They search over discrete splits instead.
• [The training loop]

  repeat until converged:
      1. forward   — compute predictions ŷ for a batch of examples
      2. loss      — compare ŷ to y, get a single scalar J
      3. backward  — backpropagate to get ∇J(θ) for every parameter
      4. step      — θ := θ − α · ∇J(θ)

< Training-loop terminology >

Term	Meaning
Batch	A group of training examples processed together in one forward and backward pass
Mini-batch	A small batch—usually only part of the training set—used for one step. In practice, “batch” usually means “mini-batch”
Iteration / step	One parameter update using one batch
Epoch	One complete pass through all training batches

< Example >

For example, suppose you have 1,000 training examples:

• Batch size = 100
• Each batch causes one parameter update
• Number of updates per epoch:

$$\frac{1000}{100}=10$$

Therefore:

• 1 epoch = process all 1,000 examples once = 10 updates
• 5 epochs = process them five times = 50 updates
• 100 epochs = process them 100 times = 1,000 updates

< Batch, stochastic, mini-batch >

Variant	Examples per step	Gradient quality	Reality
Batch (full)	All m	Exact, smooth descent	One update per epoch. Too slow, and the dataset may not fit in memory
Stochastic (SGD)	1	Very noisy	Updates constantly, but wastes hardware — no vectorization
Mini-batch	32 – 8192	Noisy but usable	What everyone actually uses. Vectorizes on GPU and the noise is useful

• The noise in mini-batch gradients is not purely a cost. It helps the optimizer escape saddle points and narrow minima, which is one reason a moderate batch size often generalizes better than an enormous one.
• Batch size and learning rate are coupled: multiplying the batch size by k roughly permits multiplying the learning rate by k (the linear scaling rule), because each gradient estimate is proportionally less noisy.
• Shuffle every epoch. Data that arrives sorted by label makes each batch's gradient point somewhere unrepresentative.
• Where to start: 32, 64, 128 or 256. The upper end of the table is reserved for large-scale distributed runs; the practical ceiling on a single device is GPU memory, since activations for the whole batch must be held for the backward pass.
• Too small has its own failure mode. Batch normalization estimates its statistics from the batch, so below roughly 8–16 examples those estimates get noisy enough to hurt, and a batch size of 1 leaves it undefined in training mode.

< The learning rate — the one knob that matters most >

α too small           α about right          α too large
      ╲                     ╲                      ╲    ╱
       ╲....                 ╲                      ╲  ╱ ↑ diverges
        ╲....•                ╲___•                  ╲╱
   creeps, may never    lands in the basin      overshoots the valley,
   reach the minimum    in few steps            loss → NaN

• Schedules decay α over training: step decay, cosine annealing, or 1/√t. Large steps early to cover ground, small steps late to settle. Cosine decay is the default for transformer training.
• Warmup ramps α up from ~0 over the first few hundred/thousand steps. Adaptive optimizers have unreliable variance estimates at step 1, and a full-size step then can wreck the initialization.
• Find a starting value by sweeping α on a log scale (1e-5 … 1e-1) for a few hundred steps and taking the largest value whose loss still falls smoothly.

Key points

< Optimizers: what each one fixes >

Optimizer	Adds	Fixes
SGD	—	Baseline. Slow through ravines, sensitive to α
+ Momentum	Velocity: an exponentially-weighted average of past gradients	Damps zig-zag across a narrow valley, accelerates along it
+ Nesterov	Look-ahead gradient	Slightly better-damped momentum
AdaGrad	Per-parameter α scaled by accumulated squared gradients	Rare features get larger steps — but α decays to zero and training stalls
RMSProp	Same, with a moving average instead of a sum	Removes AdaGrad's terminal stall
Adam	RMSProp + momentum + bias correction	The robust default. Works well with little tuning
AdamW	Decouples weight decay from the adaptive step	Makes regularization behave correctly — see Regularization

• In practice: AdamW for transformers and most NLP; SGD with momentum still matches or beats it on vision CNNs and often generalizes slightly better; plain SGD for convex problems where you can tune α properly.
• All of them are still gradient descent. They only reshape how the raw gradient is turned into a step.

< The landscape: why this works at all >

• For convex losses (linear regression's MSE, logistic regression's log-loss) there is a single minimum and gradient descent provably finds it. Everything is well behaved.
• Deep networks are non-convex — a surface with countless critical points. The classical worry was getting trapped in a bad local minimum. In high dimensions that turns out to be rare: a local minimum requires every one of millions of directions to curve upward. What you actually hit are saddle points (up in some directions, down in others), and gradient noise plus momentum escape them.
• Most minima found in practice reach a similar loss. Flat minima — wide basins where nearby parameters give similar loss — tend to generalize better than sharp ones, which is another argument for the noise in mini-batch SGD.

< Reading the loss curve >

Symptom	Likely cause	Move
Loss → NaN / inf after a few steps	α too large, or exploding gradients	Cut α 10×, add gradient clipping, add warmup
Loss flat from step 0	α far too small, dead ReLUs, or the gradient never reaches the params	Raise α; check requires_grad, zero_grad(), and that the loss is connected to the graph
Loss falls then plateaus high	Underfitting, or α now too large to settle	Add capacity, or decay α — see Bias & Variance
Very spiky loss	Batch too small, α too large, or unshuffled data	Larger batch, lower α, shuffle
Training loss falls, validation rises	Overfitting — not an optimizer problem	Early stopping, regularization
Slow zig-zag progress	Ill-conditioned surface from unscaled features	Normalize the inputs; use momentum or Adam

< Things that quietly break it >

1. [Unscaled features] elongated loss contours cap α at whatever the steepest direction tolerates — the geometry argument in Data Normalization.
2. [Exploding / vanishing gradients] deep stacks multiply gradients through many layers. Fixes are structural: gradient clipping, residual connections, careful initialization (He/Xavier([ˈzævɪr])), and normalization layers.
3. [Forgetting to reset gradients] PyTorch accumulates into .grad; without optimizer.zero_grad() every step uses the sum of all previous gradients.
4. [Loss scale under mixed precision] fp16 gradients underflow to zero. Use a gradient scaler, or bf16.
5. [Tuning α against the test set] the learning rate is a hyperparameter like any other — see the validation-set rule in Bias & Variance.

< When not to use it >

• Closed form exists. OLS linear regression has the normal equation θ = (XᵀX)⁻¹Xᵀy — exact, no α to tune. Gradient descent wins once n is large enough that inverting XᵀX is expensive.
• Second-order methods (Newton, L-BFGS) use curvature and converge in far fewer iterations, but need the Hessian or an approximation of it — quadratic in the parameter count. Practical for small convex problems, hopeless at deep-learning scale, which is why first-order methods dominate.

Implementation

< Linear Regression Using Gradient Descent >

Question: Write a Python function that performs linear regression using gradient descent. The function should take NumPy arrays X (features with a column of ones for the intercept) and y (target) as input, along with the learning rate alpha and the number of iterations. Return the learned coefficients (weights) as a NumPy array.
Requirements:

• Minimize the Mean Squared Error (MSE) loss function: $$L(\theta)=\frac{1}{2m}\sum_{i=1}^{m}\left(h_\theta\left(x^{(i)}\right)-y^{(i)}\right)^2$$ Here, $h_\theta\left(x^{(i)}\right)=\left(x^{(i)}\right)^\top\theta$ is the prediction for training example $i$, and $m$ is the number of training examples. The factor $1/2$ simplifies the gradient calculation.
• Initialize all weights to zero
• Use batch gradient descent (use all samples in each iteration)
  

The input matrix X has shape (m, n) where m is the number of training examples and n is the number of features (including the bias column of ones). The target vector y has shape (m,).

Input: X = np.array([[1, 1], [1, 2], [1, 3]]), y = np.array([3, 5, 7]), alpha = 0.1, iterations = 1000

Output: [1.0, 2.0]

Solution:
The vectorized loss is

$$L(\theta)=\frac{1}{2m}\lVert X\theta-y\rVert_2^2$$

so its gradient is

$$\nabla_\theta L(\theta)=\frac{1}{m}X^\top(X\theta-y)$$

Step-by-step derivation of the gradient

Let the prediction error vector be

$$r = X\theta-y.$$

The squared $L_2$ norm is the sum of squared errors:

$$\lVert X\theta-y\rVert_2^2 =(X\theta-y)^\top(X\theta-y) =\sum_{i=1}^{m}(\hat y_i-y_i)^2.$$

Therefore,

$$L(\theta)=\frac{1}{2m}(X\theta-y)^\top(X\theta-y).$$

Derivation by expansion

Expand the product:

$$\begin{aligned} L(\theta) &=\frac{1}{2m}(X\theta-y)^\top(X\theta-y)\\ &=\frac{1}{2m}\left(\theta^\top X^\top X\theta-2y^\top X\theta+y^\top y\right). \end{aligned}$$

Differentiate each term with respect to $\theta$:

$$\nabla_\theta\left(\theta^\top X^\top X\theta\right)=2X^\top X\theta,$$

because $X^\top X$ is symmetric, and

$$\nabla_\theta\left(-2y^\top X\theta\right)=-2X^\top y.$$

The final term does not depend on $\theta$:

$$\nabla_\theta(y^\top y)=0.$$

Putting everything together:

$$\begin{aligned} \nabla_\theta L(\theta) &=\frac{1}{2m}\left(2X^\top X\theta-2X^\top y\right)\\ &=\frac{1}{m}\left(X^\top X\theta-X^\top y\right)\\ &=\boxed{\frac{1}{m}X^\top(X\theta-y)}. \end{aligned}$$

The factor $1/2$ in the loss cancels the $2$ produced when differentiating the squared error.

Element-by-element interpretation

For parameter $\theta_j$,

$$\frac{\partial L}{\partial\theta_j} =\frac{1}{m}\sum_{i=1}^{m}(\hat y_i-y_i)x_{ij}.$$

This means:

1. Compute every prediction error, $X\theta-y$.
2. Multiply each error by its corresponding feature value.
3. Sum across all training examples using $X^\top$.
4. Divide by $m$.

Stacking all parameter derivatives into one vector gives

$$\nabla_\theta L(\theta)=\frac{1}{m}X^\top(X\theta-y).$$

The gradient-descent update is consequently

$$\theta\leftarrow\theta-\alpha\frac{1}{m}X^\top(X\theta-y).$$

import numpy as np

def linear_regression_gradient_descent(X: np.ndarray, y: np.ndarray, alpha: float, iterations: int) -> np.ndarray:
"""
Perform linear regression using gradient descent.

Args:
    X: Feature matrix of shape (m, n) where first column is all ones (for intercept)
    y: Target vector of shape (m,)
    alpha: Learning rate
    iterations: Number of gradient descent iterations

Returns:
    Learned weights as a 1D array of shape (n,)
"""
m, n = X.shape
y = y.reshape(-1, 1) # Ensure y is a column vector
theta = np.zeros((n, 1)) # Initialize weights to zeros

for _ in range(iterations):
  predictions = X @ theta
  errors = predictions - y
  gradient = (X.T @ errors) / m
  theta = theta - alpha * gradient
return theta.flatten()

Q&A

< What is convergence? >

• Convergence is when the loss stops improving. In practice, this is usually defined as a plateau in the validation loss, or a small enough gradient norm. It does not necessarily mean that the model has reached the global minimum of the loss function, especially in non-convex problems like deep learning.

< Is convergence to the global optimum guaranteed? >

• Not always.
  For convex loss functions, gradient descent is guaranteed to converge to the global optimum.
  However, for non-convex loss functions, such as those encountered in deep learning, convergence to a global optimum is not guaranteed. Instead, the algorithm may converge to a local minimum or saddle point.

• Convex Objectives

  If $f$ is convex and differentiable, every local minimum is global. Gradient descent converges to a global optimum under standard conditions, such as:

  • The gradient is Lipschitz-continuous.
  • The learning rate is sufficiently small, often $\eta \leq 1/L$, where $L$ is the gradient smoothness constant.
  • An optimum exists.

  For a smooth convex function, the typical objective-gap rate is:

  $$f(\theta_t)-f(\theta^*)=O(1/t).$$

  If the function is also strongly convex, convergence is geometric:

  $$f(\theta_t)-f(\theta^*)=O(\rho^t), \qquad 0 < \rho < 1.$$

• Non-Convex Objectives

  For non-convex objectives, including most neural networks, global optimality is not generally guaranteed. Gradient descent can encounter:

  • Multiple local minima
    
  • Saddle points
    
  • Flat plateaus
  • Poorly conditioned regions
  • Vanishing or exploding gradients
    

  Under smoothness assumptions, a common guarantee is only convergence toward a stationary point:

  $$\|\nabla f(\theta_t)\|\rightarrow 0.$$

  A stationary point can be a local minimum, local maximum, or saddle point—not necessarily a global optimum.

• Stochastic Gradient Descent

  Stochastic gradient descent (SGD) uses a noisy gradient estimate from a sample or mini-batch:

  $$\theta_{t+1}=\theta_t-\eta_t\widehat{\nabla f}(\theta_t).$$

  Its noise makes updates cheaper and can help escape some saddle points or shallow local minima. However, with a fixed learning rate, SGD often fluctuates around a solution rather than converging exactly. Suitable decreasing learning rates, such as conditions

  $$\sum_t \eta_t = \infty, \qquad \sum_t \eta_t^2 < \infty,$$

  can provide convergence guarantees under additional assumptions.

• Practical Takeaway

  • Convex problem:

    Global convergence is often guaranteed with an appropriate learning-rate schedule.

  • Strongly convex problem:

    Global convergence is usually fast and the optimum is unique.

  • Non-convex problem:

    Usually only stationary-point convergence is guaranteed.

  • Deep learning:

    There is generally no global-optimum guarantee, although overparameterization, initialization, normalization, adaptive optimizers, and SGD often produce good solutions empirically.

Crash course

• Gradient descent
  (Vanilla) Gradient Descent

  STOCHASTIC Gradient Descent

  MOMENTUM Gradient Descent

Reference

• An overview of gradient descent optimization algorithms (Sebastian Ruder)
• Why Momentum Really Works (Distill)