Depth, width and modified MLPs

Section 2.4's universal approximation theorem says that one hidden layer, wide enough, can represent any continuous function. It says nothing about how easily gradient descent finds the right weights, and for a PINN that is the whole question. This section holds the number of weights roughly fixed and changes their arrangement.

The modified MLP

Wang et al. (2021), alongside learning-rate annealing (Section 10.3), proposed an architecture built for PINNs. Two extra layers encode the input once, $U = \sigma(W_U z + b_U)$ and $V = \sigma(W_V z + b_V)$, and every hidden layer mixes them with a gate computed from the layer before:

$$H^{(1)} = \sigma(W_1 z + b_1), \qquad Z^{(k)} = \sigma\bigl(W_{k+1} H^{(k)} + b_{k+1}\bigr), \qquad H^{(k+1)} = \bigl(1 - Z^{(k)}\bigr)\odot U + Z^{(k)}\odot V,$$

with the output a linear map of the last $H$. Every layer sees the input through $U$ and $V$ directly, the way a residual connection lets a deep network pass information past its layers, and the gates choose, neuron by neuron, how much of each to use.

The comparison

Section 9.1's problem, $-u'' = (4\pi)^2\sin(4\pi x)$ with the residual divided by $(4\pi)^2$, 128 fixed points, Adam at $10^{-3}$ for 5000 epochs, three seeds; six networks of about 2000 weights each:

import math
import time
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]

k = 4 * math.pi
exact = lambda x: torch.sin(k * x)
x_test = torch.linspace(0, 1, 2001).reshape(-1, 1)

def mlp(width, depth):
    layers = [nn.Linear(1, width), nn.Tanh()]
    for _ in range(depth - 1):
        layers += [nn.Linear(width, width), nn.Tanh()]
    return nn.Sequential(*layers, nn.Linear(width, 1))

class ModifiedMLP(nn.Module):
    def __init__(self, width, depth):
        super().__init__()
        self.U, self.V, self.first = nn.Linear(1, width), nn.Linear(1, width), nn.Linear(1, width)
        self.gates = nn.ModuleList([nn.Linear(width, width) for _ in range(depth)])
        self.out = nn.Linear(width, 1)
    def forward(self, z):
        U, V = torch.tanh(self.U(z)), torch.tanh(self.V(z))
        H = torch.tanh(self.first(z))
        for gate in self.gates:
            Z = torch.tanh(gate(H))
            H = (1 - Z) * U + Z * V
        return self.out(H)

def train(make, seed, epochs=5000):
    torch.manual_seed(seed)
    net = make()
    u = lambda x: net(2 * x - 1)
    x = torch.rand(128, 1).requires_grad_(True)
    x_b = torch.tensor([[0.0], [1.0]])
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    start = time.time()
    for epoch in range(epochs):
        opt.zero_grad()
        loss = ((-d(d(u(x), x), x) - k**2 * exact(x)) / k**2).pow(2).mean() + (u(x_b) ** 2).mean()
        loss.backward()
        opt.step()
    with torch.no_grad():
        error = ((u(x_test) - exact(x_test)).norm() / exact(x_test).norm()).item()
    return error, sum(p.numel() for p in net.parameters()), time.time() - start

networks = {"1 layer of 640": lambda: mlp(640, 1), "2 layers of 45": lambda: mlp(45, 2),
            "3 layers of 32": lambda: mlp(32, 3), "6 layers of 20": lambda: mlp(20, 6),
            "12 layers of 14": lambda: mlp(14, 12), "modified MLP, 3 x 26": lambda: ModifiedMLP(26, 3)}
print("network                  weights    relative L2 error (3 seeds)    seconds per run")
for name, make in networks.items():
    runs = [train(make, seed) for seed in range(3)]
    print(f"{name:22s}  {runs[0][1]:6d}     " + "  ".join(f"{r[0]:.1e}" for r in runs)
          + f"         {sum(r[2] for r in runs) / 3:5.1f}")

network                  weights    relative L2 error (3 seeds)    seconds per run
1 layer of 640            1921     2.1e-02  1.1e-01  4.8e-02          10.9
2 layers of 45            2206     1.2e-02  1.5e-02  6.6e-03           9.9
3 layers of 32            2209     7.2e-03  1.1e-02  3.8e-03          11.6
6 layers of 20            2161     2.3e-02  3.7e-02  7.1e-02          17.1
12 layers of 14           2353     5.7e-02  4.3e-02  4.3e+00          28.5
modified MLP, 3 x 26      2289     1.3e-03  2.2e-03  1.6e-03          24.1

Among plain networks the middle of the table is best: two or three hidden layers, with errors between 0.4% and 1.5%. One very wide layer is several times worse, and depth beyond three layers hurts steadily: six layers are worse again, and twelve narrow layers are worst of all, with one run failing outright and every epoch costing two and a half times as much as with three. The modified MLP, with the same number of weights, is the best network in the table by a clear margin: 0.13% to 0.22%, three to five times better than the best plain network and the most consistent from seed to seed. Its gates and its two input encodings make each epoch about twice as expensive as for a plain network of three layers.

Choosing a size

• Start small and moderately deep. Three or four hidden layers of 32 to 64 neurons, a few thousand weights, suit most problems in one to three dimensions. The Studio's default, four layers of 40, is in this range.
• Widen before you deepen. Extra depth multiplies the length of every derivative chain, forward for the residual's $u''$ and backward for the gradient (Section 5.4), so each epoch costs more, and with tanh the signal reaching the first layers shrinks with every layer it passes.
• Bigger is not automatically better. A larger network can represent more, but the optimisation, not the capacity, is usually what limits a PINN (Chapter 9). Change the optimiser, the scaling or the features before the size.
• For hard problems, change the architecture's shape. Gated networks like the modified MLP, and Fourier features (Section 11.2), are the changes most often reported to help; Wang et al. (2023) collect these and other practices, with ablations, in a practical guide.

In the Studio

The Network block's Hidden layers takes the widths as a list, 40, 40, 40, 40 by default, and Plot network diagram draws the result. There is no modified MLP or residual connection; a generated script's PINN class can be replaced by one, since the rest of the script only calls net(pts).

Exercises

1. Show that an MLP with one input, one output and $L$ hidden layers of width $n$ has $(L - 1)n^2 + (L + 2)n + 1$ weights and biases. Check the table's counts.
2. The derivative of tanh is at most 1. Why does that make the gradient reaching the first layer of a deep tanh network smaller, and why does the modified MLP suffer less from it?
3. Doubling the width of every layer of a 3-layer network roughly multiplies its number of weights by what? And the cost of an epoch?

Recap

• At a fixed number of weights, the arrangement matters: shape changes how easily the network trains, not only what it can represent.
• Moderate depth and width is the safe start; widen before deepening.
• Gated architectures such as the modified MLP give every layer direct access to the input and are worth trying when a standard network stalls.