Gradient Descent

Overview

In the most general sense, training a model usually means choosing parameters that make a loss function small. In Linear Regression, for example, we choose weights and a bias to minimise the mean-squared error between predictions and labels.

There are, of course, many ways to minimise said loss, as exemplified by Linear Regression:

• We could solve the normal equation with a least-squares method, which yields a closed-form solution, assuming the data has full column rank (i.e. there are enough data points to determine the parameters, and the features are independent).
• We could also use gradient descent, which is the topic of this article.

How do we do Gradient Descent?

For Gradient Descent, the name of the game is to start with arbitrary parameters and iteratively improve them by moving opposite to the gradient of our loss function (thus, roughly speaking, the loss function is assumed differentiable with respect to the parameters).

An Analogy for Gradient Descent

We can intuitively understand this process by likening it to hikers descending a mountain in fog so thick they can only see a few steps ahead.

So the story goes:

The hikers cannot see the bottom of the mountain, so they can only rely on their immediate surroundings to feel their way down. They can feel the slope of the ground, or measure it with instruments, and determine the locally steepest downhill direction. They then discuss how far to move before stopping to measure the slope again. As the fog is too thick, they cannot know whether they have reached the lowest point of the mountain. They might instead stop in a local valley, make very slow progress across a nearly flat region, or repeatedly overshoot a narrow valley by taking steps that are too large.

This analogy captures a few key recurring themes of Gradient Descent that are worth highlighting:

and it should also reveal that Gradient Descent is actually a pretty crass algorithm (after all, real-life hikers would be quite cooked if they actually encountered such a scenario).

Yet, gradient descent is one of the most widely used optimisation algorithms in machine learning. We will explore how it works, and why this is the case in the next few sections.

Algorithm

To go into further details, let’s assume we are dealing with a regression task, and write some (extremely sloppy) math down:

• Suppose we had a model $f(\mathbf{x};\theta)$, accepting $\mathbf{x}$ as its input and being parameterised by $\theta$.
• We wish to train the model on an indeterminate dataset $\mathcal{D} = \{(\mathbf{x}_i, \mathbf{y}_i)\}_{i=1}^n$, where we are trying to discover $\theta$ such that $f(\mathbf{x}_i; \theta) \approx \mathbf{y}_i$ for all $i$ (i.e. we are “fitting” $f$ to $\mathcal{D}$).

In Gradient Descent (for regression tasks):

• We start by defining a loss function $\mathcal{L}(f(\mathbf{x}; \theta), \mathbf{y})$, which measures how far off our model’s predictions are from the true labels.
  • Aside from being differentiable w.r.t. $\theta$, it is mandatory for $\mathcal{L}$ to be minimized whenever $f(\mathbf{x}; \theta) = \mathbf{y}$.
• We then arbitrarily select $\theta_0$ as our initial parameters, and a learning rate $\eta$
• For every timestep $t \geq 0$ (repeating until we are satisfied):
  • We compute $\mathbf{g}_t = \nabla_{\theta_t} \mathcal{L}$, the gradient of the loss function w.r.t. the parameters
    • Indeed, $\mathbf{g}_t$ is a vector of partial derivatives of $\mathcal{L}$ w.r.t. each parameter in $\theta_t$.
  • We then collectively update our parameters using the update rule $\theta_{t+1} \leftarrow \theta_t - \eta \mathbf{g}_t$.
    • In calculus, this is a first-order approximation of our loss function, as $\mathcal{L}(\theta\_{t} - \eta \mathbf{g}\_t) \approx \mathcal{L}(\theta_t) - \eta \mathbf{g}\_t^\top \mathbf{g}\_t$.
    • Notice how it always reduces our approximated loss whenever $\mathbf{g}_t \neq \mathbf{0}$ (i.e. whenever we are not at a stationary point), but is not guaranteed to reduce the true loss.

Worked Example

Consider a toy loss $\mathcal{L}(f(\mathbf{x}; \theta), \mathbf{y}) = \frac{1}{2}(\theta - \mathbf{y})^2$ where $\theta, \mathbf{x}, \mathbf{y} \in \mathbb{R}$ and $f(\mathbf{x}; \theta) = \theta$.

• This loss function satisfies $\nabla_\theta^2 \mathcal{L} = 1 > 0$, so it is special and has a unique global minimum at $\theta = \mathbf{y}$. More on that later.

With $\theta = 0$, $\mathbf{y} = 3$, and $\eta = 0.5$, the update rule becomes

You can trace out the first few updates and check against the table below.

Baby’s First Training Dynamics Analysis

Mechanically tracing how updates evolve can familiarize us with the procedure of gradient descent, but it is arguably more useful to study how the updates change as training goes on (“training dynamics”).

Continuing with our toy example, we can observe from the above that every update moves $\theta_t$ (ever more slowly) towards 3. In fact, rearranging gives

so (at least when we have a nice loss like this), the multiplier $1-\eta$ completely determines HOW the weights evolve, as summarized below:

This shows something treated as quite obvious by ML practitioners: that learning rate is a hyperparameter that must be carefully tuned.

Convergence, Convexity, and Polyak-Lojasiewicz

Typically, we do not run gradient descent indefinitely. We stop when we are satisfied with the solution, and we then call the model “converged” (onto a set of parameters $\theta^\dag$).

What “satisfied” means varies quite wildly depending on who you ask:

• (From Boyd and Vandenberghe) It could be $\lim_{t\to\infty}\theta_t=\theta^*$, where $\theta^*$ is a true global minimizer of the loss function.
• (From Polyak) It could also be $\lim_{t\to\infty}\mathcal{L}(\theta_t)=\mathcal{L}(\theta^*)$, where $\mathcal{L}(\theta^*)$ is a true global minimum of the loss function.
• (From Nocedal and Wright) It could be $\lim_{t\to\infty}\lVert\nabla\mathcal{L}(\theta_t)\rVert=0$; that is, we reach a stationary point of the loss function, NOT even a local minimum.
• “satisfied” can also mean that the model performs “well enough” on some (optionally) held-out validation data.

However, if we are willing to put up some strong assumptions on the loss function and assume it is convex, we can guarantee every local minimum must be a global minimum.

Unfortunately, many loss functions are non-convex, so most practitioners default to the last option, and use convergence as a vague conceptual tool to guide their training process instead.

Batch, Stochastic, and Mini-Batch Gradients

Up to this point, we’ve discussed gradient descent (GD) as if we had access to the true gradient of the loss function at every step.

In this regime, we typically break up the loss function into an aggregate of per-example losses (such as “average”), then compute the gradient of this aggregate loss at every step. This is called batch GD:

This is rarely the case in practice, and with finite (or even small) datasets (relative to a presumed infinite population), we are forced to work with noisy estimates of our loss (and therefore, its gradient) at every step. We often express the gradient as

where $\epsilon_i$ is a random variable with mean $\mu$ and variance $\sigma^2$ (and, if the batch size is big enough, $\epsilon = \frac{1}{n}\sum_{i=1}^n \epsilon_i$ is approximately Gaussian by the Central Limit Theorem).

Using a single, truly-randomly selected sample (i.e. n = 1) gives us stochastic GD (SGD, not to be confused with the optimizer of the same name), while using a small, truly-randomly selected subset of samples (i.e. $n < |\mathcal{D}|$) gives us mini-batch GD.

In the vaguest of senses, here are some considerations when using these three variants of GD (which essentially vary $n$):

• Increasing $n$ reduces the mean noise variance $\epsilon = \frac{1}{n}\sum_{i=1}^n \epsilon_i$ (with variance $\sigma^2/n$) but makes each update susceptible to overfitting
• Decreasing $n$ and increasing the noise of the gradient update can help in common non-convex loss landscapes (in escaping local minima, and saddle points, and high-loss basins), but can also hurt convergence.
• However, on modern hardware (especially GPUs) where parallelism is cheap, increasing $n$ can reduce wall-clock time per step, and speed up convergence in practice.

Related Topics

Linear Regression provides a concrete convex loss on which the update and convergence ideas are easiest to see. In Neural Networks, backpropagation computes the gradient and a gradient-based optimiser applies the parameter update. Transformers inherit this training loop from neural networks, usually with mini-batches and an extension of plain gradient descent.