Scientific ML Studio

Adam (Section 3.5) is the default first optimiser for PINNs, and the Studio's. It has four settings, of which one matters most: the learning rate, which is roughly the size of every parameter's step (Section 3.5). The two decay rates $\beta_1 = 0.9$ and $\beta_2 = 0.999$ and the small constant $\varepsilon$ are almost never worth changing. What remains is how the learning rate should change during training, and which weights to keep at the end.

Schedules

A schedule changes the learning rate as training goes on. Large early steps cover ground; small late steps settle into a minimum instead of rattling around it, the effect Section 3.3 met for stochastic gradients and Section 6.2's spikes showed for Adam. PyTorch's common choices, all available in the Studio:

Schedule	Learning rate at epoch $n$ of $N$
constant	$\eta_0$
step decay	$\eta_0\,\gamma^{\lfloor n/s\rfloor}$: multiplied by $\gamma$ every $s$ epochs
exponential	$\eta_0\,\gamma^n$
cosine annealing	$\tfrac12\eta_0\bigl(1 + \cos(\pi n/N)\bigr)$: a half cosine from $\eta_0$ to zero
reduce on plateau	multiplied by $\gamma$ whenever the loss has not reached a new low for a set number of epochs

An experiment

The problem is Section 9.1's hard case, $-u'' = (4\pi)^2\sin(4\pi x)$ with $u(0) = u(1) = 0$, on 128 fixed points, for 5000 epochs and three seeds. For each run we record the error with the weights of the last epoch, the error with the weights that had the lowest loss, and how much the loss still jumps about at the end (the largest loss of the last 200 epochs divided by the smallest):

import copy
import math
import torch
import torch.nn as nn

def d(y, x):
    return torch.autograd.grad(y, x, torch.ones_like(y), create_graph=True)[0]

k = 4 * math.pi
exact = lambda x: torch.sin(k * x)
x_test = torch.linspace(0, 1, 2001).reshape(-1, 1)

def scheduler(name, opt, epochs):
    if name == "step":                                    # halve the rate every 500 epochs
        return torch.optim.lr_scheduler.StepLR(opt, step_size=500, gamma=0.5)
    if name == "cosine":                                  # half a cosine, down to zero
        return torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    if name == "exponential":                             # down by a factor of 1000, smoothly
        return torch.optim.lr_scheduler.ExponentialLR(opt, gamma=1e-3 ** (1 / epochs))
    if name == "plateau":                                 # halve it after 200 epochs with no new low
        return torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=200)
    return None

def train(lr, schedule, seed, epochs=5000):
    torch.manual_seed(seed)
    net = nn.Sequential(nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(),
                        nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
    u = lambda x: net(2 * x - 1)
    x = torch.rand(128, 1).requires_grad_(True)
    x_b = torch.tensor([[0.0], [1.0]])
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    sched = scheduler(schedule, opt, epochs)
    losses, best_loss, best_weights = [], float("inf"), None
    for epoch in range(epochs):
        opt.zero_grad()
        loss = ((-d(d(u(x), x), x) - k**2 * exact(x)) / k**2).pow(2).mean() + (u(x_b) ** 2).sum()
        if loss.item() < best_loss:                      # the weights that produced the lowest loss
            best_loss, best_weights = loss.item(), copy.deepcopy(net.state_dict())
        loss.backward()
        opt.step()
        if schedule == "plateau":
            sched.step(loss.item())
        elif sched is not None:
            sched.step()
        losses.append(loss.item())
    error = lambda: ((u(x_test) - exact(x_test)).norm() / exact(x_test).norm()).item()
    with torch.no_grad():
        last = error()
        net.load_state_dict(best_weights)
        best = error()
    return last, best, max(losses[-200:]) / min(losses[-200:])

configs = [(1e-2, "constant"), (1e-3, "constant"), (1e-4, "constant"), (1e-2, "step"),
           (1e-2, "exponential"), (1e-2, "cosine"), (1e-2, "plateau")]
print("start rate, schedule      error, last weights        error, best weights        jitter at the end")
for lr, schedule in configs:
    runs = [train(lr, schedule, seed) for seed in range(3)]
    print(f"{lr:.0e}  {schedule:12s}   " + "  ".join(f"{r[0]:.1e}" for r in runs) + "    "
          + "  ".join(f"{r[1]:.1e}" for r in runs) + "    " + "  ".join(f"{r[2]:6.0f}" for r in runs))

start rate, schedule      error, last weights        error, best weights        jitter at the end
1e-02  constant       1.4e-02  4.3e-02  1.5e-02    5.7e-03  2.8e-02  9.9e-03        37      39     268
1e-03  constant       7.3e-03  2.1e-02  7.9e-03    5.6e-03  2.8e-02  7.9e-03        11      26      50
1e-04  constant       1.8e-02  1.7e-01  5.7e-02    1.9e-02  1.7e-01  5.7e-02         1       1       1
1e-02  step           1.3e-02  2.6e-02  8.4e-02    1.3e-02  2.6e-02  8.4e-02         1       1       1
1e-02  exponential    1.1e-02  7.9e-02  7.2e-02    1.1e-02  7.9e-02  7.2e-02         1       1       1
1e-02  cosine         4.6e-03  1.5e-02  1.6e-02    4.6e-03  1.5e-02  1.6e-02         1       1       1
1e-02  plateau        4.4e-03  1.6e-02  1.4e-02    5.2e-03  1.5e-02  1.5e-02         1       1    2225

Four things stand out.

A constant rate never settles. At $10^{-2}$ the loss still jumps by factors of 40 to 270 within the last 200 epochs, at $10^{-3}$ by 10 to 50. Keeping the weights with the lowest loss helped by up to a factor of 2.5, but not always: for seed 1 at $10^{-3}$ the lowest loss came with a larger error than the last epoch's. The loss is measured at 128 points, the error everywhere.
Too small a rate is too slow. At $10^{-4}$ two of the three runs end at 6% and 17%.
A schedule calms the end; it did not make the answer more accurate. Started at $10^{-2}$, every schedule ends with the loss at rest and the last weights the best. Cosine annealing and reduce-on-plateau end between $4\times10^{-3}$ and $1.6\times10^{-2}$, the same range as a constant $10^{-3}$. Step and exponential decay end worse, up to 8%: by epoch 3000 both have cut the rate by a factor of about 60, probably before the slower runs had found the oscillation, which Section 9.1 showed can take a couple of thousand epochs.
Nothing went below about $4\times10^{-3}$. For every setting the seeds differ by a factor of two to five, more than the settings differ from each other. The floor is the conditioning of Section 9.1, which no learning rate touches; the next section breaks through it.

Settings that work

Start Adam at $10^{-3}$, or at $10^{-2}$ with a schedule that brings it down. Prefer cosine annealing when the number of epochs is known in advance, and reduce-on-plateau when it is not. With a constant rate, keep the best weights rather than the last. And do not expect any of this to buy more than a small factor in accuracy: if the error must fall by orders of magnitude, change the optimiser (Section 9.3).

Early stopping and gradient clipping

Two more switches deserve a sentence each. Early stopping ends training when the loss has not improved by a set amount for a set number of epochs. With a constant learning rate the loss keeps setting new lows between its jumps, so early stopping rarely fires; with a schedule it fires when progress has genuinely stopped, which is when it is useful. Gradient clipping caps the size of the gradient before each step. Adam already limits its steps to about the learning rate (Section 3.5), so clipping changes little for Adam; it matters for plain gradient descent and momentum, whose steps grow with the gradient.

In the Studio

The Optimiser block's Learning rate and Scheduler are the settings above: None, Step decay (with Step size and Decay factor), Cosine annealing and Reduce on plateau (with Decay factor and Patience). Gradient clip and Weight decay are there too; leave weight decay at 0 for PINNs, since pulling the weights towards zero pulls the solution away from the equation. The Train block's Early stopping is the switch described above. The generated script keeps the weights of the last epoch, so with a constant learning rate choose a schedule, or read the training log before trusting the final weights.

Exercises

Adam's steps are about the learning rate in size. With $\eta = 10^{-4}$ and 5000 epochs, roughly how far can a single weight travel? Is that enough for the network of this section?
Cosine annealing needs the total number of epochs in advance. What happens to a run that was planned for 5000 epochs and is then continued for another 5000 with the same scheduler?
In which situation would reduce-on-plateau never lower the learning rate at all?

Answers

At most about $5000 \times 10^{-4} = 0.5$, and only if every step pushes the same way. Freshly initialised weights are of order 0.2 to 1, and reaching $\sin(4\pi x)$ needs many of them to change by about that much, so $10^{-4}$ leaves the network close to where it started.
PyTorch's cosine schedule is periodic: after reaching zero at epoch $N$ it climbs back towards the starting rate over the next $N$ epochs, so the continued run starts with large steps again and loses its smooth ending. Restart it with a new schedule instead.
When the loss keeps reaching new lows, however noisily, at least once in every patience epochs: a run that jitters but still drifts downwards keeps a large rate. The plateau run above that was still jumping by a factor of 2000 at the end had been lowered only twice in 5000 epochs, to $2.5\times10^{-3}$.

Recap

The learning rate is the one Adam setting that matters; a schedule that lowers it makes the end of training calm and the last weights the best ones.
With a constant rate, the loss keeps jumping and the last epoch can be several times worse than the best: keep the best weights.
No setting of Adam removed the error floor of this problem. That takes curvature, which is the subject of the next section.