Autodiff mistakes and how to catch them

The calls of Sections 5.3 and 5.4 are short, and each argument matters. This section collects the mistakes that come up again and again when a PINN residual is written by hand. Each one is shown happening, with the message PyTorch gives, if it gives one. They fall into two groups: mistakes that stop the program, which are annoying but safe, and mistakes that let it run with the wrong gradient, which are neither. The section ends with a check that catches both kinds.

All the examples use one small network:

import math
import torch
import torch.nn as nn

torch.manual_seed(0)
net = nn.Sequential(nn.Linear(1, 16), nn.Tanh(), nn.Linear(16, 16), nn.Tanh(), nn.Linear(16, 1))
first_line = lambda err: str(err).splitlines()[0]

def gradient_of(loss):
    """The gradient of `loss` with respect to every weight, as one vector."""
    net.zero_grad()
    loss.backward()
    return torch.cat([(q.grad if q.grad is not None else torch.zeros_like(q)).flatten() for q in net.parameters()])

Mistakes that stop the program

1. The points do not track gradients. Without requires_grad=True, nothing connects the output to the input:

x = torch.rand(100, 1)
try:
    torch.autograd.grad(net(x).sum(), x)
except RuntimeError as err:
    print(first_line(err))

One of the differentiated Tensors does not require grad

2. create_graph is missing on the first derivative. The first derivative is computed, but as a plain number with no record of how it was made, so it cannot be differentiated again:

x = torch.rand(100, 1, requires_grad=True)
u_x = torch.autograd.grad(net(x).sum(), x)[0]           # create_graph forgotten
try:
    torch.autograd.grad(u_x.sum(), x)                    # asking for u_xx
except RuntimeError as err:
    print(first_line(err))

element 0 of tensors does not require grad and does not have a grad_fn

3. Evaluating the residual under torch.no_grad(). It is natural to wrap evaluation code in torch.no_grad() to save memory. A residual map, though, needs derivatives even after training:

with torch.no_grad():
    try:
        x = torch.rand(100, 1, requires_grad=True)
        torch.autograd.grad(net(x).sum(), x)
    except RuntimeError as err:
        print(first_line(err))

element 0 of tensors does not require grad and does not have a grad_fn

Compute residuals outside no_grad (or inside torch.enable_grad()), then .detach() the result for plotting.

4. Changing the points in place. Rescaling or shifting a tensor of collocation points after it has been marked for gradients is refused:

x = torch.rand(100, 1, requires_grad=True)
try:
    x *= 2.0
except RuntimeError as err:
    print(first_line(err))

a leaf Variable that requires grad is being used in an in-place operation.

Do the arithmetic first and call requires_grad_(True) last, as the Studio's sampling functions do.

Mistakes that run with the wrong gradient

These are the expensive ones, because training goes ahead, the loss falls, and the answer is quietly wrong.

5. create_graph missing on the last derivative, then "fixed" with retain_graph. Take a residual with a first and a second derivative, $u_x - 0.1\,u_{xx}$, and forget create_graph=True on the call that makes $u_{xx}$. PyTorch stops at the backward pass, and its message suggests a remedy:

x = torch.rand(100, 1, requires_grad=True)

def residual(**last_call):
    u = net(x)
    u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x.sum(), x, **last_call)[0]
    return u_x - 0.1 * u_xx

try:
    gradient_of((residual() ** 2).mean())                               # create_graph forgotten
except RuntimeError as err:
    print(first_line(err)[:72], "...")
    print("the message suggests retain_graph=True:", "retain_graph=True" in str(err))

right = gradient_of((residual(create_graph=True) ** 2).mean())        # the correct call
wrong = gradient_of((residual(retain_graph=True) ** 2).mean())        # the tempting "fix"
print(f"with retain_graph=True: no error, and the gradient is off by {(right - wrong).norm() / right.norm():.0%}")

Trying to backward through the graph a second time (or directly access s ...
the message suggests retain_graph=True: True
with retain_graph=True: no error, and the gradient is off by 10%

With retain_graph=True the error disappears, $u_{xx}$ has the right values, and training runs, but $u_{xx}$ is treated as a constant when the weights are updated: the optimiser is blind to the second-derivative term of its own residual. The fix the message does not mention is the right one, create_graph=True. The Studio's grad_of passes both arguments on every call, which is the simplest way never to meet this.

6. Breaking the record inside the residual. Any operation that leaves PyTorch, such as .detach(), .item(), .numpy() or a NumPy function applied to a tensor, cuts the record at that point. In a nonlinear residual such as Burgers', $u\,u_x$, it is easy to do by accident:

def burgers_term(cut):
    u = net(x)
    u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0]
    return (u.detach() if cut else u) * u_x

right = gradient_of((burgers_term(cut=False) ** 2).mean())
wrong = gradient_of((burgers_term(cut=True) ** 2).mean())
print(f"no error, and the gradient is off by {(right - wrong).norm() / right.norm():.0%}")

no error, and the gradient is off by 38%

Known functions of the coordinates are a different matter: a source term $f(x)$ does not depend on the weights, and computing it with NumPy, or detaching it, is harmless. The test is whether the quantity depends on the network.

7. A source term computed once from tracked points. The mirror image of mistake 6. Computing $f(x)$ once, before the training loop, from points that have requires_grad=True attaches $f$ to a record that the first backward pass frees:

x = torch.rand(100, 1, requires_grad=True)
f = torch.sin(math.pi * x)                               # still attached to x's record
optimiser = torch.optim.Adam(net.parameters(), lr=1e-3)
try:
    for step in range(2):
        optimiser.zero_grad()
        u_x = torch.autograd.grad(net(x).sum(), x, create_graph=True)[0]
        ((u_x - f) ** 2).mean().backward()
        optimiser.step()
except RuntimeError as err:
    print(f"step {step}:", first_line(err)[:72], "...")

step 1: Trying to backward through the graph a second time (or directly access s ...

The first step works and the second fails. Adding retain_graph=True to backward() would silence it and keep an ever-older record alive for no reason. The fix is f = f.detach(), or computing $f$ inside the loop.

8. Scaling the inputs outside the record. Section 5.3 showed it: the derivatives come out with respect to the scaled variable, wrong by a constant factor, with no error.

One check that catches them

Every mistake above either stops the program or changes the derivative or the gradient. A finite-difference comparison in 64-bit arithmetic, run once on a few points before a long training run, catches the second kind (Section 5.1 showed it is accurate to about $10^{-7}$ for a second derivative at $h = 10^{-4}$):

def check_derivatives(u_fn, derivatives_fn, x, h=1e-4):
    """Largest gaps between the u_x, u_xx your residual code computes and central differences, in float64."""
    x = x.detach().double()
    u_x, u_xx = derivatives_fn(x)
    with torch.no_grad():
        fd1 = (u_fn(x + h) - u_fn(x - h)) / (2 * h)
        fd2 = (u_fn(x + h) - 2 * u_fn(x) + u_fn(x - h)) / h**2
    return (u_x - fd1).abs().max().item(), (u_xx - fd2).abs().max().item()

net64 = net.double()
u_fn = lambda x: net64(2 * x - 1)                 # the network as a function of physical x in [0, 1]

def derivatives_right(x):                         # scaling inside the record
    x = x.clone().requires_grad_(True)
    u_x = torch.autograd.grad(u_fn(x).sum(), x, create_graph=True)[0]
    return u_x, torch.autograd.grad(u_x.sum(), x)[0]

def derivatives_wrong(x):                         # mistake 8: scaled first, then tracked
    z = (2 * x - 1).requires_grad_(True)
    u_z = torch.autograd.grad(net64(z).sum(), z, create_graph=True)[0]
    return u_z, torch.autograd.grad(u_z.sum(), z)[0]

points = torch.rand(50, 1)
print("right: gaps %.1e and %.1e" % check_derivatives(u_fn, derivatives_right, points))
print("wrong: gaps %.1e and %.1e" % check_derivatives(u_fn, derivatives_wrong, points))

right: gaps 5.3e-09 and 2.2e-08
wrong: gaps 2.5e-01 and 5.4e-01

The correct code agrees with finite differences to the accuracy of the differences themselves. The broken code is off by the factor its scaling dropped, which no error message would ever have told you. For the weight gradient, the counterpart is torch.autograd.gradcheck, which compares a function's backward pass with finite differences in 64-bit arithmetic; run it on the residual loss of a tiny network once, and mistakes 5 and 6 show up as a mismatch.

Exercises

1. Which of the eight mistakes raise an error, and which run silently?
2. Why is computing the source term $f(x) = -\pi^2\sin(\pi x)$ with NumPy harmless, while computing $u^2$ with NumPy is not?
3. A colleague's PINN for $u_t + u\,u_x = \nu u_{xx}$ trains, and the loss falls by three orders of magnitude, but the shock is far too smooth. Name two mistakes from this section that could explain it.

Recap

• Errors are the safe kind: missing requires_grad, missing create_graph on a first derivative, no_grad, in-place changes to the points.
• The dangerous kind runs: create_graph replaced by retain_graph, .detach() or NumPy on anything that depends on the network, scaling outside the record.
• Compare autograd's derivatives with finite differences in float64 once before training, and the dangerous kind is caught too.