A Comprehensive Guide to Gradient Descent and Its Variants (with Python Code)

Gradient descent is one of the most important and widely used optimization algorithms in machine learning. It is the workhorse that powers many models, from simple linear regression to large-scale deep neural networks with millions of parameters. The goal of gradient descent is to minimize a cost function by iteratively adjusting the model‘s parameters in the direction that reduces the cost.

While the basic concept is straightforward, there are many variants of gradient descent that are commonly used, each with its own pros and cons. Understanding these variants and how to implement them is crucial for getting good performance out of your ML models. In this post, we‘ll take an in-depth look at batch gradient descent, stochastic gradient descent, mini-batch gradient descent, and gradient descent with momentum. We‘ll explain the key ideas behind each variant and walk through a complete Python implementation.

The Need for Optimization in Machine Learning

Before diving into the details of gradient descent, let‘s take a step back and consider why we need optimization algorithms in the first place. Suppose you are building a machine learning model to predict housing prices based on features like square footage, number of bedrooms, etc. You have a dataset of past home sales that you can use to train your model.

The goal is to learn a function that maps the input features to the output price, by tuning the model parameters to minimize some cost function, such as the mean squared error between the predicted and actual prices. There are many possible functions that could fit the data. Optimization algorithms like gradient descent provide an automated way to search for the parameters that give the best fit, as measured by the cost function.

Batch Gradient Descent

The most basic form of gradient descent is batch gradient descent. In this approach, we calculate the gradient of the cost function with respect to the model parameters, using the entire training dataset. We then update the parameters by taking a step in the negative direction of the gradient, scaled by a learning rate hyperparameter.

Here are the key steps in batch gradient descent:

1. Initialize the model parameters (randomly or using a pre-specified scheme)
2. Repeat until convergence:
   • Calculate the predicted outputs using the current parameters
   • Compute the cost function on the entire dataset
   • Calculate the gradient of the cost function w.r.t. the parameters
   • Update the parameters by taking a step in the negative gradient direction

Let‘s see how this works in practice with a Python implementation. We‘ll use a simple linear regression model as an example. First we‘ll generate a toy dataset:

import numpy as np

# Generate synthetic data
np.random.seed(0)
X = np.random.rand(100, 1)
y = 2 + 3 * X + np.random.rand(100, 1)

Next we‘ll define functions to compute the model outputs, cost function, and gradient:

def model(X, w, b):
    return X * w + b

def cost(y_pred, y):
    N = len(y) 
    return np.sum((y_pred - y)**2) / N

def gradient(X, y, w, b):
    N = len(y)
    y_pred = model(X, w, b)
    dw = (2/N) * np.sum(X * (y_pred - y))
    db = (2/N) * np.sum(y_pred - y)
    return dw, db

Finally, we‘ll implement batch gradient descent and train the model:

def batch_gradient_descent(X, y, learning_rate=0.01, num_iters=100):
    w = 0
    b = 0
    for i in range(num_iters):
        y_pred = model(X, w, b)
        dw, db = gradient(X, y, w, b)
        w -= learning_rate * dw
        b -= learning_rate * db
        if i % 10 == 0:
            print(f"Iteration {i}: Cost {cost(y_pred, y):.3f}")
    return w, b

w, b = batch_gradient_descent(X, y) 
print(f"Final w: {w:.3f}, b: {b:.3f}")

This will print the cost every 10 iterations so we can monitor convergence:

Iteration 0: Cost 4.601
Iteration 10: Cost 0.490
Iteration 20: Cost 0.153
...
Iteration 90: Cost 0.052
Final w: 2.998, b: 2.024

The cost decreases over time as the model parameters are optimized. After 100 iterations, we‘ve arrived at a pretty good fit to the true parameters (w=3, b=2).

Batch gradient descent is guaranteed to converge to the global minimum for convex cost functions like the mean squared error. However, it can be very slow for large datasets, since we have to process the entire dataset to compute each gradient step. This leads us to consider stochastic and mini-batch variants that work with subsets of data.

Stochastic Gradient Descent

Stochastic gradient descent (SGD) is a popular alternative to batch gradient descent that is much faster and can escape local minima. The key idea is that instead of computing the gradient on the entire dataset, we estimate it based on just a single randomly sampled data point. We then update the parameters based on this stochastic estimate of the gradient.

Since each gradient step only touches one data point, SGD can start making progress very quickly, often in the first few iterations. It also has a better chance of escaping local minima since the noisy gradient estimates can "jump out" of a poor local optimum. The downside is that SGD can have very high variance, which may cause it to bounce around or diverge if the learning rate is too high.

Here‘s how we can modify the linear regression code to use SGD:

def stochastic_gradient_descent(X, y, learning_rate=0.01, num_iters=100):
    w = 0
    b = 0
    for i in range(num_iters):
        idx = np.random.choice(len(X))
        x_i = X[idx]
        y_i = y[idx]
        y_pred_i = model(x_i, w, b)
        dw, db = gradient(x_i, y_i, w, b)
        w -= learning_rate * dw
        b -= learning_rate * db
        if i % 10 == 0:
            y_pred = model(X, w, b)
            print(f"Iteration {i}: Cost {cost(y_pred, y):.3f}")
    return w, b

w, b = stochastic_gradient_descent(X, y)
print(f"Final w: {w:.3f}, b: {b:.3f}") 

Notice that we randomly sample a single data point at each iteration and compute the gradient only on that point. We‘ve also moved the cost calculation and printout outside the loop since the cost will fluctuate a lot across iterations.

Running this code, we see that SGD indeed makes rapid progress at first, but then tends to bounce around:

Iteration 0: Cost 1.830
Iteration 10: Cost 0.137
Iteration 20: Cost 0.218
...
Iteration 90: Cost 0.080
Final w: 3.027, b: 1.792

The final parameters are still pretty close to the true values, but not as good as what we got with batch gradient descent. This is a common trade-off with SGD – faster initial convergence but more erratic behavior.

Mini-Batch Gradient Descent

Mini-batch gradient descent is a happy medium between batch gradient descent and SGD. Instead of computing the gradient on the entire dataset (batch) or a single point (SGD), we compute it on a small subset or "mini-batch" of data, usually 64-512 data points.

The benefits of mini-batch gradient descent are:

• Computationally more efficient than batch (since we only touch a subset of data per iteration)
• More stable convergence than SGD (since we average gradient over a mini-batch)
• Can take advantage of vectorized hardware operations for even greater efficiency

Here‘s the code for mini-batch gradient descent:

def mini_batch_gradient_descent(X, y, learning_rate=0.01, batch_size=32, num_iters=100):
    w = 0 
    b = 0
    for i in range(num_iters):
        idx = np.random.choice(len(X), batch_size)
        X_batch = X[idx]
        y_batch = y[idx]
        y_pred_batch = model(X_batch, w, b)
        dw, db = gradient(X_batch, y_batch, w, b)
        w -= learning_rate * dw
        b -= learning_rate * db
        if i % 10 == 0:
            y_pred = model(X, w, b)
            print(f"Iteration {i}: Cost {cost(y_pred, y):.3f}")
    return w, b

w, b = mini_batch_gradient_descent(X, y)
print(f"Final w: {w:.3f}, b: {b:.3f}")

The key difference is that we now sample a mini-batch of data points at each iteration rather than just one. We then compute the gradient averaged over this mini-batch.

In practice, mini-batch gradient descent is the most commonly used variant since it strikes a good balance between efficiency and stability. The batch size is a key hyperparameter that can be tuned – smaller batch sizes give noisier gradients that can help with exploration, while larger batch sizes give more stable gradients for better convergence.

Gradient Descent with Momentum

Gradient descent with momentum is a powerful extension that almost always outperforms "vanilla" mini-batch gradient descent. The key idea is to maintain a "velocity" vector that accumulates gradients over time. This helps the optimizer gain momentum in directions that consistently reduce the cost, while dampening oscillations in other directions.

Here are the equations for gradient descent with momentum:

v = beta * v - learning_rate * gradient
parameters += v

where v is the velocity vector (initialized to zero), beta is the momentum hyperparameter (typically 0.9), and gradient is the current mini-batch gradient. The velocity is essentially an exponentially decaying moving average of the gradients.

Intuitively, the velocity accumulates gradients in directions that consistently decrease the cost, speeding up progress. The beta hyperparameter controls how much the previous velocity is retained – a higher beta retains more momentum.

Here‘s how we can modify the mini-batch code to include momentum:

def mini_batch_gradient_descent_momentum(X, y, learning_rate=0.01, batch_size=32, num_iters=100, beta=0.9):
    w = 0
    b = 0
    v_w = 0
    v_b = 0
    for i in range(num_iters):
        idx = np.random.choice(len(X), batch_size)
        X_batch = X[idx]
        y_batch = y[idx]
        y_pred_batch = model(X_batch, w, b)
        dw, db = gradient(X_batch, y_batch, w, b)
        v_w = beta * v_w - learning_rate * dw
        v_b = beta * v_b - learning_rate * db
        w += v_w
        b += v_b
        if i % 10 == 0:
            y_pred = model(X, w, b)
            print(f"Iteration {i}: Cost {cost(y_pred, y):.3f}")
    return w, b

w, b = mini_batch_gradient_descent_momentum(X, y) 
print(f"Final w: {w:.3f}, b: {b:.3f}")

We‘ve introduced velocity variables v_w and v_b that are updated according to the momentum equations. The parameters w and b are then updated using the velocities rather than the raw gradients.

Running this code, we see that momentum helps convergence significantly:

Iteration 0: Cost 2.022
Iteration 10: Cost 0.204
Iteration 20: Cost 0.082
...
Iteration 90: Cost 0.052
Final w: 2.994, b: 2.012

The cost decreases faster and smoother than with standard mini-batch gradient descent, and the final parameter values are very close to the ground truth. Momentum is a simple but effective addition that is used in almost all neural network optimizers.

Advanced Optimizers

Building on the concepts of mini-batches and momentum, researchers have developed many more advanced optimization algorithms tailored for deep learning. Some of the most popular include:

• Adagrad: Adapts the learning rate per-parameter based on the historical gradients. Useful for sparse data.

• RMSprop: Modifies Adagrad to avoid shrinking the learning rate too aggressively. Works well for non-convex problems.

• Adam: Combines ideas from momentum and RMSprop. Good default choice that works well across many problems.

These optimizers further refine the gradient descent update rule to speed up convergence and improve stability. They are often the methods of choice for training large-scale deep learning models. However, the basic concepts of gradients, mini-batches, and momentum still form the foundation of these techniques.

Conclusion and Recommendations

Gradient descent is a cornerstone of machine learning that underlies most optimization problems. While the core idea is simple, there are many variants that build on it to improve the speed and stability of convergence. Here are some recommendations on when to use each variant:

• Batch gradient descent: Use for small datasets where computational cost is not an issue. Gives the most stable convergence.

• Stochastic gradient descent: Use for very large datasets or streaming data where it‘s not feasible to load the entire dataset at once. Can escape local minima but very erratic convergence.

• Mini-batch gradient descent: Use most of the time as the default. Good balance of efficiency and stability. Often used with a batch size of 32-512.

• Gradient descent with momentum: Use to accelerate mini-batch gradient descent. Especially helpful for noisy or ill-conditioned problems. Start with a momentum of 0.9.

• Advanced optimizers (Adagrad, RMSprop, Adam): Use for training deep neural networks. Adam is a good default choice.

I hope this post has demystified the various flavors of gradient descent and how they can be implemented in practice. The code examples provide a concrete starting point that you can adapt to your own models and datasets.

Remember, while it‘s important to choose a good optimizer, the most important factors are still going to be the model architecture and the quality of the data. Gradient descent is just a tool that helps bring it all together. Happy optimizing!