sklearn and Gradient Descent

Gradient Descent¶

At this point, we’ve grown quite familiar with the process of choosing a model and a corresponding loss function and optimizing parameters by choosing the values of $\theta$ that minimize the loss function. So far, we’ve optimized $\theta$ by

1. Using calculus to take the derivative of the loss function with respect to $\theta$, setting it equal to 0, and solving for $\theta$.

2. Using the geometric argument of orthogonality to derive the OLS solution $\hat{\theta} = (\mathbb{X}^T \mathbb{X})^{-1}\mathbb{X}^T \mathbb{Y}$.

One thing to note, however, is that the techniques we used above can only be applied if we make some big assumptions. For the calculus approach, we assumed that the loss function was differentiable at all points and that we could algebraically solve for the zero points of the derivative; for the geometric approach, OLS only applies when using a linear model with MSE loss. What happens when we have more complex models with different, more complex loss functions? The techniques we’ve learned so far will not work, so we need a new optimization technique: gradient descent.

BIG IDEA: use an iterative algorithm to numerically compute the minimum of the loss.

Minimizing an Arbitrary 1D Function¶

Let’s consider an arbitrary function. Our goal is to find the value of $x$ that minimizes this function.

How do we find the value of $x$ that minimizes the following function?

How does a droplet of water “know” how to reach the bottom?

It will follow the path of least resistance:

Eventually, if the droplet reaches a local minimum, the path of least resistance will be to remain where it is.

The driving instructions for the droplet are as follows:

1. Face the path of least resistance.

2. Move in that direction a little bit.

3. Iterate until the path of least resistance is to stay put.

Digging into Gradient Descent¶

Looking at the function across this domain, it is clear that the function’s minimum value occurs around $\theta = 5.3$. Let’s pretend for a moment that we couldn’t see the full view of the cost function. How would we guess the value of $\theta$ that minimizes the function?

It turns out that the first derivative of the function can give us a clue. In the graph below, the function and its derivative are plotted, with points where the derivative is equal to 0 plotted in light green.

import plotly.graph_objects as go

def derivative_arbitrary(x):
    return (4*x**3 - 45*x**2 + 160*x - 180)/10

fig = go.Figure()
roots = np.array([2.3927, 3.5309, 5.3263])

fig.add_trace(go.Scatter(x = xs, y = arbitrary(xs), 
                         mode = "lines", name = "f"))
fig.add_trace(go.Scatter(x = xs, y = derivative_arbitrary(xs), 
                         mode = "lines", name = "df", line = {"dash": "dash"}))
fig.add_trace(go.Scatter(x = np.array(roots), y = 0*roots, 
                         mode = "markers", name = "df = zero", marker_size = 12))
fig.update_layout(font_size = 20, yaxis_range=[-1, 3])
fig.update_layout(autosize=False, width=800, height=600)
fig.show()

In the plots below, the colored tangent line shows the local slope of the curve at the current value of $\theta$.
A red tangent line means the slope is negative, and a green tangent line means the slope is positive.

To lower the cost, we want to move in the direction that goes downhill.

The plots show the key intuition:

• when the slope is negative, the curve goes down as we move right

• when the slope is positive, the curve goes down as we move left

So:

• negative slope $\rightarrow$ right is downhill

• positive slope $\rightarrow$ left is downhill

In this example, the slope is negative, so increasing $\theta$ moves us downhill.

In this example, the slope is positive, so decreasing $\theta$ moves us downhill.

So the derivative tells us the direction of each update:

• negative derivative $\rightarrow$ move right

• positive derivative $\rightarrow$ move left

Algorithm Attempt 1¶

Armed with this knowledge, let’s try to see if we can use the derivative to optimize the function.

We start by making some guess for the minimizing value of $x$. Then, we look at the derivative of the function at this value of $x$, and step downhill in the opposite direction. We can express our new rule as a recurrence relation:

Translating this statement into English: we obtain our next guess for the minimizing value of $x$ at timestep $t+1$ ($x^{(t+1)}$) by taking our last guess ($x^{(t)}$) and subtracting the derivative of the function at that point ($\frac{d}{dx} f(x^{(t)})$).

A few steps are shown below, where the old step is shown as a transparent point, and the next step taken is the green-filled dot.

Looking pretty good! We do have a problem though – once we arrive close to the minimum value of the function, our guesses “bounce” back and forth past the minimum without ever reaching it.

In other words, each step we take when updating our guess moves us too far. We can address this by decreasing the size of each step.

Algorithm Attempt 2¶

Let’s update our algorithm to use a learning rate (also sometimes called the step size), which controls how far we move with each update. We represent the learning rate with $\alpha$.

A small $\alpha$ means that we will take small steps; a large $\alpha$ means we will take large steps. When do we stop updating? We stop updating either after a fixed number of updates or after a subsequent update doesn’t change much.

Updating our function to use $\alpha=0.3$, our algorithm successfully converges (settles on a solution and stops updating significantly, or at all) on the minimum value.

Convexity¶

In our analysis above, we focused our attention on the global minimum of the loss function. You may be wondering: what about the local minimum that’s just to the left?

If we had chosen a different starting guess for $\theta$, or a different value for the learning rate $\alpha$, our algorithm may have gotten “stuck” and converged on the local minimum, rather than on the true optimum value of loss.

If the loss function is convex, gradient descent is guaranteed to converge and find the global minimum of the objective function. Formally, a function $f$ is convex if:

for all $a, b$ in the domain of $f$ and $t \in [0, 1]$.

To put this into words: if you drew a line between any two points on the curve, all values on the curve must be on or below the line. Importantly, any local minimum of a convex function is also its global minimum so we avoid the situation where the algorithm converges on some critical point that is not the minimum of the function.

In summary, non-convex loss functions can cause problems with optimization. This means that our choice of loss function is a key factor in our modeling process. It turns out that MSE is convex, which is a major reason why it is such a popular choice of loss function. Gradient descent is only guaranteed to converge (given enough iterations and an appropriate step size) for convex functions.

Gradient Descent in 1 Dimension¶

Terminology clarification: In past lectures, we have used “loss” to refer to the error incurred on a single datapoint. In applications, we usually care more about the average error across all datapoints. Going forward, we will take the “model’s loss” to mean the model’s average error across the dataset. This is sometimes also known as the empirical risk (R), cost function, or objective function.

$$L(\theta) = R(\theta) = \frac{1}{n} \sum_{i=1}^{n} L(y_{i}, \hat{y}_{i})$$

In our discussion above, we worked with some arbitrary function $f$. As data scientists, we will almost always work with gradient descent in the context of optimizing models – specifically, we want to apply gradient descent to find the minimum of a loss function. In a modeling context, our goal is to minimize a loss function by choosing the minimizing model parameters.

Recall our modeling workflow from the past few lectures:

1. Define a model with some parameters $\theta_i$

2. Choose a loss function

3. Select the values of $\theta_i$ that minimize the loss function on the data

Gradient descent is a powerful technique for completing this last task. By applying the gradient descent algorithm, we can select values for our parameters $\theta_i$ that will lead to the model having minimal loss on the training data.

When using gradient descent in a modeling context, we:

1. Make guesses for the minimizing $\theta_i$

2. Compute the derivative of the loss function $L$

We can “translate” our gradient descent rule from before by replacing $x$ with $\theta$ and $f$ with $L$:

Gradient Descent on the tips Dataset¶

To see this in action, let’s consider a case where we have a linear model with no offset. We want to predict the tip (y) given the price of a meal (x). To do this, we

• Choose a model: $\hat{y} = \theta_1 x$,

• Choose a loss function: $L(\theta) = MSE(\theta) = \frac{1}{n} \sum_{i=1}^n (y_i - \theta_1x_i)^2$.

Let’s apply our gradient_descent function from before to optimize our model on the tips dataset. We will try to select the best parameter $\theta_i$ to predict the tip $y$ from the total_bill $x$.

We can visualize the value of the MSE on our dataset for different possible choices of $\theta_1$. To optimize our model, we want to select the value of $\theta_1$ that leads to the lowest MSE.

To apply gradient descent, we need to compute the derivative of the loss function with respect to our parameter $\theta_1$.

• Given our loss function,

  $$L(\theta) = MSE(\theta) = \frac{1}{n} \sum_{i=1}^n (y_i - \theta_1x_i)^2$$

• We take the derivative with respect to $\theta_1$

  $$\frac{\partial}{\partial \theta_{1}} L(\theta_1^{(t)}) = \frac{-2}{n} \sum_{i=1}^n (y_i - \theta_1^{(t)} x_i) x_i$$

• Which results in the gradient descent update rule

  $$\theta_1^{(t+1)} = \theta_1^{(t)} - \alpha \frac{d}{d\theta_1}L(\theta_1^{(t)})$$

for some learning rate $\alpha$.

Implementing this in code, we can visualize the MSE loss on the tips data. MSE is convex, so there is one global minimum.

def gradient_descent(df, initial_guess, alpha, n):
    """Performs n steps of gradient descent on df using learning rate alpha starting
       from initial_guess. Returns a numpy array of all guesses over time."""
    guesses = [initial_guess]
    current_guess = initial_guess
    while len(guesses) < n:
        current_guess = current_guess - alpha * df(current_guess)
        guesses.append(current_guess)
        
    return np.array(guesses)

def mse_single_arg(theta_1):
    """Returns the MSE on our data for the given theta1"""
    x = df["total_bill"]
    y_obs = df["tip"]
    y_hat = theta_1 * x
    return np.mean((y_hat - y_obs) ** 2)

def mse_loss_derivative_single_arg(theta_1):
    """Returns the derivative of the MSE on our data for the given theta1"""
    x = df["total_bill"]
    y_obs = df["tip"]
    y_hat = theta_1 * x
    
    return np.mean(2 * (y_hat - y_obs) * x)

loss_df = pl.DataFrame({"theta_1":np.linspace(-1.5, 1), "MSE":[mse_single_arg(theta_1) for theta_1 in np.linspace(-1.5, 1)]})

trajectory = gradient_descent(mse_loss_derivative_single_arg, -0.5, 0.0001, 100)

plt.plot(loss_df["theta_1"], loss_df["MSE"])
plt.scatter(trajectory, [mse_single_arg(guess) for guess in trajectory], c="white", edgecolor="firebrick")
plt.scatter(trajectory[-1], mse_single_arg(trajectory[-1]), c="firebrick")
plt.xlabel(r"$\theta_1$")
plt.ylabel(r"$L(\theta_1)$");

print(f"Final guess for theta_1: {trajectory[-1]}")

Gradient Descent on Multi-Dimensional Models¶

The function we worked with above was one-dimensional – we were only minimizing the function with respect to a single parameter, $\theta$. However, models usually have a cost function with multiple parameters that need to be optimized. For example, simple linear regression has 2 parameters:

and multiple linear regression has $p+1$ parameters:

We’ll need to expand gradient descent so we can update our guesses for all model parameters all in one go.

With multiple parameters to optimize, we consider a loss surface, or the model’s loss for a particular combination of possible parameter values.

import seaborn as sns
import numpy as np
import polars as pl
df = pl.from_pandas(sns.load_dataset("tips"))

We can also visualize a bird’s-eye view of the loss surface from above using a contour plot:

contour = go.Contour(x=u[0], y=v[:, 0], z=np.reshape(MSE, u.shape))
fig = go.Figure(contour)
fig.update_layout(
    xaxis_title = "theta0",
    yaxis_title = "theta1", autosize=False, width=800, height=600)

fig.show()

The Gradient Vector¶

As before, the derivative of the loss function tells us the best way towards the minimum value.

On a 2D (or higher) surface, the best way to go down (gradient) is described by a vector.

For the vector of parameter values $\vec{\theta} = \begin{bmatrix} \theta_{0} \\ \theta_{1} \\ \end{bmatrix}$, we take the partial derivative of loss with respect to each parameter: $\frac{\partial L}{\partial \theta_0}$ and $\frac{\partial L}{\partial \theta_1}$.

For example, consider the 2D function:

$$f(\theta_0, \theta_1) = 8 \theta_0^2 + 3\theta_0\theta_1$$

For a function of 2 variables $f(\theta_0, \theta_1)$, we define the gradient

The gradient vector of a generic function of $p+1$ variables is therefore

Recall the 1D case, for example, if $\frac{dy}{dx} = -5$, then increasing $x$ will decrease $y$, while decreasing $x$ will increase $y$.

The gradient in 1-D can be written as:

In vector form, the gradient can be written as:

This vector points in the direction of steepest ascent of $y$. The corresponding unit ascent direction is:

The direction of steepest descent is given by the negative gradient:

which points in the positive $x$-direction.

Let $L(\theta_1, \theta_2, \dots, \theta_n)$ be a loss function. Its gradient is:

If $\frac{\partial L}{\partial \theta_i} > 0$, then increasing $\theta_i$ increases $L$, so to decrease $L$, we should decrease $\theta_i$. We can accomplish this by subtracting the value of $\frac{\partial L}{\partial \theta_i}$, which is positive.

If $\frac{\partial L}{\partial \theta_i} < 0$, then increasing $\theta_i$ decreases $L$, so to decrease $L$, we should increase $\theta_i$. We can accomplish this by subtracting the value of $\frac{\partial L}{\partial \theta_i}$, which is negative.

Visually, the gradient can be represented as shown below on a 3D loss surface:

For a sufficiently small learning rate $\eta > 0$, moving in the direction of the gradient increases the loss. Thus, to minimize the loss, we move in the opposite direction of the gradient:

At this point, we see that the core idea of using the gradient to guide the next step is the same in both 1D and multidimensional settings. Therefore, we can naturally extend the same update rule from the 1D case.

• Recall our 1D update rule:

  $$\theta^{(t+1)} = \theta^{(t)} - \alpha \frac{d}{d\theta}L(\theta^{(t)})$$

• For models with multiple parameters, we work in terms of vectors:

  $$\begin{bmatrix} \theta_{0}^{(t+1)} \\ \theta_{1}^{(t+1)} \\ \vdots \end{bmatrix} = \begin{bmatrix} \theta_{0}^{(t)} \\ \theta_{1}^{(t)} \\ \vdots \end{bmatrix} - \alpha \begin{bmatrix} \frac{\partial L}{\partial \theta_{0}} \\ \frac{\partial L}{\partial \theta_{1}} \\ \vdots \\ \end{bmatrix}$$

• Written in a more compact form,

  $$\vec{\theta}^{(t+1)} = \vec{\theta}^{(t)} - \alpha \nabla_{\vec{\theta}} L(\theta^{(t)})$$

  • $\theta$ is a vector with our model weights

  • $L$ is the loss function

  • $\alpha$ is the learning rate (ours is constant, but other techniques use an $\alpha$ that decreases over time)

  • $\vec{\theta}^{(t)}$ is the current value of $\theta$

  • $\vec{\theta}^{(t+1)}$ is the next value of $\theta$

  • $\nabla_{\vec{\theta}} L(\vec{\theta}^{(t)})$ is the gradient of the loss function evaluated at the current $\theta$:

    $$\frac{1}{n}\sum_{i=1}^{n}\nabla_{\vec{\theta}} l(y_i, f_{\vec{\theta}^{(t)}}(X_i))$$

Note that $-\nabla_{\vec{\theta}} L$ points in the direction we should move $\vec{\theta}$ to make the loss decrease most quickly.

Let’s now walk through an example of calculating and updating the gradient vector. Say our model and loss are:

Plugging in $f_{\vec{\theta}}(\vec{x})$ for $\hat{y}$, our loss function becomes $l(\vec{\theta}, \vec{x}, y_i) = (y_i - \theta_0x_0 - \theta_1x_1)^2$.

To calculate our gradient vector, we can start by computing the partial derivative of the loss function with respect to $\theta_0$:

Let’s now do the same but with respect to $\theta_1$:

Putting this together, our gradient vector is:

Remember that we need to keep updating $\theta$ until the algorithm converges to a solution and stops updating significantly (or at all). When updating $\theta$, we’ll have a fixed number of updates and subsequent updates will be quite small (we won’t change $\theta$ by much).

Stochastic (Mini-batch) Gradient Descent¶

Formally, the algorithm we derived above is called batch gradient descent. For each iteration of the algorithm, the derivative of loss is computed across the entire batch of all $n$ datapoints. While this update rule works well in theory, it is not practical in most circumstances. For large datasets (with perhaps billions of datapoints), finding the gradient across all the data is incredibly computationally taxing; gradient descent will converge slowly because each individual update is slow.

Stochastic (mini-batch) gradient descent tries to address this issue. In stochastic descent, only a sample of the full dataset is used at each update. We estimate the true gradient of the loss surface using just that sample of data. The batch size is the number of data points used in each sample. The sampling strategy is generally without replacement (data is shuffled and batch size examples are selected one at a time.)

Each complete “pass” through the data is known as a training epoch. After shuffling the data, in a single training epoch of stochastic gradient descent, we

• Compute the gradient on the first x% of the data. Update the parameter guesses.

• Compute the gradient on the next x% of the data. Update the parameter guesses.

• $\dots$

• Compute the gradient on the last x% of the data. Update the parameter guesses.

Every data point appears once in a single training epoch. We then perform several training epochs until we’re satisfied.

Batch gradient descent is a deterministic technique – because the entire dataset is used at each update iteration, the algorithm will always advance towards the minimum of the loss surface. In contrast, stochastic gradient descent involves an element of randomness. Since only a subset of the full data is used to update the guess for $\vec{\theta}$ at each iteration, there’s a chance the algorithm will not progress towards the true minimum of loss with each update. Over the longer term, these stochastic techniques should still converge towards the optimal solution.

The diagrams below represent a “bird’s eye view” of a loss surface from above. Notice that batch gradient descent takes a direct path towards the optimal $\hat{\theta}$. Stochastic gradient descent, in contrast, “hops around” on its path to the minimum point on the loss surface. This reflects the randomness of the sampling process at each update step.

Batch, Mini-batch, Stochastic Gradient Descent¶

The time complexity of a single gradient descent step takes only $O(nd)$ time where $n$ is the number of samples (rows) and $d$ is the number of features (columns).

Suppose we run $T$ iterations. The final complexity would then be $O(Tnd)$. Typically, $n$ is much larger than $T$ and $d$. How can we reduce the cost of this algorithm using a technique from Data 100? Do we really need to use $n$ data points? We don’t! Instead, we can use stochastic gradient descent.

We know that our true gradient of $\nabla_{\vec{\theta}} L (\vec{\theta^{(t)}}) = \frac{1}{n}\sum_{i=1}^{n}\nabla_{\vec{\theta}} l(y_i, f_{\vec{\theta}^{(t)}}(X_i))$ has a time complexity of $O(nd)$. Instead of using all $n$ samples to calculate the true gradient of the loss surface, let’s use a probability sample of our data to approximate the gradient!

Say we sample $b$ records ($s_1, \cdots, s_b$) from our $n$ datapoints. Our new (stochastic) gradient descent function will be:

and will now have a time complexity of $O(bd)$, which is much faster! For more on computational complexity of the normal equation, see the bonus section in the OLS chapter.

Summarizing our two gradient descent techniques:

• (Batch) Gradient Descent: GD computes the true descent and always descends towards the true minimum of the loss based on a batch (a full dataset). While accurate, it can often be computationally expensive.

• (Mini-batch) Stochastic Gradient Descent: SGD approximates the true gradient descent based off of mini-batch (a subset of a dataset). It may not descend towards the true minimum with each update, but it’s often less computationally expensive than batch gradient descent. Note that it is mini-batch, because we create a smaller, randomized sample of the batch to run SGD on it.

To summarize the tradeoffs of batch size:

The typical solution is to set batch size to ensure sufficient hardware utilization.

[Bonus] Stochastic Gradient Descent in PyTorch¶

While this material is out of scope for Data 100, it is useful if you plan to enter a career in data science!

In practice, you will use software packages such as PyTorch when computing gradients and implementing gradient descent. You’ll often follow three main steps:

1. Sample a batch of the data.

2. Compute the loss and the gradient.

3. Update your gradient until you reach an appropriate estimate of the true gradient.

If you want to learn more, this Intro to PyTorch tutorial is a great resource to get started!