Understanding the Gradient Descent Algorithm: A Deep Dive into the Math Behind It

Gradient descent is one of the most fundamental and widely used optimization algorithms in machine learning. It is a first-order iterative optimization algorithm used to find the local minimum of a cost function by iteratively adjusting the parameters in the direction of steepest descent as defined by the negative of the gradient. In this blog post, we‘ll dive deep into the math behind the gradient descent algorithm, develop intuition around its variants like batch, mini-batch and stochastic gradient descent, and demonstrate how to implement it from scratch in Python.

Intuitive Explanation

Before diving into the mathematical details, let‘s develop some intuition for what the gradient descent algorithm does. Imagine you are standing at the top of a hill and your goal is to reach the bottom of the valley as quickly as possible. A reasonable strategy would be to look around, identify the direction of steepest descent, take a step in that direction, and repeat the process until you reach the bottom of the valley.

This is essentially what the gradient descent algorithm does – it iteratively adjusts parameters of a model in the direction that minimizes a cost function until it converges to a local minimum. The size of the steps is controlled by a hyperparameter called the learning rate. The direction of descent is determined by computing the gradient of the cost function with respect to each parameter.

There are three main variants of gradient descent:

1. Batch Gradient Descent: Parameters are updated after computing the gradient on the entire training set.
2. Stochastic Gradient Descent: Parameters are updated after computing the gradient on a single randomly selected training example.
3. Mini-batch Gradient Descent: Parameters are updated after computing the average gradient of a mini-batch of n training examples.

The following table summarizes the key differences between the variants of gradient descent:

Variant	Definition	Advantages	Disadvantages
Batch GD	Use all m examples in each iteration	Computationally efficient, less noisy convergence	Requires full pass through data to update, memory intensive
Stochastic GD	Use 1 example in each iteration	Frequent updates, less memory required	Noisy convergence, may never converge to exact minimum
Mini-batch GD	Use n examples in each iteration	Frequent updates, smoother convergence, computationally efficient	Need to tune batch size n

Batch gradient descent performs redundant computations for large datasets, as it recomputes gradients for similar examples before each parameter update. Stochastic GD performs frequent updates with a high variance that cause the loss function to fluctuate heavily, but it can get us to a reasonably good solution very fast. Mini-batch gradient descent is typically the algorithm of choice when training a neural network and the batch size is a hyperparameter that needs to be tuned.

Mathematical Derivation

Let‘s now derive the gradient descent update rule from first principles using some fundamental concepts from calculus like partial derivatives and the chain rule. We‘ll start with the case of a simple linear regression model which predicts an output y as an affine function of an input x:

$y = mx + b$

Here, m and b are the parameters of the model that we wish to learn from data. Given a training set of N input-output pairs $(x_i, y_i)$, we define a cost function $J(m, b)$ that measures how well the model fits the training data. A common choice is the mean squared error (MSE) which is the average squared difference between the predicted and true values:

$J(m, b) = \frac{1}{2N} \sum_{i=1}^N (y_i – (mx_i + b))^2$

To find the values of the parameters m and b that minimize the cost function, we use a technique called gradient descent. The key idea is to iteratively adjust the parameters in the direction of the negative gradient of the cost function w.r.t. to the parameters. Mathematically, this is expressed by the following update rule:

$\begin{align}
m &= m – \alpha \frac{\partial J}{\partial m} \
b &= b – \alpha \frac{\partial J}{\partial b}
\end{align}$

Here, $\alpha$ is a hyperparameter called the learning rate which controls the size of the steps we take in the direction of the negative gradient. The partial derivatives $\frac{\partial J}{\partial m}$ and $\frac{\partial J}{\partial b}$ give us the direction of steepest ascent with respect to parameters m and b respectively. To compute them, we use the chain rule and power rule from calculus:

$\begin{align}
\frac{\partial J}{\partial m} &= \frac{1}{N} \sum_{i=1}^N (-(y_i – (mx_i + b)) \cdot xi) \
\frac{\partial J}{\partial b} &= \frac{1}{N} \sum{i=1}^N (-(y_i – (mx_i + b)))
\end{align}$

Plugging the above expressions into the update rule gives us the batch gradient descent algorithm which updates the parameters as follows:

$\begin{align}
m &= m – \alpha \frac{1}{N} \sum_{i=1}^N (-(y_i – (mx_i + b)) \cdot xi) \
b &= b – \alpha \frac{1}{N} \sum{i=1}^N (-(y_i – (mx_i + b)))
\end{align}$

For stochastic and mini-batch gradient descent, the update rules are similar except the gradient is computed on a single example or a mini-batch of examples respectively at each iteration instead of averaging over all examples. The following code snippet shows how to implement batch gradient descent in Python from scratch:

def gradient_descent(X, y, m_init, b_init, alpha, num_iters):
    m = m_init
    b = b_init

    for _ in range(num_iters):
        m_gradient = np.mean(-(y - (m * X + b)) * X)
        b_gradient = np.mean(-(y - (m * X + b)))
        m = m - (alpha * m_gradient)
        b = b - (alpha * b_gradient)

    return m, b

Here, X and y are the input and output variables respectively, m_init and b_init are the initial values of the parameters, alpha is the learning rate, and num_iters is the number of iterations to run the algorithm for. At each iteration, we compute the average gradients w.r.t. m and b over all examples, and update the parameters by taking a step of size alpha in the negative direction of the gradient.

Momentum and Nesterov Accelerated Gradient

While gradient descent is guaranteed to converge to a local minimum for convex loss surfaces, the convergence can be slow for ill-conditioned problems where the loss surface has a steep gradient in one dimension and a shallow gradient in another. A simple extension of gradient descent called momentum accelerates convergence by adding a fraction of the previous update to the current update at each iteration:

$\begin{align}
v &= \beta v – \alpha \frac{\partial J}{\partial \theta} \
\theta &= \theta + v
\end{align}$

Here, $v$ is the velocity, $\beta$ is a hyperparameter between 0 and 1 that controls the momentum, $\alpha$ is the learning rate, and $\theta$ represents the model parameters. Intuitively, momentum helps the algorithm build up velocity in directions of low curvature that consistently reduce the loss and dampens oscillations in high curvature directions.

Nesterov accelerated gradient (NAG) is a variant of momentum that computes the gradient at the approximate future location of the parameters rather than at the current location:

$\begin{align}
v &= \beta v – \alpha \frac{\partial J}{\partial \theta}(\theta + \beta v) \
\theta &= \theta + v
\end{align}$

NAG has been shown to provide better convergence rates than standard momentum in certain cases. Momentum and NAG are widely used in practice when training deep neural networks.

Gradient Descent for Training Neural Networks

Gradient descent lies at the core of training deep neural networks. Neural networks are composed of multiple layers of interconnected nodes that perform weighted linear combinations of their inputs followed by non-linear activation functions. The goal is to learn the weights of the network that minimize a loss function on the training data.

For a neural network with $L$ layers, the output of the $l$-th layer is given by:

$a^{(l)} = g^{(l)}(z^{(l)})$ where $z^{(l)} = W^{(l)} a^{(l-1)} + b^{(l)}$

Here, $g^{(l)}$ is the activation function (e.g. sigmoid, ReLU) of the $l$-th layer, $W^{(l)}$ is the weight matrix, $b^{(l)}$ is the bias vector, and $a^{(0)} = x$ is the input. The goal is to learn the parameters $W^{(l)}$ and $b^{(l)}$ for each layer $l$ that minimize a loss function $J(W, b)$ on the training data.

The backpropagation algorithm uses the chain rule to efficiently compute the gradients of the loss function w.r.t. the parameters of each layer. The key step is to compute the error $\delta^{(l)}$ at each layer $l$, which measures how much that layer is responsible for the errors in the output. The errors are computed in a backward pass:

$\delta^{(L)} = \nabla_a J \odot g‘^{(L)}(z^{(L)})$
$\delta^{(l)} = ((W^{(l+1)})^T \delta^{(l+1)}) \odot g‘^{(l)}(z^{(l)})$

Here, $\odot$ represents element-wise multiplication and $g‘^{(l)}$ is the derivative of the activation function. Once we have the errors, we can compute the gradients w.r.t. the parameters as follows:

$\frac{\partial J}{\partial W^{(l)}} = \delta^{(l)} (a^{(l-1)})^T$
$\frac{\partial J}{\partial b^{(l)}} = \delta^{(l)}$

Finally, we update the parameters using gradient descent:

$W^{(l)} = W^{(l)} – \alpha \frac{\partial J}{\partial W^{(l)}}$
$b^{(l)} = b^{(l)} – \alpha \frac{\partial J}{\partial b^{(l)}}$

This process is repeated for a certain number of iterations or until convergence. The backpropagation algorithm is a special case of a more general technique called reverse-mode automatic differentiation. Modern deep learning frameworks like PyTorch and TensorFlow provide autograd packages that automatically compute the gradients based on the computational graph of the network which makes it very convenient to train neural networks without having to manually implement backpropagation.

Challenges and Best Practices

While gradient descent is a powerful technique, there are a few challenges to be aware of when using it in practice:

• Choosing the learning rate: The learning rate $\alpha$ is a critical hyperparameter that needs to be tuned. If $\alpha$ is too small, convergence is slow. If $\alpha$ is too large, the algorithm may oscillate around the minimum or diverge completely. Common strategies to tune $\alpha$ include grid search with exponential decay, or more advanced adaptive methods like Adam or RMSProp that adapt a separate learning rate for each parameter.

• Overfitting: Deep neural networks have a large number of parameters and can easily overfit the training data. Regularization techniques are used to constrain the model complexity and improve generalization to unseen data. Common regularization techniques include:

  • L1/L2 Regularization: Adding a penalty term to the loss function that encourages the weights to be small.
  • Dropout: Randomly setting a fraction of activations to zero during training which prevents co-adaptation of neurons.
  • Early Stopping: Monitoring performance on a validation set and stopping training when it starts degrading.

• Vanishing and Exploding Gradients: In very deep networks, the gradients can sometimes become very small (vanish) or very large (explode) during backpropagation which makes training difficult. Careful initialization of the weights is important to keep the gradients well-behaved. Common initialization strategies are:

  • Xavier Initialization: Initializing weights from a distribution with zero mean and variance $2/(n{in} + n{out})$ where $n{in}$ and $n{out}$ are the number of input and output units of a layer.
  • He Initialization: Initializing weights from a distribution with zero mean and variance $2/n_{in}$. Commonly used for ReLU activations.

• Local Minima and Saddle Points: The loss surfaces of deep neural networks are highly non-convex and can have many local minima and saddle points which makes optimization difficult. While this is still an active area of research, a few strategies that have proven effective in practice are:

  • Using a momentum term to help escape saddle points.
  • Gradient clipping to prevent gradients from exploding.
  • Second-order methods like L-BFGS that use curvature information to escape flat regions.

Conclusion

In this blog post, we took a deep dive into the math behind the gradient descent optimization algorithm. We derived the update rules for batch, mini-batch, and stochastic gradient descent from first principles using partial derivatives and the chain rule. We also discussed some advanced variants like momentum and Nesterov accelerated gradient that are commonly used in practice.

We then looked at how gradient descent forms the core of training deep neural networks and walked through the backpropagation algorithm to compute gradients efficiently. We also highlighted some of the key challenges and best practices when using gradient descent to train neural networks, including learning rate tuning, regularization, initialization, and optimization difficulties.

There have been many advances in optimization algorithms for machine learning in recent years, like Adam, AdaGrad, RMSProp, and L-BFGS, but gradient descent remains the workhorse algorithm that underlies them all. I hope this post gave you a solid foundation and understanding for how it works. For further reading, I recommend checking out:

• An overview of gradient descent optimization algorithms by Sebastian Ruder
• Why Momentum Really Works by Gabriel Goh
• A Recipe for Training Neural Networks by Andrej Karpathy