NTK and self-adaptive weights

Sections 10.1 to 10.3 balanced whole terms. This section looks at the balance through the kernel that governs how a network's outputs change during training, and then at a scheme that balances individual points.

The neural tangent kernel

Collect all the residuals of a PINN, at collocation points and boundary points, in one vector $\mathbf{r}$, and let the weights move by gradient descent with a tiny step on $L = \tfrac12\|\mathbf{r}\|^2$. Then $\theta$ changes by $-\eta J^\top\mathbf{r}$, with $J$ the Jacobian of Section 9.4, and so the residuals change by

$$\Delta\mathbf{r} \approx J\,\Delta\theta = -\eta\,JJ^\top\mathbf{r} = -\eta\,K\,\mathbf{r}, \qquad K = JJ^\top.$$

$K$ is the neural tangent kernel of the residuals (Jacot et al., 2018), the same matrix that Levenberg–Marquardt inverted in Section 9.4. Its entry $K_{ij}$ says how much a step that reduces residual $j$ also changes residual $i$. Split it into the block for the equation's residuals, $K_{rr}$, and the block for the boundary's, $K_{bb}$: residuals whose block has large eigenvalues shrink fast, those whose block has small ones shrink slowly. Wang et al. (2022) used this to explain why PINNs fail to train: in their examples the eigenvalues of $K_{rr}$ were orders of magnitude larger than those of $K_{bb}$, so the equation was fitted long before the boundary, and they proposed weights that equalise the two blocks' traces.

The diagonal of $K$ is cheap to compute, one gradient per residual: $K_{ii} = \|\nabla_\theta r_i\|^2$. Its average over each block measures how fast that block's residuals can move. For Section 10.1's problem, with the residual divided by $(4\pi)^2$ as in Chapter 9, at the start of training and after 5000 epochs with that section's fixed weights:

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]

def mean_diagonal(residuals, params):
    """The mean of K_ii = |grad r_i|^2 over a block of residuals."""
    total = 0.0
    for r in residuals:
        grads = torch.autograd.grad(r, params, retain_graph=True, allow_unused=True)
        total += sum((g ** 2).sum().item() for g in grads if g is not None)
    return total / len(residuals)

k = 4 * math.pi
print("K_rr / K_bb, mean diagonal     at the start       after 5000 epochs (fixed weights)")
for seed in range(3):
    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))
    params = list(net.parameters())
    u = lambda x: net(2 * x - 1)
    x = torch.rand(128, 1).requires_grad_(True)
    x_b = torch.tensor([[0.0], [1.0]])
    residuals = lambda: (-d(d(u(x), x), x) - k**2 * torch.sin(k * x)) / k**2     # scaled, as in Chapter 9
    ratio = lambda: mean_diagonal(residuals().flatten(), params) / mean_diagonal(u(x_b).flatten(), params)
    start = ratio()
    opt = torch.optim.Adam(params, lr=1e-3)
    for epoch in range(5000):
        opt.zero_grad()
        ((residuals() ** 2).mean() + (u(x_b) ** 2).mean()).backward()
        opt.step()
    print(f"seed {seed}                         {start:.1e}            {ratio():.1e}")

K_rr / K_bb, mean diagonal     at the start       after 5000 epochs (fixed weights)
seed 0                         7.7e-05            5.8e-02
seed 1                         5.7e-05            3.6e-02
seed 2                         9.1e-05            1.4e-01

At the start of training the boundary block's diagonal is about ten thousand times larger than the equation's: the boundary residuals are the ones that can move fastest. The kernel describes the network as it is, a smooth function whose second derivative barely responds to its weights, not the solution it has to reach. As the network develops the oscillation its $u''$ becomes sensitive, and the ratio grows about a thousandfold, to between 0.04 and 0.14 after 5000 epochs. Dividing the residual by a constant changes these numbers by the square of that constant and nothing else (Exercise 1). Here, then, the boundary is the fast block, the opposite of the situation Wang et al. (2022) describe for their networks: which block is larger depends on the network, the scaling and the stage of training, which is a reason to measure it rather than assume it.

Weights that equalise the blocks' traces, as Wang et al. (2022) proposed, balance the rates at which the two kinds of residual shrink, whatever the units: the balance they produce is unchanged by any rescaling of the terms. That is their attraction, and on this problem their trouble. At the start they would make the boundary about ten thousand times lighter than the weight of 1 that Chapter 9 trained with, close to the light-boundary failure of Section 10.1, because the fast boundary residuals look over-served. The kernel measures how fast each part can move, not how far it has to go. Trace weights are better switched on once the network has the outline of the solution, and are in any case a useful diagnostic: weights you can compute and inspect before you trust them.

Self-adaptive weights

McClenny & Braga-Neto (2023) took the idea from terms to points. Every residual gets its own weight, $\lambda_i$, and the loss becomes

$$L(\theta, \boldsymbol{\lambda}) = \frac{1}{N}\sum_i m(\lambda_i)\,r_i(\theta)^2 + \frac{1}{N_b}\sum_j m(\mu_j)\,b_j(\theta)^2,$$

with $m$ an increasing function that keeps the weights positive (here $m(\lambda) = \lambda^2$). The network's weights $\theta$ are trained to decrease this loss and the point weights to increase it: gradient ascent in $\lambda$, descent in $\theta$. A point whose residual stays large has its weight raised fastest, so the network is pushed hardest where it is failing. It is a min-max game, the same idea as the adversarial training of other parts of machine learning, and it is a close cousin of Section 7.2's adaptive sampling: one moves points to where the residual is large, the other weights them.

The test is Section 7.2's sharp layer, $u = \tanh\bigl((x - 0.4)/0.02\bigr)$, with 64 fixed uniform points, which that section's uniform runs could not resolve in 5000 epochs. Here three runs with and without self-adaptive weights, the point weights trained by Adam at a learning rate of $10^{-2}$:

delta, x0 = 0.02, 0.4
exact = lambda x: torch.tanh((x - x0) / delta)
f = lambda x: 2 / delta**2 * torch.tanh((x - x0) / delta) / torch.cosh((x - x0) / delta) ** 2
x_test = torch.linspace(0, 1, 4001).reshape(-1, 1)

def train_layer(adaptive, 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(64, 1).requires_grad_(True)
    x_b = torch.tensor([[0.0], [1.0]])
    lam = torch.ones(64, 1, requires_grad=True)                 # one weight per collocation point...
    mu = torch.ones(2, 1, requires_grad=True)                   # ...and per boundary point
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    opt_weights = torch.optim.Adam([lam, mu], lr=1e-2, maximize=True)    # gradient ascent
    for epoch in range(epochs):
        opt.zero_grad()
        opt_weights.zero_grad()
        r = (-d(d(u(x), x), x) - f(x)) * delta**2
        b = u(x_b) - exact(x_b)
        if adaptive:
            loss = (lam**2 * r**2).mean() + (mu**2 * b**2).mean()
        else:
            loss = (r**2).mean() + (b**2).mean()
        loss.backward()
        opt.step()
        if adaptive:
            opt_weights.step()
    with torch.no_grad():
        error = ((u(x_test) - exact(x_test)).norm() / exact(x_test).norm()).item()
        weights = (lam**2).flatten() if adaptive else torch.ones(64)
        near = ((x.flatten() - x0).abs() < 0.1).float()
        share = (weights * near).sum().item() / weights.sum().item()
        boundary = (mu**2).mean().item() if adaptive else 1.0
    return error, share, weights.mean().item(), boundary

print("weights          relative L2 error (3 seeds)    share within 0.1 of the layer    mean weight: points, boundary")
for adaptive in (False, True):
    runs = [train_layer(adaptive, seed) for seed in range(3)]
    print(f"{'self-adaptive' if adaptive else 'fixed, equal':15s}  " + "  ".join(f"{r[0]:.1e}" for r in runs)
          + "       " + "  ".join(f"{r[1]:4.0%}" for r in runs) + "         "
          + "   ".join(f"{r[2]:.0f}, {r[3]:.1f}" for r in runs))

weights          relative L2 error (3 seeds)    share within 0.1 of the layer    mean weight: points, boundary
fixed, equal     4.8e-02  5.4e-01  4.9e-01        22%   23%   33%         1, 1.0   1, 1.0   1, 1.0
self-adaptive    1.5e-02  1.2e-01  2.9e-03        22%   10%   30%         143, 1.6   844, 1.7   283, 1.6

With equal weights one run of three found the layer; with self-adaptive weights all three improved, two of them to 1.5% and 0.3%. Where did the weight go? Not, on the whole, to the layer: the share of the points' weight within 0.1 of it is no larger than a uniform share, and in one run much smaller. What changed is the balance between the terms. The points' weights grew on average by factors of 140 to 840, while the two boundary weights rose only to about 1.6, so the equation came to count a hundred times or more against the boundary, as if a fixed weight had been changed. Self-adaptive weights adjust the balance between terms and between points at once; here the first effect was the one that mattered.

In the Studio

None of these schemes is a Studio setting. The diagonal of the kernel can be computed for a generated script from its residual functions, as mean_diagonal does above, and is a cheap check of a set of weights before a long run: after weighting, the blocks' mean diagonals should be within an order of magnitude or two of each other. Self-adaptive weights need the residuals point by point, which the generated loss_pde1 reduces to a mean; a patch would replace reduce_term by a weighted mean with its own trainable weights, and the optimiser by one that also takes those weights.

Exercises

1. Why does dividing the residual by $(4\pi)^2$ divide its kernel block's entries by $(4\pi)^4$?
2. In the self-adaptive loss, what stops the point weights from growing for ever?
3. Self-adaptive weights and Section 7.2's RAD both concentrate the loss where the residual is large. Name one practical difference between them.

Recap

• The neural tangent kernel $K = JJ^\top$ governs how fast each residual can shrink; the mean of its diagonal over each term measures an imbalance before training, without running it.
• Weights that equalise the kernel's blocks balance rates, not sizes, and are blind to the scales; early in training they can badly misjudge the balance.
• Self-adaptive weights give each point a weight that training raises where the network fails; on the sharp layer they improved every run, mostly by shifting the balance between the equation and the boundary.