Sine networks and adaptive activations

Fourier features (Section 11.2) put oscillations at the network's input. Two other ideas put them, or the means of making them, inside: a network whose every activation is a sine, and an activation whose steepness is itself trained.

Sine networks

A SIREN (Sitzmann et al., 2020) uses $\sin(\omega_0 z)$ as its activation in every hidden layer, with $\omega_0 = 30$. A sine is smooth, so all its derivatives exist, and every derivative of a sine network is again built of sines and cosines: well suited to a residual that needs $u''$. What makes it train is its initialisation, which is designed to keep the pre-activations spread evenly through the depth of the network:

• the first layer has weights uniform in $[-1/n, 1/n]$ for $n$ inputs, so that with the factor $\omega_0$ the first sines oscillate up to about 30 radians across the input range;
• every later layer has weights uniform in $\bigl[-\sqrt{6/n}/\omega_0,\ \sqrt{6/n}/\omega_0\bigr]$, which cancels the factor $\omega_0$ and keeps each layer's output distributed like the one before.

Adaptive activations

Jagtap et al. (2020) put a trainable number inside the activation, $\tanh(n\,a\,z)$, with a fixed factor $n \ge 1$ and $a$ starting at $1/n$, so that training begins with the ordinary tanh. The optimiser can then make every neuron steeper or gentler at once, which changes how quickly the network can bend, and they report faster convergence for PINNs. A single $a$ for the whole network is the simplest version; one per layer or per neuron are the others.

The comparison

Section 11.2's two-scale problem, with the same training, for four networks of three hidden layers of 32: tanh; a SIREN with the initialisation above; a SIREN with every layer initialised by the later-layer rule, including the first, and zero biases, which is how the Studio's generated script does it; and tanh with one adaptive slope, $n = 10$:

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]

exact = lambda x: torch.sin(2 * math.pi * x) + 0.1 * torch.sin(20 * math.pi * x)
f = lambda x: (2 * math.pi) ** 2 * torch.sin(2 * math.pi * x) + 0.1 * (20 * math.pi) ** 2 * torch.sin(20 * math.pi * x)
f_size = (2 * math.pi) ** 2 + 0.1 * (20 * math.pi) ** 2
x_test = torch.linspace(0, 1, 4001).reshape(-1, 1)

class Sine(nn.Module):
    def __init__(self, w0=30.0):
        super().__init__()
        self.w0 = w0
    def forward(self, z):
        return torch.sin(self.w0 * z)

class AdaptiveTanh(nn.Module):
    def __init__(self, a, n=10.0):
        super().__init__()
        self.a, self.n = a, n               # one Parameter shared by every layer; the optimiser sees it once
    def forward(self, z):
        return torch.tanh(self.n * self.a * z)

def make(kind, width=32):
    layers = [nn.Linear(1, width), nn.Linear(width, width), nn.Linear(width, width), nn.Linear(width, 1)]
    if kind == "tanh":
        acts = [nn.Tanh() for _ in range(3)]
    elif kind == "adaptive tanh":
        a = nn.Parameter(torch.tensor(0.1))
        acts = [AdaptiveTanh(a) for _ in range(3)]
    else:
        acts = [Sine() for _ in range(3)]
        for i, layer in enumerate(layers):
            n = layer.in_features
            if kind == "SIREN" and i == 0:
                nn.init.uniform_(layer.weight, -1 / n, 1 / n)     # first layer: fast sines from the start
            else:
                bound = math.sqrt(6 / n) / 30
                nn.init.uniform_(layer.weight, -bound, bound)
            if kind == "SIREN, Studio's initialisation":
                nn.init.zeros_(layer.bias)
    return nn.Sequential(layers[0], acts[0], layers[1], acts[1], layers[2], acts[2], layers[3])

def train(kind, seed, epochs=5000):
    torch.manual_seed(seed)
    net = make(kind)
    u = lambda x: net(2 * x - 1)
    x = torch.rand(400, 1).requires_grad_(True)
    x_b = torch.tensor([[0.0], [1.0]])
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    for epoch in range(epochs):
        opt.zero_grad()
        loss = (((-d(d(u(x), x), x) - f(x)) / f_size) ** 2).mean() + (u(x_b) ** 2).mean()
        loss.backward()
        opt.step()
    x_fresh = torch.rand(4000, 1, generator=torch.Generator().manual_seed(7)).requires_grad_(True)
    with torch.no_grad():
        error = ((u(x_test) - exact(x_test)).norm() / exact(x_test).norm()).item()
    held_out = ((-d(d(u(x_fresh), x_fresh), x_fresh) - f(x_fresh)) / f_size).pow(2).mean().item()
    return error, loss.item(), held_out

print("network                           relative L2 error (2 seeds)   training loss       loss at fresh points")
for kind in ("tanh", "SIREN", "SIREN, Studio's initialisation", "adaptive tanh"):
    runs = [train(kind, seed) for seed in range(2)]
    print(f"{kind:32s}  " + "  ".join(f"{r[0]:.1e}" for r in runs) + "        "
          + "  ".join(f"{r[1]:.0e}" for r in runs) + "       " + "  ".join(f"{r[2]:.0e}" for r in runs))

network                           relative L2 error (2 seeds)   training loss       loss at fresh points
tanh                              8.9e-02  3.7e-01        3e-04  4e-02       6e-04  5e-02
SIREN                             9.8e-01  1.0e+00        2e-03  6e-06       6e-02  7e-02
SIREN, Studio's initialisation    2.0e-01  6.9e-02        6e-03  7e-05       3e-03  3e-04
adaptive tanh                     1.4e-02  1.6e-02        1e-04  7e-05       5e-05  8e-05

The adaptive tanh is the clear winner here: 1.4% and 1.6%, six to twenty times better than plain tanh, with the residual at fresh points as small as at the training points. The SIREN with its published initialisation failed completely, and the last two columns show how. In the second run its training loss is the smallest in the table, $6\times10^{-6}$, and its residual at fresh points ten thousand times larger: its first layer's sines oscillate fast enough to satisfy the equation at the 400 points and do something else between them, the aliasing of Section 6.1. With the slower first layer of the Studio's initialisation, the same network did better, 7% and 20%, with its residual at fresh points within a few times its training loss. A fast start is what sine networks are for, on solutions with fine detail and enough points to pin it down; for this solution and 400 points it was too fast.

In the Studio

The Network block's Activation includes sin (SIREN), with $\omega_0 = 30$ in every hidden layer, and its Initialisation includes SIREN, which the help text rightly says is required with sine activations. As noted above, the generated initialisation applies the later-layer rule to the first layer as well and sets the biases to zero, so its first sines are about twelve times slower than in the original method for a one-dimensional input; the table shows what that does on this problem. There is no adaptive activation. GELU, SiLU and Softplus are on offer too; all are smooth, which is what the second derivatives of a residual need, unlike ReLU (Section 5.4).

Exercises

1. With the first-layer rule, $W \in [-1/n, 1/n]$ and one input, what is the largest frequency, in radians per unit of $z$, of the first layer's sines? And with the later-layer rule applied to the first layer?
2. Why would a sine network be a poor choice for a solution that is a slowly varying ramp?
3. In the adaptive tanh, the factor $n = 10$ and the starting value $a = 0.1$ cancel. What, then, does $n$ change?

Recap

• Sine networks represent oscillations and all their derivatives naturally, but only with their initialisation, which treats the first layer differently from the others.
• Adaptive activations let training tune the steepness of the network; on the two-scale problem a single trained slope did nearly as well as Fourier features of the right width, and more reliably than either sine network.
• Both make it easier to oscillate between collocation points: always check the residual at fresh points.