Run an optimizer step

Learn it

Words to know

parameter

Any number in the model that training is allowed to change. Weights and biases are parameters.

optimizer

The part of the training code that changes the parameters, using their gradients.

SGD

Stochastic gradient descent, the plainest optimizer: each parameter moves downhill by the learning rate times its gradient.

backpropagation

Working out the gradient of every weight by going backwards from the loss, one step of the calculation at a time.

How it works

1. optimizer.zero_grad() clears the gradients left from the last step. PyTorch adds new gradients to old ones unless you clear them.
2. Compute the loss from the current parameter: here ((w * x - y) ** 2).mean().
3. loss.backward() fills in w.grad, the gradient of the loss for w.
4. optimizer.step() updates the parameter. With SGD that is w = w - rate * w.grad.

Worked example

Input
xs = [1, 2]
ys = [2, 4]
rate = 0.1
steps = 2

Output
1.5

The loss falls from 10 to 2.5 as w moves toward 2, the true answer. The gradient halves each step, so the steps shrink too.

Why it matters

Training a GPT is this same loop: clear, loss, backward, step. The workshop's GPT uses another optimizer, AdamW, in place of SGD, and adds one line that limits the size of the gradients.

Your turn

Using PyTorch, fit one weight w so that w * x matches y, with mean squared error and SGD. Start from w = 0.

xs

the inputs; at least one number

ys

the targets, one per input

rate

the learning rate

steps

how many optimizer steps to run; 0 or more

Returns

The value of w after the last step, as a Python float.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
xs = [1, 2]
ys = [2, 4]
rate = 0.1
steps = 1

Output
1

Input
xs = [1, 2]
ys = [2, 4]
rate = 0.1
steps = 2

Output
1.5

Input
xs = [1]
ys = [2]
rate = 0.1
steps = 0

Output
0

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   Set up three things before the loop: the data as tensors, the weight as a torch.nn.Parameter, and an SGD optimizer that holds that weight.

2. Hint 2: The key step

   Inside the loop the order is fixed: clear the gradients, compute the loss, call backward, call step.

3. Hint 3: The outline

   Make x and y float64 tensors and w = torch.nn.Parameter(torch.zeros((), dtype=torch.float64)). Set optimizer = torch.optim.SGD([w], lr=rate). Repeat steps times: optimizer.zero_grad(), loss = ((w * x - y) ** 2).mean(), loss.backward(), optimizer.step(). Return w.item().

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. A parameter is a tensor that PyTorch tracks gradients for. torch.zeros(()) is a single number with no dimensions. Wrapping it in Parameter tells PyTorch to record what is done with it.
2. Clear the gradient left over from the last step. backward adds to whatever is already stored. Without clearing, step 2 would use the sum of two gradients.
3. Predict, and measure the mean squared error. w * x multiplies every input by the one weight. The loss must be a single number for backward to start from.
4. Fill in w.grad. PyTorch walks back through the calculation of the loss and applies the chain rule, as you did by hand in chapter 2.
5. Move w against its gradient. SGD subtracts rate times the gradient, the same update as in the first lesson.

The same numbered steps are the comments in the solution.

import torch


def solve(xs, ys, rate, steps):
    x = torch.tensor(xs, dtype=torch.float64)
    y = torch.tensor(ys, dtype=torch.float64)
    # 1. A parameter is a tensor that PyTorch tracks gradients for
    w = torch.nn.Parameter(torch.zeros((), dtype=torch.float64))
    optimizer = torch.optim.SGD([w], lr=rate)
    for _ in range(steps):
        # 2. Clear the gradient left over from the last step
        optimizer.zero_grad()
        # 3. Predict, and measure the mean squared error
        loss = ((w * x - y) ** 2).mean()
        # 4. Fill in w.grad
        loss.backward()
        # 5. Move w against its gradient
        optimizer.step()
    return w.item()

Watch out

• Leaving out zero_grad() makes gradients pile up. In the worked example the second step would use -15 instead of -5 and jump to 2.5, past the true answer of 2.
• Return w.item(), a plain number. Returning w itself gives a tensor, which the checks do not accept.

Check yourself