L-BFGS

Section 9.1 showed that PINN losses are badly conditioned and that a method using curvature would not mind. Newton's method uses it exactly: $\theta \leftarrow \theta - H^{-1}\nabla L$, which on a quadratic loss lands on the minimum in one step, whatever the condition number. For a network it is impractical: the Hessian $H$ has $P^2$ entries for $P$ parameters (five million for the 2209 weights of this chapter's network, ten billion for a network a hundred times larger), and away from a minimum it need not even point downhill.

From Newton to L-BFGS

A quasi-Newton method builds an approximation to $H^{-1}$ as it goes, from what it has already computed. After a step $\mathbf{s} = \theta_{\text{new}} - \theta_{\text{old}}$ the gradient changes by $\mathbf{y} = \nabla L_{\text{new}} - \nabla L_{\text{old}}$, and for a quadratic loss $H\mathbf{s} = \mathbf{y}$ exactly. BFGS updates its approximation so that it satisfies this for every new pair. L-BFGS, the limited-memory version, keeps only the last $m$ pairs and applies the approximate inverse in about $4mP$ operations, never storing a matrix (Liu & Nocedal, 1989). Each iteration then searches along the resulting direction for a step length that lowers the loss enough and flattens the slope enough, the Wolfe conditions (Nocedal & Wright, 2006).

Both halves compare the loss, or its gradient, at different points. That is the method's one strict requirement: the loss must be the same function every time it is evaluated. The full set of collocation points is used at every evaluation (no mini-batches, no resampling), which is what Section 3.3 meant by a full-batch optimiser needing the exact gradient.

The experiment

Section 9.2's problem, 128 fixed points, three seeds, four ways: Adam alone for 5000 epochs; the standard recipe, Adam for 2000 epochs and then L-BFGS for up to 500 iterations (memory 50, strong Wolfe line search, the Studio's settings); L-BFGS alone from the start, for up to 2000 iterations; and the recipe again with one mistake, fresh random collocation points at every evaluation of the loss.

import math
import time
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 run(adam_epochs, lbfgs_iters, seed, resample=False):
    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]])
    loss_at = lambda pts: ((-d(d(u(pts), pts), pts) - k**2 * exact(pts)) / k**2).pow(2).mean() + (u(x_b) ** 2).sum()
    start = time.time()
    error = lambda: ((u(x_test) - exact(x_test)).norm() / exact(x_test).norm()).item()
    adam = torch.optim.Adam(net.parameters(), lr=1e-3)
    for epoch in range(adam_epochs):
        adam.zero_grad()
        loss_at(x).backward()
        adam.step()
    handover = error()                                    # where L-BFGS takes over (if it does)
    iterations = 0
    if lbfgs_iters:
        lbfgs = torch.optim.LBFGS(net.parameters(), lr=1.0, max_iter=lbfgs_iters, history_size=50,
                                  line_search_fn="strong_wolfe")
        def closure():                                    # L-BFGS calls this whenever it needs the loss
            lbfgs.zero_grad()
            pts = torch.rand(128, 1).requires_grad_(True) if resample else x
            loss = loss_at(pts)
            loss.backward()
            return loss
        lbfgs.step(closure)
        iterations = lbfgs.state[lbfgs._params[0]]["n_iter"]
    return handover, error(), iterations, time.time() - start

recipes = {"Adam, 5000 epochs": (5000, 0, False), "Adam 2000, then L-BFGS": (2000, 500, False),
           "L-BFGS alone": (0, 2000, False), "Adam 2000, L-BFGS, resampled": (2000, 500, True)}
print("recipe                         error at hand-over     final error (3 seeds)    L-BFGS iterations  seconds")
for name, (adam_epochs, lbfgs_iters, resample) in recipes.items():
    runs = [run(adam_epochs, lbfgs_iters, seed, resample) for seed in range(3)]
    before = "  ".join(f"{r[0]:.1e}" for r in runs) if adam_epochs and lbfgs_iters else " " * 25
    print(f"{name:29s}  {before}  " + "  ".join(f"{r[1]:.1e}" for r in runs) + "   "
          + "  ".join(f"{r[2]:4d}" for r in runs) + f"    {sum(r[3] for r in runs) / 3:5.1f}")

recipe                         error at hand-over     final error (3 seeds)    L-BFGS iterations  seconds
Adam, 5000 epochs                                         7.3e-03  2.1e-02  7.9e-03      0     0     0     11.7
Adam 2000, then L-BFGS         9.5e-03  3.7e-01  1.9e-02  2.5e-03  1.3e-02  5.3e-03    229   488   442      5.9
L-BFGS alone                                              5.4e-04  1.6e-02  1.5e-03    585   654   455      1.9
Adam 2000, L-BFGS, resampled   9.5e-03  3.7e-01  1.9e-02  1.0e-02  3.7e-01  7.1e-03      3     1     4      4.6

L-BFGS after 2000 epochs of Adam beat 5000 epochs of Adam in every run, by a factor of 1.5 to 3, in half the time. Measured from where Adam handed over, the gain is larger: the seed-1 run, which Adam had left at 37%, ended at 1.3%. L-BFGS alone, from the random start, was the most accurate in two runs of three and took two seconds; in the third it stopped at 1.6%. The last row is the mistake: with new points at every evaluation, L-BFGS stopped after one to four iterations and left the errors where Adam had left them.

The gains on this small, smooth problem are modest. On the two-dimensional problem of the Studio section below, L-BFGS reduced Adam's error about tenfold, once its loss was kept fixed.

Why it stops, and what "iterations" means

PyTorch's L-BFGS runs until it has done max_iter iterations, or until the gradient or the change in the loss falls below its tolerances (tolerance_grad = 1e-7, tolerance_change = 1e-9 by default), whichever comes first. In single precision, on a problem like this one, the tolerances are usually reached first, so a run finishes before max_iter; calling it again resumes from where it stopped and, as a rule, stops again almost at once. Each iteration may evaluate the loss more than once, because the line search tries several step lengths. And the step length lr is only the first one the line search tries; leave it at 1.

Adam first, or not at all?

On this problem L-BFGS alone was the most accurate, and the quickest. That is not a rule. From a random start on a harder problem, L-BFGS can settle in a poor region that Adam's noisier steps would have left, and Rathore et al. (2024) found Adam followed by L-BFGS better than either alone across their test problems, with a theoretical argument for why the combination helps. The two-stage recipe is the safe default: Adam to get into a good region, L-BFGS to finish. When a problem is small and smooth, trying L-BFGS from the start costs seconds.

In the Studio

Tick Second stage in the Optimiser block and choose L-BFGS: the generated script runs the first stage for Epochs, then calls L-BFGS once with max_iter set to Iterations (500), history_size=50 and the strong Wolfe line search, using the stage's Learning rate (1.0) as the first trial step. It adds one entry to the loss history for the whole second stage, not one per iteration.

One detail of the generated script matters a great deal here: its boundary and initial-condition terms draw new points every time they are evaluated, including inside L-BFGS's closure. On a line without time this is harmless, because the boundary is two fixed points. On a rectangle or a circle, and on a line with time, where the boundary points have random times and the initial points random positions, it makes the loss random, with the result the warning above describes. We generated the 2D Poisson notebook with Second stage ticked (2000 Adam epochs, then L-BFGS) and ran it three times. Each time L-BFGS stopped after 3 or 4 iterations, and the relative error stayed between $1.7$ and $2.7\times10^{-3}$. With the boundary points held fixed during the second stage, it ran all 500 iterations and the error fell to $1.6$ to $1.9\times10^{-4}$, about ten times smaller. The Studio's heat-equation starter, run the same way with three seeds, stopped after 2 to 7 iterations as generated, with errors of $0.7$ to $2.7\times10^{-3}$. With its points fixed, by either remedy below, L-BFGS ran 135 to 184 iterations in two of the three runs, which ended between $2.7$ and $4.0\times10^{-4}$; in the third it stopped after 3 iterations all the same, at $1.0\times10^{-3}$. A fixed loss is necessary for L-BFGS, not sufficient.

Until the generator draws these points once, there are two remedies. Without editing the script, set the Distribution of every boundary to Regular grid, and for a problem in time the Interior distribution too, since the initial points use it: a grid is the same at every call. The collocation points then lie on a grid as well, with the drawback Section 7.1 describes. Or patch the generated script, which keeps the points random but draws them once:

# not run
import poisson_2d as p                  # a generated script with Second stage: L-BFGS ticked

def drawn_once(sampler):
    points = []
    def sample(*args, **kwargs):        # the first draw, returned at every later call
        if not points:
            points.append(sampler(*args, **kwargs))
        return points[0]
    return sample

for name in ("sample_bc1", "sample_bc2", "sample_bc3", "sample_bc4"):   # add "sample_initial" for a problem in time
    setattr(p, name, drawn_once(getattr(p, name)))
history = p.train()

It fixes the boundary points for the Adam stage too, which costs nothing with 80 points per edge; in a test run the patched 2D script's error was $1.0\times10^{-4}$.

Exercises

1. Why can L-BFGS take a step of length 1, and Adam not?
2. A colleague redraws collocation points every 100 epochs during the Adam stage, then runs L-BFGS. Which points should L-BFGS use?
3. L-BFGS keeps $m$ pairs of vectors of length $P$. How much memory does that take for $m = 50$ and the 2209 weights here, in single precision? For a network with a million weights?

Recap

• L-BFGS approximates Newton's method from the gradients it has seen, with a line search for the step length. It removes much of the conditioning that slows Adam.
• It needs a deterministic, full-batch loss. Random points inside the loss make it stop almost at once, silently.
• Adam then L-BFGS is the safe default. Expect L-BFGS to stop on its tolerances before max_iter, and do not expect a second call to help.