Scientific ML Studio

L-BFGS treats the loss as a black box and learns its curvature from gradients. A PINN loss is not a black box: it is the squared length of a vector of residuals, one entry per collocation point and per boundary point,

$$L(\theta) = \|\mathbf{r}(\theta)\|^2, \qquad \mathbf{r} = \bigl(r_1(\theta), \dots, r_m(\theta)\bigr),$$

and for problems of that form there is a good approximation to the curvature that costs only first derivatives.

Gauss–Newton and Levenberg–Marquardt

Let $J$ be the $m \times P$ Jacobian of the residuals, $J_{ij} = \partial r_i/\partial\theta_j$. Then $\nabla L = 2J^\top\mathbf{r}$, and the Hessian is $2J^\top J$ plus a term that contains the residuals themselves, which is small near a good solution. Dropping it gives the Gauss–Newton step, the exact minimiser of the loss of the linearised residuals:

$$\boldsymbol{\delta} = -\bigl(J^\top J\bigr)^{-1}J^\top\mathbf{r}.$$

Far from the solution the linearisation is poor and $J^\top J$ may be singular, so Levenberg–Marquardt adds a multiple of the identity, $\bigl(J^\top J + \mu I\bigr)\boldsymbol{\delta} = -J^\top\mathbf{r}$, and adjusts $\mu$ as it goes: after a step that lowers the loss, $\mu$ is reduced and the method behaves more like Gauss–Newton; after a step that does not, the step is rejected, $\mu$ is raised, and the method behaves more like a short gradient step (Marquardt, 1963; Nocedal & Wright, 2006). When there are fewer residuals than parameters, the same step is cheaper to compute as $\boldsymbol{\delta} = -J^\top\bigl(JJ^\top + \mu I\bigr)^{-1}\mathbf{r}$, which needs only an $m \times m$ system. $JJ^\top$ is the neural tangent kernel of the residuals, which Section 10.4 meets again.

On Section 9.1's sine model, where the residuals are linear in the parameters, one Gauss–Newton step solves the problem exactly, whatever its condition number.

The experiment

The same problem, network, points and seeds as Sections 9.2 and 9.3, now in double precision (the reason is below). The residuals are written with torch.func, which composes derivatives with respect to the input (for $u''$) and with respect to the weights (for $J$). For comparison, L-BFGS alone runs on the same double-precision problem:

import math
import time
import torch
import torch.nn as nn
from torch.func import functional_call, grad, jacrev, vmap

torch.set_default_dtype(torch.float64)
k = 4 * math.pi
x_test = torch.linspace(0, 1, 2001).reshape(-1, 1)

def make_problem(seed, n=128):
    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))
    x = torch.rand(n)
    x_b = torch.tensor([0.0, 1.0])
    def u_at(params, xi):                                       # the network at one point, as a number
        return functional_call(net, params, (2 * xi.reshape(1, 1) - 1,)).squeeze()
    u_xx = grad(grad(u_at, argnums=1), argnums=1)               # its second derivative in x
    def residuals(params):                                      # the loss is the squared length of this
        r = (-vmap(u_xx, in_dims=(None, 0))(params, x) - k**2 * torch.sin(k * x)) / k**2 / math.sqrt(n)
        return torch.cat([r, vmap(u_at, in_dims=(None, 0))(params, x_b)])
    params = {name: p.detach().clone() for name, p in net.named_parameters()}
    return net, params, residuals

def error_of(net, params):
    u = functional_call(net, params, (2 * x_test - 1,))
    return ((u - torch.sin(k * x_test)).norm() / torch.sin(k * x_test).norm()).item()

def levenberg_marquardt(net, params, residuals, iterations=200, mu=1e-3, report=(25, 50, 100, 200)):
    names = list(params)
    sizes = [params[n].numel() for n in names]
    unflat = lambda v: {n: t.reshape(params[n].shape) for n, t in zip(names, torch.split(v, sizes))}
    theta = torch.cat([params[n].reshape(-1) for n in names])
    r = residuals(params)
    loss = (r ** 2).sum()
    errors = []
    for it in range(1, iterations + 1):
        jac = jacrev(residuals)(unflat(theta))
        J = torch.cat([jac[n].reshape(len(r), -1) for n in names], dim=1)       # m x P
        while mu < 1e10:
            step = -J.T @ torch.linalg.solve(J @ J.T + mu * torch.eye(len(r)), r)
            r_new = residuals(unflat(theta + step))
            if (r_new ** 2).sum() < loss:                       # accept, and trust the linear model more
                theta, r, loss, mu = theta + step, r_new, (r_new ** 2).sum(), mu / 3
                break
            mu *= 2                                             # reject, and step more cautiously
        if it in report:
            errors.append(error_of(net, unflat(theta)))
    return errors

def lbfgs_alone(net, params, residuals, iterations=2000):
    theta = {n: p.clone().requires_grad_(True) for n, p in params.items()}
    opt = torch.optim.LBFGS(list(theta.values()), lr=1.0, max_iter=iterations, history_size=50,
                            line_search_fn="strong_wolfe")
    def closure():
        opt.zero_grad()
        loss = (residuals(theta) ** 2).sum()
        loss.backward()
        return loss
    opt.step(closure)
    return error_of(net, {n: p.detach() for n, p in theta.items()})

print("seed   L-BFGS alone    Levenberg-Marquardt after 25, 50, 100, 200 iterations    seconds")
for seed in range(3):
    net, params, residuals = make_problem(seed)
    e_lbfgs = lbfgs_alone(net, params, residuals)
    start = time.time()
    e_lm = levenberg_marquardt(net, params, residuals)
    print(f"{seed:3d}    {e_lbfgs:.1e}         " + "   ".join(f"{e:.1e}" for e in e_lm)
          + f"                     {time.time() - start:4.0f}")

seed   L-BFGS alone    Levenberg-Marquardt after 25, 50, 100, 200 iterations    seconds
  0    1.4e-02         4.7e-04   1.7e-04   5.2e-05   3.7e-06                       18
  1    8.8e-03         3.7e-03   3.3e-03   1.7e-04   3.4e-05                       18
  2    3.6e-03         1.1e-02   3.2e-03   2.1e-04   3.6e-05                       17

After 200 iterations, about eighteen seconds per run, Levenberg–Marquardt's errors are between $4\times10^{-6}$ and $4\times10^{-5}$: a hundred to four thousand times smaller than those of L-BFGS on the same double-precision problem, and two to three orders of magnitude below anything Adam reached in Section 9.2. The first iterations are nothing special; after 25 the errors are about where L-BFGS ends. The difference is that it keeps going where the others stop, because every step solves the linearised least-squares problem instead of following the slope.

The price

Each iteration computes the whole Jacobian, $m \times P$ numbers, and solves a linear system of size $\min(m, P)$. Here $m = 130$ and $P = 2209$, and an iteration takes a fraction of a second. A two-dimensional problem with $10^4$ collocation points and a network of $10^5$ weights has a Jacobian of $10^9$ entries, 8 GB in double precision, and a system of size $10^4$ to solve at every step: out of reach without further ideas. Those ideas are the current research on PINN optimisers:

Matrix-free solves. The system can be solved approximately by conjugate gradients, which needs only products with $J$ and $J^\top$, each the cost of a forward or backward pass. Rathore et al. (2024) combine this with a Nyström approximation of the curvature as a preconditioner (NysNewton-CG) and report large gains over Adam and L-BFGS.
Natural gradients. Müller & Zeinhofer (2023) precondition the gradient with the Gram matrix of the equation's own inner product, the energy natural gradient, which is closely related to Gauss–Newton for linear equations, and report errors several orders of magnitude below those of Adam and L-BFGS given more time.

Double precision matters for these methods. The matrix $JJ^\top$ squares the conditioning of the Jacobian, as Section 9.1 showed for the loss, and single precision keeps only about seven significant digits. With torch.float32 in the set_default_dtype line of the code above, the same three seeds ended at errors of $4\times10^{-5}$ to $1.5\times10^{-4}$, up to forty times worse.

Which optimiser

For a first look, Adam with a decaying learning rate. To finish, L-BFGS on a fixed set of points. When the problem is small (a few hundred residuals, a few thousand weights) and the accuracy matters, a Levenberg–Marquardt or Gauss–Newton solve in double precision can go further than either, at a cost that grows quickly with the number of points and weights.

In the Studio

The Studio offers Adam, AdamW, SGD with momentum and RMSprop, and L-BFGS as a second stage. Gauss–Newton methods need the residuals as a vector, not just the summed loss, so they are not a setting but a few dozen lines that take the generated script's residual functions (residual_pde1, and the boundary terms) and stack them, as residuals does above.

Exercises

Show that $\bigl(J^\top J + \mu I\bigr)^{-1}J^\top = J^\top\bigl(JJ^\top + \mu I\bigr)^{-1}$. Which side is cheaper here?
Why would mini-batches, which suit Adam, make no sense for Levenberg–Marquardt?
Levenberg–Marquardt drove the loss far lower than L-BFGS on 128 points. What check from Section 6.5 must you apply before believing that the error between the points is as small?

Answers

Multiply both sides on the left by $J^\top J + \mu I$ and on the right by $JJ^\top + \mu I$: both sides become $J^\top JJ^\top + \mu J^\top$. The right-hand form needs a $130 \times 130$ solve instead of a $2209 \times 2209$ one.
Each step solves for the minimiser of a model of the whole loss, built from the residuals it was given. A different mini-batch at every step would give a different loss each time, the same problem as for L-BFGS, with a more expensive step thrown away at every iteration.
The residual at fresh points, many more than were trained on. With 130 residuals and 2209 weights the network could in principle satisfy the equation at the points and nowhere else; the small error on the dense test grid above says that here it did not, and the held-out residual says it without a reference.

Recap

A PINN loss is a sum of squared residuals. Gauss–Newton uses $J^\top J$ as its curvature, and Levenberg–Marquardt makes it safe far from the solution.
On a small problem in double precision, it reached errors far below those of L-BFGS within a couple of hundred iterations.
Its cost grows with the number of residuals times the number of weights; matrix-free and natural-gradient methods are how current research scales the idea up.

Beyond L-BFGS