Fixed weights from the scales

Section 10.1 found the good boundary weights around the ratio of the natural sizes of the terms. This section turns that into a rule, shows that the rule is the nondimensionalisation of Section 4.5 in another form, and tries it on the steel rod of that section, in real units, where Section 4.5 predicted an imbalance of $2\times10^7$.

The rule

Give every term a scale $s$: the size its residual would have for an error of 100%. For a condition $u = g$, it is the size of $g$, or of the variation of $u$ it controls. For the equation, it is the size of its terms for the true solution: $(4\pi)^2$ for Section 10.1's source, $a\,\Delta T/L^2$ for a rod whose temperature rise $\Delta T$ decays over the diffusion time $L^2/a$. Then weight each term by

$$w = \frac{1}{s^2},$$

so that each term measures the error relative to its own size, and a 1% mistake costs about the same in every term. For Section 10.1's problem the rule's ratio is $s_r^2/s_b^2 = (4\pi)^4 \approx 2.5\times10^4$, at the upper end of the good range that the sweep found. The rule gets the order of magnitude right, which is what matters: the range of good weights was two orders of magnitude wide, the range of possible ones eight.

The same thing as dimensionless variables

Section 4.5 wrote the rod's problem in the variables $\hat x = x/L$, $\hat t = t\,a/L^2$ and $\hat u = (u - u_e)/\Delta T$, and every coefficient disappeared. Write the network as $u = u_e + \Delta T\,N(x, t)$, so that $N$ plays the role of $\hat u$. Then, by the chain rule,

$$u_t - a\,u_{xx} = \frac{a\,\Delta T}{L^2}\bigl(N_{\hat t} - N_{\hat x\hat x}\bigr), \qquad u - u_e = \Delta T\,N,$$

and the loss with weights $1/s^2$, with $s_r = a\Delta T/L^2$ for the equation and $\Delta T$ for the conditions, is term for term the loss of the dimensionless problem. Weighting by the scales and scaling the unknown is nondimensionalisation, done inside the loss instead of on paper.

The rod, four ways

The rod of Section 4.5: $L = 0.5$ m, $a = 1.2\times10^{-5}$ m²/s, ends at 20 °C, an initial rise of 280 K, one hour. Every version scales its inputs to $[-1, 1]$ inside the network, as the Studio does, and uses the same network, points and optimiser. They differ in what the loss sees:

• raw: temperatures in °C, the equation in K/s, all weights 1;
• weights from scales: the same, with $w = 1/s^2$ for each term;
• weights and output scale: as above, with the network's output scaled, $u = u_e + \Delta T\,N$;
• dimensionless: Section 4.5's problem in $\hat x$, $\hat t$, $\hat u$, all weights 1.

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]

L, a, u_e, dT, t_end = 0.5, 1.2e-5, 20.0, 280.0, 3600.0       # m, m^2/s, degC, K, s
tau = L**2 / a                                                  # diffusion time, s
exact = lambda x, t: u_e + dT * torch.exp(-math.pi**2 * t / tau) * torch.sin(math.pi * x / L)

def train(version, seed=0, epochs=5000):
    g = torch.Generator().manual_seed(seed)
    x, t = L * torch.rand(1000, 1, generator=g), t_end * torch.rand(1000, 1, generator=g)
    x0 = L * torch.rand(200, 1, generator=g)                    # points at t = 0
    tb = t_end * torch.rand(200, 1, generator=g)                # points on the two ends
    xb = L * (torch.rand(200, 1, generator=g) < 0.5).float()
    torch.manual_seed(seed)
    net = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(),
                        nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
    N = lambda x, t: net(torch.cat([2 * x / L - 1, 2 * t / t_end - 1], dim=1))   # inputs scaled to [-1, 1]
    if version == "dimensionless":                              # x^ = x/L, t^ = t/tau, u^ = (u - u_e)/dT
        X, T = (x / L).requires_grad_(True), (t / tau).requires_grad_(True)
        u_hat = lambda X, T: N(X * L, T * tau)                  # the same network on the same inputs
        def terms():
            U = u_hat(X, T)
            return (d(U, T) - d(d(U, X), X),
                    u_hat(x0 / L, 0 * x0) - torch.sin(math.pi * x0 / L),
                    u_hat(xb / L, tb / tau))
        weights = (1.0, 1.0, 1.0)
        temperature = lambda x, t: u_e + dT * u_hat(x / L, t / tau)
    else:
        u = (lambda x, t: u_e + dT * N(x, t)) if version == "weights and output scale" else N
        xr, tr = x.clone().requires_grad_(True), t.clone().requires_grad_(True)
        def terms():
            U = u(xr, tr)
            return (d(U, tr) - a * d(d(U, xr), xr),                              # K/s
                    u(x0, 0 * x0) - (u_e + dT * torch.sin(math.pi * x0 / L)),   # K
                    u(xb, tb) - u_e)                                              # K
        weights = (1.0, 1.0, 1.0) if version == "raw" else ((L**2 / (a * dT)) ** 2, 1 / dT**2, 1 / dT**2)
        temperature = u
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    for epoch in range(epochs):
        opt.zero_grad()
        r, ic, bc = terms()
        loss = weights[0] * (r ** 2).mean() + weights[1] * (ic ** 2).mean() + weights[2] * (bc ** 2).mean()
        loss.backward()
        opt.step()
    xs, ts = torch.meshgrid(torch.linspace(0, L, 101), torch.linspace(0, t_end, 101), indexing="ij")
    xs, ts = xs.reshape(-1, 1), ts.reshape(-1, 1)
    with torch.no_grad():
        pred = temperature(xs, ts)
        rise = exact(xs, ts) - u_e
        return ((pred - exact(xs, ts)).norm() / rise.norm()).item(), pred.max().item()

print("version                     error in the temperature rise (2 seeds)    hottest predicted point (degC)")
for version in ("raw", "weights from scales", "weights and output scale", "dimensionless"):
    runs = [train(version, seed) for seed in range(2)]
    print(f"{version:26s}  " + "  ".join(f"{e:.2e}" for e, _ in runs) + "                     "
          + "  ".join(f"{m:6.1f}" for _, m in runs))

version                     error in the temperature rise (2 seeds)    hottest predicted point (degC)
raw                         5.42e-01  5.14e-01                      130.1   146.8
weights from scales         5.37e-01  4.95e-01                      131.2   127.0
weights and output scale    1.50e-02  1.13e-02                      297.2   297.6
dimensionless               1.44e-02  1.16e-02                      294.4   297.7

Raw units fail as Section 4.5 predicted, with an error of about 50%: the hottest point the networks predict is between 127 and 147 °C, in a rod that starts at 300. Weighting the terms by their scales changes almost nothing on its own. The network still has to output numbers of order 300 from weights that produce numbers of order 1, and 5000 Adam steps of about $10^{-3}$ each do not get it there (Section 11.1 returns to this). Scale the output as well and the error falls to 1.1–1.5%, the same, seed for seed, as for the dimensionless problem: 1.50 against 1.44 per cent, and 1.13 against 1.16. The two losses are equal term for term; the last digits differ because rounding takes slightly different paths through the two computations, and training amplifies the difference.

When the scales are not known

Sometimes the solution's size, or the size of the equation's terms, is not known in advance: an inverse problem whose coefficient may be anywhere in a range, or a source with a sharp peak of unknown height. Two practical options: run briefly with a rough guess, measure the terms on the result, and set the weights from those measurements; or let the weights adjust themselves during training, which is what the next two sections do.

In the Studio

The weights go into the Loss block: PDE weight $1/s_r^2$, and $1/s^2$ for each entry of BC weights and IC weights. The Studio has no setting that scales the network's output, so for a problem in real units the equivalent of the last two versions above is to enter the dimensionless problem: a Domain from 0 to 1 in $\hat x$ and from 0 to 0.173 in $\hat t$, the equation $u_t = u_{xx}$, an initial value sin(pi*x) and zero at both ends, and to convert the answer back with $u = u_e + \Delta T\,\hat u$.

Exercises

1. For the oscillator of Section 6.2, $u'' + 2\zeta\omega u' + \omega^2u = 0$ with $u(0) = 1$, $u'(0) = 0$, what are the natural scales of the equation term and of the two initial conditions? Compare with what the scaled branch of Section 6.2 did.
2. Why did weighting the terms not rescue the raw version on its own?
3. Section 8.4's inverse problem had readings of size about 1 and an equation term of size $a\pi^2 \approx 1$. Was any weighting needed there?

Recap

• Weight each term by $1/s^2$, with $s$ its natural size, so that every term measures a relative error.
• Weights from the scales together with a scaled output are nondimensionalisation, exactly; on the rod the two gave the same answer.
• Weights cannot fix an unscaled output. Scale the unknown as well as the terms.