Derivatives with respect to the inputs in PyTorch

Chapter 3 asked PyTorch for the derivative of a loss with respect to the weights, with loss.backward(). A PINN also needs the derivative of the network's output $u$ with respect to its inputs, the coordinates $x$ and $t$. It is the same reverse sweep (Section 5.2) aimed at a different target. This section shows the call, checks it against a formula, and explains the three details that decide whether it can be used in a PINN: which tensors are tracked, how a batch of points is handled, and create_graph.

The call, checked against a formula

A network with one input and one hidden layer of $H$ tanh neurons is $u(x) = \sum_j c_j \tanh(w_j x + b_j) + c_0$, so its derivative can be written down: $u'(x) = \sum_j c_j w_j\bigl(1 - \tanh^2(w_j x + b_j)\bigr)$. PyTorch should agree with it exactly.

import torch
import torch.nn as nn

torch.manual_seed(0)
net = nn.Sequential(nn.Linear(1, 8), nn.Tanh(), nn.Linear(8, 1)).double()   # float64 for a clean comparison

x = torch.linspace(-1, 1, 5, dtype=torch.float64).reshape(-1, 1).requires_grad_(True)
u = net(x)
u_x, = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)

w, b, c = net[0].weight[:, 0], net[0].bias, net[2].weight[0]
formula = (c * w * (1 - torch.tanh(x * w + b) ** 2)).sum(dim=1, keepdim=True)
print("u_x by autograd:", u_x.detach().ravel().numpy().round(6))
print("u_x by formula :", formula.detach().ravel().numpy().round(6))
print("largest difference:", f"{(u_x - formula).abs().max().item():.1e}")

u_x by autograd: [0.468531 0.619015 0.635236 0.512404 0.352809]
u_x by formula : [0.468531 0.619015 0.635236 0.512404 0.352809]
largest difference: 0.0e+00

They agree to the last bit. Three arguments made this work.

1. x.requires_grad_(True). PyTorch records operations only on tensors that ask for it. Parameters ask by default; data points do not. Without this line there is no record connecting u to x, and the call fails.
2. grad_outputs=torch.ones_like(u). torch.autograd.grad differentiates a scalar. When u holds many values, grad_outputs gives the weight of each in a weighted sum, and ones mean "the plain sum". Writing torch.autograd.grad(u.sum(), x) is the same thing.
3. create_graph=True. It records the derivative computation itself, so that u_x can be differentiated again: for $u_{xx}$, and for the gradient of a residual loss with respect to the weights. We come back to it below.

Why one call gives every point's derivative

u holds the network's output at five points, and one call returned five derivatives. Why does differentiating the sum of the outputs give each point's own derivative? Because the network processes every row independently: $u_i$ depends on $x_i$ and on no other point, so

$$\frac{\partial}{\partial x_i}\sum_k u_k = \sum_k \frac{\partial u_k}{\partial x_i} = \frac{\partial u_i}{\partial x_i}.$$

The matrix of all derivatives $\partial u_k/\partial x_i$ is diagonal, and summing collects the diagonal (Figure 5.3). We can confirm that the batched call matches one call per point, and also see the trick break when the points are not independent, here by a network whose output is measured from its batch average:

one_at_a_time = torch.cat([
    torch.autograd.grad(net(xi).sum(), xi)[0]
    for xi in (x[i:i + 1].detach().clone().requires_grad_(True) for i in range(len(x)))
])
print("one call vs one point at a time:", f"{(one_at_a_time - u_x).abs().max().item():.1e}")

centred = lambda x: net(x) - net(x).mean()          # every output now depends on every input
v = centred(x)
v_x, = torch.autograd.grad(v.sum(), x)
print("sum trick on the centred network:", v_x.ravel().numpy().round(6))

one call vs one point at a time: 0.0e+00
sum trick on the centred network: [0. 0. 0. 0. 0.]

For the centred network the outputs always sum to zero, so the trick returns zero for every point, although each output certainly changes with its own input. Anything that mixes the points of a batch, such as batch normalisation, breaks input derivatives in this way. PINN networks are plain multilayer perceptrons and never mix points, which is why the trick is safe there.

Several inputs

With inputs $(x, t)$ the points are rows of an $(N, 2)$ tensor, and one call returns an $(N, 2)$ tensor whose columns are $u_x$ and $u_t$. Slicing picks one, keeping the column shape:

torch.manual_seed(1)
net2 = nn.Sequential(nn.Linear(2, 16), nn.Tanh(), nn.Linear(16, 1))
pts = torch.rand(4, 2, requires_grad=True)          # columns: x, t
u2 = net2(pts)
g = torch.autograd.grad(u2, pts, torch.ones_like(u2), create_graph=True)[0]
u_x2, u_t2 = g[:, 0:1], g[:, 1:2]
print("gradient shape:", tuple(g.shape), "  u_x shape:", tuple(u_x2.shape), "  u_t shape:", tuple(u_t2.shape))

gradient shape: (4, 2)   u_x shape: (4, 1)   u_t shape: (4, 1)

This is exactly what the Studio's generated script does. Its helper grad_of(y, pts) wraps the call with create_graph=True, and the residual takes columns from the result, as in Read the script.

Scale inside the network, not before it

Section 4.5 recommended feeding the network inputs scaled to $[-1, 1]$. Where the scaling happens matters. On a domain $[0, 10]$, the scaled input is $z = 2x/10 - 1$, so $du/dx = (du/dz)\,(dz/dx) = 0.2\,du/dz$. If the scaling is part of the recorded computation, autograd includes the factor. If the points are scaled first and only the scaled tensor is tracked, the derivative is taken with respect to $z$ and the factor is lost:

xr = torch.linspace(0, 10, 5, dtype=torch.float64).reshape(-1, 1).requires_grad_(True)

inside = net(2 * xr / 10 - 1)                                  # scaling recorded
d_inside, = torch.autograd.grad(inside.sum(), xr)

z = (2 * xr / 10 - 1).detach().requires_grad_(True)           # scaled first, then tracked
d_outside, = torch.autograd.grad(net(z).sum(), z)

print("scaled inside the graph :", d_inside.ravel().numpy().round(5))
print("scaled outside the graph:", d_outside.ravel().numpy().round(5))
print("ratio:", (d_outside / d_inside).ravel().numpy().round(3))

scaled inside the graph : [0.09371 0.1238  0.12705 0.10248 0.07056]
scaled outside the graph: [0.46853 0.61901 0.63524 0.5124  0.35281]
ratio: [5. 5. 5. 5. 5.]

The second version is wrong by a factor of 5 everywhere, and for a second derivative it would be wrong by 25. Nothing raises an error; the PINN would simply solve a different equation. The Studio's network does the scaling inside its forward method, which is why its derivatives come out in the units of the Domain.

Exercises

1. What shape does torch.autograd.grad(u, pts, torch.ones_like(u))[0] have when pts is $(1000, 3)$ for $(x, y, t)$? How do you get $u_y$?
2. Why must the points be created with requires_grad=True while the network's parameters need nothing extra?
3. The domain is $x \in [-2, 6]$ and someone scales the points to $[-1, 1]$ before tracking them. By what factor will their $u_x$ and $u_{xx}$ be wrong?

Recap

• torch.autograd.grad(u, x, torch.ones_like(u), create_graph=True) gives $\partial u/\partial x$ at every point in one call, provided x was created with requires_grad=True.
• The sum trick works because a pointwise network's matrix of input derivatives is diagonal; anything that mixes points breaks it.
• Scale the inputs inside the recorded computation, so that derivatives come out in the original units.