The optimizer (SGD)

You have a number measuring the error and gradients pointing from it into every weight. What is missing is the step that touches those weights. It is so short that on first reading you ask whether that was really all of it: push every weight a little bit against its gradient.

Against, not with. The gradient points in the direction in which the loss grows

• it is the direction of steepest ascent. We want to go downhill, so we go backwards. That minus sign is the entire idea of the method, and the name says so too: gradient descent.

The goal

The learning step is one line. p.data -= lr * p.grad, for every parameter. The rest of the optimizer is administration: a list of parameters and a step size.

The step size lr (learning rate) is the only knob, and it has two wrong settings. Too small: training crawls and needs thousands of epochs. Too large: the step overshoots the valley, lands higher on the other side than it started, and the loss grows instead of falling. Between the two lies a surprisingly wide zone in which it simply works - which is why rates like 0.01 or 0.05 get you surprisingly far.

Zero magic in this step

NumPy is allowed, torch, jax, tinygrad and autograd are forbidden. torch.optim.SGD above all - a class that is not thirty lines in PyTorch and whose core you are about to write in one.

The exercise

A new file, meintorch/optim.py:

class SGD:
    """Stochastic gradient descent: the simplest optimizer there is."""

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

    def step(self):
        # TODO(du): push every parameter a little against its gradient.
        # On the raw array, not via tensor operations.
        raise NotImplementedError

    def zero_grad(self):
        # TODO(du): reset all gradients to zero. Without that, the next
        # backward pass would pile up on top of the old one.
        raise NotImplementedError

Hint 1 - the direction

step() is a loop over self.parameters with exactly one line inside. The sign is the only decision: the loss should get smaller, the gradient points where it gets bigger.

zero_grad() is a loop over the same list. Your tensor already knows how to do what has to happen inside it - you wrote that method in chapter 02.

Why zero_grad() has to exist at all: gradients are accumulated (+=, see chapter 02). Without cleaning up, the gradient in epoch 10 would hold the sum of all ten backward passes, the step would be ten times too large, and training derails. This is by far the most common bug in hand-built training loops - and in PyTorch code, for that matter.

Hint 2 - more concrete

• step(): for p in self.parameters: p.data -= self.lr * p.grad.
• zero_grad(): for p in self.parameters: p.zero_grad().
• list(parameters) in the constructor is deliberate: model.parameters() returns a list, but a generator would be empty after the first pass.

class SGD:
    def __init__(self, parameters, lr=0.01):
        self.parameters = list(parameters)
        self.lr = lr

    def step(self):
        for p in self.parameters:
            p.data -= self.lr * p.grad

    def zero_grad(self):
        for p in self.parameters:
            p.zero_grad()

In goal mode

The autonomous run does not only check that the loss falls, but also that the step has exactly the right size and leaves the graph alone:

Create `meintorch/optim.py` with a class `SGD(parameters, lr=0.01)`. `step()`
moves every parameter by `lr * grad` against its gradient, in-place on the
NumPy array, so that no new graph node appears. `zero_grad()` resets all
gradients.
Mandatory: NumPy allowed, torch/jax/tinygrad/autograd forbidden. No momentum,
no weight decay - bare gradient descent.
Write `tests/test_05_optimizer.py`: one step on a parameter with a known
gradient has to be exactly `lr * grad`, a step after a backward pass has to
measurably lower the loss of a small network, `zero_grad()` has to empty all
gradients, and after `step()` no parameter may have children in the graph.
Run `python -m pytest` until everything is green.

The milestone test

tests/test_05_optimizer.py. The second test is the milestone: a single step, and the loss is smaller than before. That is the moment where arithmetic turns into learning.

import numpy as np

from meintorch.loss import mse_loss
from meintorch.nn import Linear, ReLU, Sequential
from meintorch.optim import SGD
from meintorch.tensor import Tensor


def network_and_data():
    net = Sequential(Linear(3, 4, seed=0), ReLU(), Linear(4, 1, seed=1))
    x = Tensor(np.random.default_rng(7).normal(size=(8, 3)))
    y = Tensor(np.random.default_rng(8).normal(size=(8, 1)))
    return net, x, y


def test_the_step_is_exactly_lr_times_gradient():
    p = Tensor([[1.0, -2.0]])
    p.grad = np.array([[2.0, 4.0]])
    SGD([p], lr=0.1).step()
    assert np.allclose(p.data, [[0.8, -2.4]])


def test_one_step_lowers_the_loss():
    net, x, y = network_and_data()
    opt = SGD(net.parameters(), lr=0.05)
    before = float(mse_loss(net(x), y).data)

    loss = mse_loss(net(x), y)
    opt.zero_grad()
    loss.backward()
    opt.step()

    assert float(mse_loss(net(x), y).data) < before


def test_zero_grad_cleans_up_before_the_next_step():
    net, x, y = network_and_data()
    opt = SGD(net.parameters(), lr=0.05)
    mse_loss(net(x), y).backward()
    opt.zero_grad()
    for p in net.parameters():
        assert np.allclose(p.grad, 0.0)


def test_the_step_adds_nothing_to_the_graph():
    net, x, y = network_and_data()
    opt = SGD(net.parameters(), lr=0.05)
    mse_loss(net(x), y).backward()
    opt.step()
    for p in net.parameters():
        assert p._children == ()