Sweep the learning rate

LabEasy20 min15 checks

learning-rate gradient-descent optimization

Lesson: Learning Rate Schedules and Warmup Cheat sheet Track: Optimizers Track: Experiment labs

The learning rate decides if gradient descent converges, crawls or blows up. A rate sweep is the first experiment on any new model: train with each rate, keep the final loss, and pick the rate. In this lab you run the sweep on a small linear regression and compare the result with the bound that theory gives.

support.py gives the data, the grid of rates and the trainer. make_regression(seed) returns X of shape (64, 4) and y of shape (64,). train(X, y, lr, steps) runs full batch gradient descent on the mean squared error from $inline$ . It returns the loss before each step and after the last step, an array of steps + 1 values. When the loss passes 1e10, the trainer stops and fills the rest of the curve with inf.

Task

Write five helper functions and answers().

run_sweep(X, y, lrs, steps=STEPS) returns a float array of shape (len(lrs), steps + 1). Row $inline$ is train(X, y, lrs[i], steps).
is_diverged(losses) returns True when the final loss is not finite or is larger than the first loss. Otherwise it returns False.
best_lr(lrs, curves) returns the rate with the lowest final loss among the runs that did not diverge. If two runs have the same final loss, return the smaller rate. If every run diverged, return None.
smallest_diverging_lr(lrs, curves) returns the smallest rate, by value, whose run diverged. If no run diverged, return None.
stability_bound(X) returns $inline$ , where $inline$ is the largest eigenvalue of the Hessian $inline$ of the mean squared error.

Then do the experiment in answers().

Make the data with make_regression(seed=0).
Change the learning rate over LEARNING_RATES. Train each rate for STEPS steps.
Measure the final loss of each run. Report the best rate as best_lr and its final loss as best_final_loss.
Report the smallest rate that diverged as smallest_diverging_lr.
Compute the theory bound on the same X. Report it as stability_bound.

Input and output

Name	Type	Shape	Meaning
`X`	float array	`(n, d)`	the inputs, one row for each example
`y`	float array	`(n,)`	the targets
`lrs`	list of floats	`(k,)`	the learning rates, in any order
`steps`	int	scalar	the number of gradient steps
`losses`	float array	`(steps + 1,)`	one loss curve
`curves`	float array	`(k, steps + 1)`	one loss curve for each rate, in the order of `lrs`

run_sweep returns a float array of shape (k, steps + 1).
is_diverged returns a bool.
best_lr and smallest_diverging_lr return one element of lrs or None.
stability_bound returns a float.
answers() returns a dict with the four keys above. Each value is a Python float.

Constraints

A curve can hold inf or nan.
Do not change the input arrays.
Do not round the answers. The checks accept a relative difference of 1e-6 from the reference run.
Use numpy and the standard library only.

Examples

pythonis_diverged(np.array([1.0, 2.0, 3.0]))

textTrue

pythonbest_lr([0.1, 0.2, 0.3], np.array([[1.0, 0.5], [1.0, 0.2], [1.0, np.inf]]))

text0.2

pythonstability_bound(np.array([[1.0, 0.0], [0.0, 2.0]]))

text0.5

pythonX, y = make_regression(0)
train(X, y, 0.1, 3)

text[3.95117454 0.8025197  0.6543398  0.57749216]

Provided code: support.py

Your code and the checks import this module with from support import .... You do not change it.

import numpy as np

LEARNING_RATES = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1.0]
STEPS = 100
BLOW_UP = 1e10


def make_regression(seed=0, n=64):
    rng = np.random.default_rng(seed)
    scales = np.array([1.0, 0.5, 2.0, 0.25])
    X = rng.normal(size=(n, scales.size)) * scales
    w = rng.normal(size=scales.size)
    y = X @ w + 0.1 * rng.normal(size=n)
    return X, y


def mse(X, y, w):
    residual = X @ w - y
    return float(np.mean(residual ** 2))


def train(X, y, lr, steps=STEPS):
    n = X.shape[0]
    w = np.zeros(X.shape[1])
    losses = []
    for t in range(steps + 1):
        loss = mse(X, y, w)
        if not np.isfinite(loss) or loss > BLOW_UP:
            losses.extend([np.inf] * (steps + 1 - len(losses)))
            break
        losses.append(loss)
        if t < steps:
            w = w - lr * (2.0 / n) * (X.T @ (X @ w - y))
    return np.array(losses)

Hints

4 hints

Solution and notes

Locked

Pass every check with Submit to open the reference solution and the notes.

Run it on your computer

Clone the repository, install numpy and run the same checks in a terminal:

python3 scripts/practice/run.py lab-learning-rate-sweep my_answer.py
python3 scripts/practice/run.py lab-learning-rate-sweep my_answer.py --examples

Start my_answer.py from starter.py.

Your code is saved in this browser as you type. On a phone you can read and edit the code, but the checks are easier to use on a desktop.

Keyboard shortcuts

Ctrl Enter: Run the examples
Ctrl Shift Enter: Submit
Tab / Shift Tab: Indent or outdent the lines
Esc then Tab: Move the focus out of the editor
Alt H: Show the next hint

On a Mac, use Cmd in place of Ctrl.

Select Run to check the examples, or Submit to run every check.

Problem