Learning Rate - From Neuron to GPT

In the previous two chapters, we trained the model with gradient descent for the first time. The direction was right, but a single step moved the weights from [0.5, 0.5] to [6854, 14147], the loss exploded, and training failed.

The root of the problem is this: the gradient only tells us which way is downhill. It does not tell us how big a step to take.

A large gradient only means that the slope in the current direction is steep. It does not mean we should take a big step. If the step is too large, we jump straight over the valley and land on the opposite slope, even higher than where we started. This phenomenon is called divergence.

The solution is to introduce a proportionality coefficient that scales down the size of each update step. This coefficient is called the learning rate, written $\eta$ .

With the learning rate, the gradient descent update formulas become:

w_{\text{new}} = w_{\text{old}} - \eta \cdot \frac{\partial L}{\partial w}

(1)

b_{\text{new}} = b_{\text{old}} - \eta \cdot \frac{\partial L}{\partial b}

(2)

The learning rate has to be chosen carefully:

Too large: the steps are too big, and the loss oscillates from side to side or even diverges;
Too small: the steps are too small, so the loss decreases extremely slowly and training takes a very long time to converge;
Just right: the loss decreases steadily and finally converges to the neighborhood of the optimum.

Hyperparameters¶

The learning rate is the first hyperparameter we have met.

The model parameters (the weights and the bias) are what the model learns automatically: after each training step, they are updated according to the gradients.

Hyperparameters are different: they are not updated by model training. We have to set them before training starts, based on experience and the progress of our experiments. The learning rate controls the size of the training steps, and we will meet other hyperparameters in later chapters.

The process of repeatedly testing and adjusting hyperparameters to get the best training result is called hyperparameter tuning. There are two kinds of tuning, which are usually combined:

Manual tuning: before training, set the hyperparameters by hand based on test data and experience;
Automatic tuning: during training, the program code adjusts the hyperparameters according to the training progress.

Optimizer¶

With the learning rate introduced, in practice we usually use an optimizer to manage the use of the learning rate (the parameter update) and automatic tuning in one place.

import numpy as np

Tensor¶

class Tensor:

    def __init__(self, data):
        self.data = np.array(data)
        self.grad = np.zeros_like(self.data)
        self.gradient_fn = None
        self.parents = []

    def backward(self):
        if self.gradient_fn is not None:
            self.gradient_fn()

        for p in self.parents:
            p.backward()

    def __str__(self):
        return f'Tensor({self.data})'

Data¶

feature = Tensor([28.1, 58.0])
label = Tensor([165])

Model¶

With the optimizer, we no longer need the model’s backward function to update the parameters one by one. Instead, the model provides a new attribute:

parameters (the parameter list): all the parameters of this model that need to be updated. For a linear model, these are the weights and the bias.

class Linear:

    def __init__(self, in_size, out_size):
        self.weight = Tensor(np.ones((out_size, in_size)) / in_size)
        self.bias = Tensor(np.zeros(out_size))

    def __call__(self, x: Tensor):
        return self.forward(x)

    def forward(self, x: Tensor):
        p = Tensor(x.data @ self.weight.data.T + self.bias.data)

        def gradient_fn():
            self.weight.grad += p.grad * x.data
            self.bias.grad += np.sum(p.grad)

        p.gradient_fn = gradient_fn
        return p

    # List of trainable model parameters
    @property
    def parameters(self):
        return [self.weight, self.bias]

Loss function (mean squared error)¶

class MSELoss:

    def __call__(self, p: Tensor, y: Tensor):
        return self.loss(p, y)

    def loss(self, p: Tensor, y: Tensor):
        mse = Tensor(np.mean(np.square(y.data - p.data)))

        def gradient_fn():
            p.grad += -2 * (y.data - p.data)

        mse.gradient_fn = gradient_fn
        mse.parents = [p]
        return mse

Optimizer (stochastic gradient descent)¶

The first optimizer we implement is the stochastic gradient descent (SGD) optimizer. Its job is very simple: after each training step (one forward propagation and one backpropagation), it immediately uses the gradients computed by backpropagation, together with the learning rate, to update the values of all the parameters.

To create an optimizer, we need to know:

parameters (the parameter list): all the parameters of the network model that need to be updated, such as the weights and the bias. Note that only the trainable parameters are included here. Intermediate values and predictions are tensors too, and their gradients are also computed during backpropagation, but they do not need to be updated by the optimizer;
lr (the learning rate): the proportionality coefficient of the parameter update.

Step method¶

The step method (step) implements the most basic function of the optimizer: updating the parameters.

The step method of the stochastic gradient descent optimizer is simple: before updating the parameters, it multiplies the gradients by the learning rate to reduce the size of each gradient descent step.

class SGDOptimizer:

    def __init__(self, parameters, lr):
        self.parameters = parameters
        self.lr = lr

    # Step method (update the parameters)
    def step(self):
        for p in self.parameters:
            p.data -= p.grad * self.lr

Hyperparameters¶

Learning rate¶

The stochastic gradient descent optimizer is simple and practical, but it has no automatic tuning. So we rely mainly on manual tuning.

Manually tuning the learning rate is a process of experimenting. Usually we pick an initial value based on the model, for example 0.01. Then we increase it (for example, to 0.1) or decrease it (for example, to 0.001) according to the training results.

💡 Tip: The learning rate we use here is very small, and that is the result of manual tuning. The main reason is that we have not applied standardization to the training data, which makes the computed gradients too large.

LEARNING_RATE = 0.00001

Modeling¶

We need to create an instance of the optimizer. Creating it requires passing in the list of model parameters and the learning rate hyperparameter.

model = Linear(2, 1)
loss_fn = MSELoss()
optimizer = SGDOptimizer(model.parameters, lr=LEARNING_RATE)

Training¶

In model training, we no longer call the model’s backward function (it has been removed). Instead, we call the optimizer’s step method (step) to update the parameters.

# Forward propagation (build the computational graph)
prediction = model(feature)
# Compute the loss
loss = loss_fn(prediction, label)
# Backpropagation (compute the gradients)
loss.backward()
# Update the parameters
optimizer.step()

Inference¶

prediction = model(feature)
print(f'prediction:\t{prediction}')

prediction:	Tensor([53.18309379])

Evaluation¶

loss = loss_fn(prediction, label)
print(f'loss:\t{loss}')

loss:	Tensor(12503.020514375934)

The evaluation shows that the loss dropped from 1,026,548,766,283 to 12,503. It is back in the normal range, and it is even smaller than the 14,872 from before training. This shows that introducing the learning rate works, and that the value we chose (by manual tuning) is reasonable.

Exercises¶

The design of the parent list (parents) and the parameter list (parameters) makes the traversal of the computational graph, backpropagation, and the parameter update fully automatic. Try removing different parameters from the parent list or the parameter list. What happens to the result of backpropagation?
Change the learning rate to 0.001, 0.0001, and 0.000001 respectively, train once with each, and observe how the loss relates to the learning rate.