Scientific ML Studio

Section 4.5 showed what raw inputs do to a fresh layer: in metres and seconds almost every neuron is saturated before training starts. Section 10.2 found the same trouble at the other end: a network asked to output temperatures of 300 had not got there after 5000 epochs. This section measures both effects on one problem and shows how to remove them without touching the equation.

Scaling inside the network

The fix for both ends is an affine map built into the trial function:

$$u_\theta(x) = U\,N_\theta\!\left(\frac{2(x - x_\text{min})}{x_\text{max} - x_\text{min}} - 1\right),$$

where $N_\theta$ is the network, which now sees inputs in $[-1, 1]$, and $U$ is the size of the solution. Because the map is part of $u_\theta$, automatic differentiation includes its factors in every derivative (Section 5.3): $u_\theta''$ comes out in the units of $u$ per unit of $x$ squared, and the residual is written exactly as before.

An experiment at each end

Inputs. The problem $-u'' = (\pi/100)^2\sin(\pi x/100)$ on $0 < x < 100$, with $u = 0$ at both ends: the solution $\sin(\pi x/100)$ has the gentlest possible shape, on a domain a hundred units long. The network sees either $x$ itself or $x$ mapped to $[-1, 1]$.

Outputs. The problem $-u'' = A\pi^2\sin(\pi x)$ on $0 < x < 1$, with solution $A\sin(\pi x)$, for $A = 1000$ and $A = 0.001$. The network's output is used either as it is or multiplied by $A$.

In every case the residual and the boundary term are divided by their natural sizes (Section 10.2), so the loss is balanced and only the scaling differs:

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 train(length, amplitude, scale_inputs, scale_output, seed, epochs=2000):
    k = math.pi / length                                       # u = amplitude * sin(k x) on [0, length]
    exact = lambda x: amplitude * torch.sin(k * x)
    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))
    def u(x):
        z = 2 * x / length - 1 if scale_inputs else x
        return amplitude * net(z) if scale_output else net(z)
    x = (length * torch.rand(64, 1)).requires_grad_(True)
    x_b = torch.tensor([[0.0], [length]])
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    for epoch in range(epochs):
        opt.zero_grad()
        residual = (-d(d(u(x), x), x) - amplitude * k**2 * torch.sin(k * x)) / (amplitude * k**2)
        loss = (residual ** 2).mean() + ((u(x_b) / amplitude) ** 2).mean()
        loss.backward()
        opt.step()
    x_test = torch.linspace(0, length, 1001).reshape(-1, 1)
    with torch.no_grad():
        return ((u(x_test) - exact(x_test)).norm() / exact(x_test).norm()).item()

cases = [("x in [0, 100], raw input",        100, 1.0,   False, False),
         ("x in [0, 100], scaled input",     100, 1.0,   True,  False),
         ("u of size 1000, raw output",      1,   1e3,   True,  False),
         ("u of size 1000, scaled output",   1,   1e3,   True,  True),
         ("u of size 0.001, raw output",     1,   1e-3,  True,  False),
         ("u of size 0.001, scaled output",  1,   1e-3,  True,  True)]
print("case                              relative L2 error after 2000 epochs (3 seeds)")
for name, length, amplitude, scale_inputs, scale_output in cases:
    errors = [train(length, amplitude, scale_inputs, scale_output, seed) for seed in range(3)]
    print(f"{name:32s}  " + "  ".join(f"{e:.1e}" for e in errors))

case                              relative L2 error after 2000 epochs (3 seeds)
x in [0, 100], raw input          1.4e+00  1.7e+00  1.0e+00
x in [0, 100], scaled input       5.6e-03  2.5e-04  1.3e-04
u of size 1000, raw output        9.7e-01  9.8e-01  9.9e-01
u of size 1000, scaled output     1.8e-04  2.5e-04  2.3e-04
u of size 0.001, raw output       2.4e-02  2.6e-02  1.1e-01
u of size 0.001, scaled output    1.9e-04  6.4e-04  2.3e-04

Every scaled version reaches errors of a few times $10^{-4}$ in 2000 epochs (one input run, $5.6\times10^{-3}$). Raw inputs on $[0, 100]$ fail outright, with errors above 100%: the saturated network of Section 4.5 cannot even represent the gentle curve it is asked for. A solution of size 1000 without output scaling is 97–99% wrong: after 2000 steps of about $10^{-3}$ each, the network's output is still only a few per cent of the size it must reach. A solution of size 0.001 is less dramatic, since shrinking an output is easier than growing it, but it is still about a hundred times less accurate than with the output scaled. None of this has anything to do with the equation; it is all in what the network is asked to produce from what it is given.

Initialisation

The scale of the starting weights matters for the same reason: it decides whether the signals reaching the deep layers are of order one. The rules of Section 2.2 and 2.5 keep the variance of each layer's output close to that of its input: Xavier (Glorot) initialisation for tanh and sine-like activations (Glorot & Bengio, 2010), Kaiming (He) initialisation for ReLU-like ones (He et al., 2015). With inputs and outputs of order one, a standard initialisation and a network a few layers deep need no further thought. The exception is the sine network of Section 11.3, which needs an initialisation of its own.

In the Studio

The generated network scales its inputs to $[-1, 1]$ from the Domain block's bounds, always when no Encoding block is connected, and when Normalise inputs to [-1, 1] is ticked when one is. The Network block's Initialisation offers Xavier / Glorot (the default, for tanh), Kaiming / He, SIREN and PyTorch default. There is no setting for the output scale: for a solution far from size one, either scale the problem itself (Section 4.5) or patch the trial function after loading the script, as Section 6.4 did for hard constraints:

# not run
import poisson_1d as p                      # any generated script

U = 1000.0                                  # the size of the solution
network_field = p.field
p.field = lambda pts, component=0: U * network_field(pts, component)
history = p.train()

Exercises

With raw inputs on $[0, 100]$ and first-layer weights of order 1 (PyTorch's default for a layer with one input), roughly how large are the first layer's pre-activations at $x = 50$? What does tanh do with them?
The output scale $U$ multiplies $u_\theta$, so it multiplies every derivative too. Why does that not upset the residual?
A solution has size about 1 in most of the domain but a peak of 50 in a small region. What scale would you choose for the output, and what else might help?

Answers

With one input the first layer's weights are of order 1, so the pre-activations are of order $50$: tanh is flat there to many digits, every neuron outputs $\pm1$, and the network starts as almost a constant with almost no gradient.
The residual is linear in $u$ here, so scaling $u$ by $U$ scales the residual by $U$, and dividing the residual by its natural size, which also contains $U$, removes it again. What changes is only what the network itself has to produce: numbers of order one.
Scale by the typical size, 1, so that most of the solution is of order one, and accept that the network's output must reach 50 in the peak; or write the solution as a known peak shape plus a correction, the idea of the enriched trial function in Section 7.3; and put collocation points in the peak (Section 7.2).

Recap

Map the inputs to $[-1, 1]$ and multiply the output by the solution's size, inside the trial function, so that automatic differentiation keeps every derivative in the problem's units.
Unscaled inputs saturate the network; unscaled outputs leave it a long way to travel. Both cost far more than any choice of optimiser can recover.
A standard initialisation suffices once inputs and outputs are of order one.

Input and output scaling