Gradient descent

Examples

We use torch because it computes gradients for us. We write down the loss, and backward() gives the partial derivatives with respect to every parameter.

The algorithm

import numpy as np
import torch
import matplotlib.pyplot as plt

torch.manual_seed(6)


def gradient_descent(loss, x, eta, T, track=False):
    path = []
    for t in range(T):
        loss(x).backward()
        with torch.no_grad():
            x -= eta * x.grad
        x.grad.zero_()
        if track:
            path.append(x.detach().clone().numpy())
    return (x, np.array(path)) if track else x

1: Compute the gradient of the loss with respect to every element of x that has requires_grad=True.
2: Take a step against the gradient multiplied by the learning rate eta.
3: Reset the gradient. torch accumulates gradients, so without this the next step would use the sum of all previous gradients.
4: Use track=True to keep the whole history of x values along the gradient descent path. detach() is needed to extract only the values from x, clone() is needed to copy the values, and numpy() to convert them to numpy arrays.

Linear regression

For linear regression the loss is the mean squared error, \[ \mathcal L(\beta_0, \beta_1) = \frac1n\sum_{i=1}^n \big(y_i - \beta_0 - \beta_1 x_i\big)^2 , \] here on 40 points from \(y = 0.1 - 0.4x + \varepsilon\), with \(\varepsilon\) normal with standard deviation 0.1. The panel runs 100 steps of the loop above from a starting point of your choice.

Top left, the data and the fit at step \(t\). Bottom left, the loss at every step, with step \(t\) in red. Right, the level lines of the loss and the whole path of gradient descent, with step \(t\) in red.

If the learning rate is too small, we need many steps. If it is too large, the loss goes up instead of down: here, somewhere between a learning rate of 0.8 and 0.9, the path starts to jump from one side of the valley to the other, further out each time. There is no formula for the right value. We try a few and look at the learning curve.

Logistic regression

The model is the one from week 2, \(p(x) = \sigma(\beta_0 + \beta_1 x)\), the probability of class 1. The loss is the negative log-likelihood of the Bernoulli model, averaged over the data: \[ \mathcal L(\beta_0, \beta_1) = -\frac1n\sum_{i=1}^n \Big(y_i\log p(x_i) + (1 - y_i)\log\big(1 - p(x_i)\big)\Big) . \] With the labels coded as \(y_i \in \{0, 1\}\), one of the two terms is zero for every point: each point contributes \(-\log\) of the probability the model gives to its observed label. There is no closed form for the minimum, but gradient descent only needs the loss.

def log_reg_loss(X, y):
    def loss(beta):
        p = torch.sigmoid(beta[0] + beta[1] * X)
        return -torch.mean(y * torch.log(p) + (1 - y) * torch.log(1 - p))
    return loss


rng = np.random.default_rng(61)
x2 = rng.standard_normal(500)
y2 = 1 / (1 + np.exp(-(-0.1 + 1.7 * x2))) > rng.random(500)
X2 = torch.tensor(x2)
Y2 = torch.tensor(y2, dtype=torch.float64)

beta2 = torch.tensor([0.1, 0.2], dtype=torch.float64, requires_grad=True)
beta2 = gradient_descent(log_reg_loss(X2, Y2), beta2, eta=0.5, T=500)
print("gradient descent:", beta2.detach().numpy().round(3))

1: Labels from a logistic generator with \(\beta_0 = -0.1\) and \(\beta_1 = 1.7\): a point is labelled 1 with probability \(\sigma(-0.1 + 1.7x)\).
2: The labels as the numbers 0 and 1, so that the loss can be written as in the formula above.

gradient descent: [-0.141  1.643]

With 500 points we get close to the generator’s \(-0.1\) and \(1.7\). The panel shows the same data and the same loss.

The same panel for the logistic regression on the 500 points above. Top left, the labels and the fitted probability at step \(t\).

Convex and non-convex

The losses of linear and logistic regression are convex. They have one minimum, and gradient descent finds it from any starting point, if the learning rate is small enough.

Not every loss is convex. Take 50 inputs from a standard normal distribution and noise-free outputs \(y = 0.3\sin(2x) + 0.7\sin(3.5x + 1)\), and fit the model \(\theta_1\sin(\theta_2 x + \theta_3) + \theta_4\sin(\theta_5 x + \theta_6)\) with the loss \[ \mathcal L(\theta) = \frac1n\sum_{i=1}^n\Big(y_i - \theta_1\sin(\theta_2 x_i + \theta_3) - \theta_4\sin(\theta_5 x_i + \theta_6)\Big)^2 . \] The loss is zero at \(\theta = (0.3, 2, 0, 0.7, 3.5, 1)\), the generator. But it is also zero at \((0.7, 3.5, 1, 0.3, 2, 0)\), where the two terms swap places, at every setting that flips the signs inside and outside a sine, since \(\sin(-u) = -\sin(u)\), and at every setting that shifts a phase by \(2\pi\). Which of these does gradient descent find?

10⁴ steps of gradient descent with learning rate 0.1 from a random starting point, 2 times a standard normal for every parameter. Left, the data and the fit at step \(t\). Top right, the loss at every step. Bottom right, the loss against \(\theta_2\) and \(\theta_5\) with the other four parameters held at their values at step \(t\), and the current \(\theta_2\), \(\theta_5\) in red. The six parameters cannot be drawn at once.

14 of the 20 seeds end at a loss below \(10^{-7}\), but not at the same parameters: seed 5 finds \(-0.70\sin(-3.50x - 1.00) + 0.30\sin(2.00x)\), the generator with the terms swapped and the signs flipped. The other 6 seeds end in a local minimum, with a loss of 0.008 to 0.02 and a fit that misses the data in places. Where gradient descent ends depends on where it starts.

In practice non-convexity matters less than it sounds. In high dimensions most local minima of a model with many parameters give a similar loss. But it does mean that two runs with different initial values give different models, and that we have to set a seed if we want to reproduce a result.

The next page computes the gradient on a part of the data in every step.