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 tensor that has requires_grad=True.
2: Take a step against the gradient.
3: Reset the gradient. torch accumulates gradients, so without this the next step would use the sum of all previous gradients.

Linear regression

def lin_reg_loss(X, y):
    def loss(beta):
        return torch.mean((y - beta[0] - beta[1] * X) ** 2)
    return loss


X = torch.randn(40, dtype=torch.float64)
y = 0.1 - 0.4 * X + 0.1 * torch.randn(40, dtype=torch.float64)

beta = torch.tensor([0.1, 0.2], requires_grad=True)
beta, path = gradient_descent(lin_reg_loss(X, y), beta, eta=0.1, T=100, track=True)
print("gradient descent:", beta.detach().numpy().round(4))

1: We want the partial derivatives with respect to these numbers, so requires_grad=True.

gradient descent: [ 0.09   -0.4373]

from sklearn.linear_model import LinearRegression

m = LinearRegression().fit(X.numpy().reshape(-1, 1), y.numpy())
print("least squares:   ", np.array([m.intercept_, m.coef_[0]]).round(4))

least squares:    [ 0.09   -0.4373]

The two agree. Gradient descent found the same answer as the closed form, just more slowly.

The path of gradient descent on the loss surface, dark where the loss is small. Move the learning rate and the number of steps. Below about 0.9 it converges, above 1 it diverges.

The learning rate

If the learning rate is too small, we need many steps. If it is too large, the loss goes up instead of down. There is no formula for the right value. We try a few and look at the learning curve.

Logistic regression

Nothing changes except 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) + (~y) * torch.log(1 - p))
    return loss


X2 = torch.randn(500, dtype=torch.float64)
y2 = torch.sigmoid(-0.1 + 1.7 * X2) > torch.rand(500, dtype=torch.float64)

beta2 = torch.tensor([0.1, 0.2], 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: The negative log-likelihood of the Bernoulli model. ~y negates the boolean tensor.

gradient descent: [-0.063  1.478]

The generator used \(-0.1\) and \(1.7\). With 500 points we get close.

Convex and non-convex

The loss of linear and logistic regression is convex. It has one minimum, and gradient descent finds it from any starting point.

The loss of a neural network is not convex. It has many local minima and many saddle points. Gradient descent finds one of them, and which one depends on where we started.

The same algorithm on a logistic regression. The loss is convex here too, so every starting point ends in the same place, but the path and the number of steps depend on where we begin.

In practice this matters less than it sounds. In high dimensions most local minima of a large network 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 same algorithm on a model that is not linear in its parameters. Change the seed and watch two runs end in visibly different fits, at visibly different losses.