Multiwavelet-based Operator Learning for Differential Equations

In plain words

Many physical laws are written as partial differential equations. A neural network that learns the map from an equation's input, such as a starting state or a material property, to its solution would save many expensive simulations. Most such networks are tied to one grid size and must be retrained when the resolution changes.

This paper looks at the mathematical kernel that links inputs to outputs. For smooth problems that kernel can be written compactly in a family of bases called multiwavelets. A multiwavelet splits a signal into a coarse summary and a set of fine details at each scale, similar to a zoomable map that shows both the whole country and a single street.

The network applies this split repeatedly, learns a small mapping for each scale, and then puts the pieces back together. Because the same small pieces are reused at every scale, the authors say the model can also run on grids of a different size from the one it was trained on. They report that it beats several earlier neural operator models on the equations in their experiments.

TL;DR

• A neural operator that learns input-to-solution maps by applying a fixed multiwavelet decomposition and learned kernels at each scale §2.2 §2.3.
• The filter matrices H and G come from orthogonal polynomials (Legendre or Chebyshev) and are computed, not trained §2.2 app. B app. C.
• Author-reported: the Legendre variant has the lowest relative L2 error of the listed baselines on KdV, Burgers, Darcy and Navier-Stokes tab. 1 tab. 2 tab. 3 tab. 7.
• Before implementing, read Pitfalls: the code does not match the paper on the kernel type (A, B, C), and the printed filter equation for d contains a typo.

Problem

Learning a PDE solution operator from data needs a model that maps whole input functions to output functions. Earlier neural networks for this task are tied to one resolution, or to one PDE instance, and must be retrained when either changes §1. Neural operators (GNO, MGNO, LNO, FNO) learn the map without knowing the PDE, and they are discretisation-independent, but the paper argues that they do not exploit the structure of the operator kernel §1 §2.4.

The paper's premise: the integral kernel K of an operator with a smooth kernel, apart from finitely many singularities, is sparse in a multiwavelet basis §2.1 §2.4. Pseudo-differential operators with smooth coefficients have a kernel that is smooth away from the diagonal, and a multiwavelet basis with k orthogonal polynomials removes the first k terms of the Taylor expansion (vanishing moments) §2.4. The paper's stated condition is k > T - 1, where T + 1 is the highest derivative order §2.4.

Method

Problem setup. The target operator T maps an input a to an output u, Ta = u, and is an integral operator with kernel K §2.1 eq. 1. Data are point values a_i = a(x_i) and u_i = u(x_i) on a grid of M points §3.

Multiwavelet transform. Piecewise polynomials of degree below k on dyadic intervals form nested spaces V_n §2.2 eq. 2. Each level adds a wavelet space W_n orthogonal to V_n §2.2 eq. 3. A basis at scale n is a shifted and scaled copy of the scale-0 basis §2.2 eq. 4. The wavelets psi_i have vanishing moments §2.2 eq. 5. The paper uses orthogonal polynomials w.r.t. a measure mu. The Legendre basis uses the uniform measure, and the Chebyshev basis uses weight 1/sqrt(1 - (2x-1)^2) on [0, 1] app. C.2 app. C.3.

Decomposition and reconstruction §2.2. Scaling coefficients s and wavelet coefficients d at scale n are obtained from the next finer scale by the filters H and G eq. 6 eq. 7. The inverse uses the same filters, corrected by the Gram matrices Sigma, which are the identity for Legendre eq. 8 eq. 9 app. C.1. The 2D and 3D versions use Kronecker products of the 1D filters §2.2.

Non-standard form §2.2. The operator is written in the multiwavelet basis with three blocks: A (wavelet to wavelet), B (scaling to wavelet) and C (wavelet to scaling), so that the output is U_d = A d + B s, U_s = C d, and the coarsest scale is handled by T-bar eq. 10 eq. 11.

MWT model §2.3. The paper replaces A, B, C and T-bar by learned maps. For a 1D input, the model runs the decomposition for log2(N) levels, applies the learned kernels at each level, and then runs the reconstruction back to the finest scale §2.3. The paper states that the same A, B, C are reused at every scale, which is why the model is resolution-independent §2.3. It uses a two-layer cascade in 1D and four layers with ReLU in 2D §3.

# MWT operator, 1D  [§2.3] [§3]
# Input:  samples a (B, N), N a power of two  [code: models/models.py#L121]
# Output: predicted samples u (B, N)

def mwt_operator(a):
    # Lift each sample to c * k channels  [code: models/models.py#L193-L194]
    x = lift(a)                                   # (B, N, c, k)
    for layer in range(num_layers):               # 2 layers in 1D  [§3]
        x = multiwavelet_layer(x)                 # [§2.3]
        # ReLU between layers, not after the last  [code: models/models.py#L196-L199]
        if layer < num_layers - 1:
            x = relu(x)
    # Decode each point to one value  [code: models/models.py#L201-L205]
    return decoder(x)                             # Linear, ReLU, Linear (128)

# One multiwavelet layer  [§2.3] [eq. 6] [eq. 7] [eq. 8] [eq. 9] [eq. 11]
# Input:  x (B, N, c, k) on the finest scale
# Output: x (B, N, c, k) on the finest scale

def multiwavelet_layer(x):
    # Decomposition, one scale per step, from fine to coarse  [eq. 6] [eq. 7]
    detail_terms = []
    mixed_terms = []
    for scale in range(log2(N) - L):
        # Pair neighbouring points: even and odd indices
        # [code: models/models.py#L140-L146]
        pairs = concat(x[::2], x[1::2], axis=channels)
        d = pairs @ G_stack                       # wavelet coefficients
        x = pairs @ H_stack                       # coarser scaling coefficients
        # Learned kernels act on the coarse scaling and detail terms  [eq. 11]
        mixed_terms.append(A(d) + B(x))
        detail_terms.append(C(d))
    # Coarsest scale: a learned k x k map  [§2.3] [code: models/models.py#L130]
    x = T0(x)
    # Reconstruction, coarse to fine  [eq. 8] [eq. 9] [app. C.1]
    for scale in reversed(range(log2(N) - L)):
        # Code: models/models.py#L133-L136
        x = x + detail_terms[scale]
        x = concat(x, mixed_terms[scale], axis=channels)
        x = interleave(x @ R_even, x @ R_odd)  # R_* from H, G, Sigma
    return x

The code's kernels A, B and C in 1D are Fourier-domain maps, shown below. The paper mentions a Fourier-based convolution only for its fluctuation study §3.1 (see Paper vs code).

# Spectral kernel, 1D  [code: models/models.py#L62-L76]
# Input:  x (B, N, c, k), real
# Output: y (B, N, c, k), real

def spectral_kernel(x):
    # Real FFT along the grid, channels flattened
    # [code: models/models.py#L65-L67]
    spectrum = rfft(flatten(x), axis=grid)    # (B, c*k, N // 2 + 1)
    # Keep the lowest `modes` frequencies and mix channels per frequency
    low = min(modes, N // 2 + 1)              # [code: models/models.py#L69]
    out = zeros_like(spectrum)
    out[:, :, :low] = channel_mix(spectrum[:, :, :low], weights[:, :, :low])
    # Inverse FFT back to the grid  [code: models/models.py#L74]
    return reshape(irfft(out, n=N), (B, N, c, k))

Legendre filter matrices, computed by Gaussian quadrature app. C.2:

# Legendre filter matrices  [app. C.2] [app. B]
# Input:  k (number of polynomials)
# Output: H0, H1, G0, G1 (each k x k)

def legendre_filters(k):
    # Nodes: roots of P_k(2x - 1) on [0, 1], k of them  [app. C.2]
    nodes = roots(P_k(2 * x - 1))
    # Quadrature weights  [app. C.2]  [code: models/utils.py#L151]
    weights = 1 / (k * dP_k(2*nodes-1) * P_{k-1}(2*nodes-1))
    # Entries are integrals over the half-intervals  [app. C.1] [app. C.2]
    for i in range(k):
        for j in range(k):
            H0[i, j] = sum(weights * phi[i](nodes / 2) * phi[j](nodes))
            H1[i, j] = sum(weights * phi[i]((nodes + 1) / 2) * phi[j](nodes))
            G0[i, j] = sum(weights * psi[i](nodes / 2) * phi[j](nodes))
            G1[i, j] = sum(weights * psi[i]((nodes + 1) / 2) * phi[j](nodes))
    # Multiply by 1/sqrt(2) on all four  [code: models/utils.py#L155-L158]
    return H0 / sqrt(2), H1 / sqrt(2), G0 / sqrt(2), G1 / sqrt(2)

The phi_i are the orthonormal shifted Legendre polynomials, phi_i(x) = sqrt(2i+1) P_i(2x-1) app. C.2. The psi_i are the wavelets from Gram-Schmidt orthogonalisation app. B app. C.2. For k = 3 the paper gives H^(0) = [[1/√2, 0, 0], [-√3/(2√2), 1/(2√2), 0], [0, -√15/(4√2), 1/(4√2)]] app. C.2, which is a useful check for an implementation.

Key equations (from §2.2, with the printed typo fixed):

s_{l}^{n} = H^{(0)} s_{2l}^{n+1} + H^{(1)} s_{2l+1}^{n+1} eq. 6

d_{l}^{n} = G^{(0)} s_{2l}^{n+1} + G^{(1)} s_{2l+1}^{n+1} eq. 7

s_{2l}^{n+1} = \Sigma^{(0)} ( H^{(0)T} s_{l}^{n} + G^{(0)T} d_{l}^{n} ) eq. 8

s_{2l+1}^{n+1} = \Sigma^{(1)} ( H^{(1)T} s_{l}^{n} + G^{(1)T} d_{l}^{n} ) eq. 9

U_{d,l}^{n} = A\,d_{l}^{n} + B\,s_{l}^{n}, \qquad U_{\hat{s},l}^{n} = C\,d_{l}^{n}, \qquad U_{s,l}^{L} = \bar{T}\,s_{l}^{L} eq. 11

Hyperparameters

Values are from the paper unless marked code. Code values come from the notebooks at commit 4b0770c, not from a stated paper value. "not stated" means neither the paper nor the code gives it.

Results

Author-reported, relative L2 error on the test set. Setup: 1000 training and 200 test samples unless stated §3. Data generation for each equation is described in §3.1, §3.3 and §3.4. Burgers and Navier-Stokes baselines are taken from published results of the earlier neural-operator papers §3.3 app. D.1.

KdV equation, varying input resolution s §3.1 tab. 1

Darcy flow, varying resolution s §3.4 tab. 2

Navier-Stokes, four configurations (nu, T, N) app. D.1 tab. 3

Train at low resolution, test at high resolution, Burgers app. D.2 tab. 4

KdV with squared-exponential inputs app. D.4.1 tab. 5

Burgers equation, varying resolution s app. D.5 tab. 7 (also Fig. 6)

The paper's "nearly an order of magnitude" claim §3.1 §3.3 holds against GNO, MGNO and LNO: for example, KdV at s=64 gives 0.0429 / 0.00338 = 12.7 against LNO. Against FNO, the Legendre margin is smaller, about 1.2 to 3.7 times across the tables here (KdV at s=64: 0.0125 / 0.00338 = 3.7; Burgers at s=256: 0.00332 / 0.00199 = 1.7; Navier-Stokes at nu=1e-4, N=10000: 0.0820 / 0.0667 = 1.2 against FNO-3D). The abstract's "2X-10X" range does not name a baseline.

Paper vs code

The code is the official repository (commit 4b0770c, 2021-09-29). It has no LICENSE file, so do not copy code from it. Only the paths and lines above are cited. The repository's README and notebooks give the training setups.

1. Kernels A, B and C. The paper says A, B and C are "single-layered CNNs following a linear layer" §2.3 §3, and for NS "a single-layered 3D convolution" app. D.1. The code uses spectral (Fourier) kernels instead. In 1D, all of A, B and C are spectral models.py. The 1D convolutional class sparseKernel1d is defined but not used by MWT_CZ1d models.py. In 2D, A is spectral and B and C are 2D convolutions models.py. In 3D, the library model has spectral A and 3D convolutions for B and C models.py. The NS test scripts, which load the pretrained checkpoints, define spectral kernels for all three test_NS_MWT_N_10000.py NS.ipynb. The NS notebook also applies a ReLU inside the spectral kernel after the inverse FFT NS.ipynb. So the paper's "CNN" wording describes the 1D and NS models only partly, and the 2D and 3D library models only in B and C. Implement what the code does, and say so.
2. Printed equation 7. The paper prints d = G^(0) s_{2l} + H^(1) s_{2l+1} §2.2 eq. 7. The matrix form in the appendix uses G^(1) app. C.1, and the code uses G1 utils.py models.py. The card uses G^(1).
3. Loss reduction. The paper says the loss is relative L2 §3. The code uses LpLoss(size_average=False), which sums the per-sample relative errors over the batch kDV.ipynb Darcy.ipynb NS.ipynb utils.py. The reported train and test numbers are means over samples utils.py utils.py.
4. Quadrature. Appendix C.3 says a k-point rule suffices for Chebyshev, and C.4 says the paper uses 2k points app. C.3 app. C.4. The code uses 2k points for Chebyshev utils.py and k points for Legendre utils.py. Legendre psi are computed in closed form by exact polynomial integration, not by quadrature utils.py.
5. Architecture details the paper omits. The code adds a linear lift from the input channels to c·k models.py, a decoder of linear, ReLU and linear with hidden width 128 models.py models.py, and ReLU between multiwavelet layers only, not after the last one models.py. The NS scripts add BatchNorm after each layer test_NS_MWT_N_10000.py test_NS_MWT_N_10000.py.
6. Inputs and normalisation. The paper does not state input features or normalisation. The code uses a grid channel in Darcy and in NS, and per-point Gaussian normalisers on Darcy and NS Darcy.ipynb NS.ipynb. KdV is not normalised kDV.ipynb.
7. NS training set. The paper's N=10000 configuration app. D.1 uses 9800 training and 200 test samples in the code test_NS_MWT_N_10000.py.
8. Data and checkpoints are not in the repository. The KdV loader expects Data/KDV/kdv_train_test.mat kDV.ipynb. The repository directory is Data/kDV (case-sensitive) and contains only MATLAB generators gen_kdv_smooth.m. The Darcy loader expects a file outside the repository Darcy.ipynb. The NS checkpoints named in test_NS_MWT_N_10000.py are not included. The Burgers and NS datasets come from the FNO authors (README).
9. Darcy resolution. The Darcy notebook uses the full 512x512 grid with no subsampling Darcy.ipynb. The lower-resolution rows of tab. 2 need a subsampling step that is not in the repository.
10. Random-filter experiment. MWT Rnd tab. 2 is not implemented in the repository. The paper's claim that random filters do not learn is not checked here.

Pitfalls

• Power-of-two grids. The number of levels is floor(log2(N)), and each level halves the grid models.py models.py. A grid size that is not a power of two breaks the concatenation at the coarse levels, so the paper's "resolution independence" holds only for 2^k grids in this code. Test with a non-power-of-two input before relying on it.
• Sum-reduced loss. Gradient size grows with the batch size because the loss sums over samples utils.py. The paper does not say which reduction it uses. Keep the code's reduction if you want to match its results, and note that the reported error is still a mean.
• Sigma on the reconstruction side only. For Chebyshev, the decomposition uses the raw H and G, and the reconstruction uses H·Sigma and G·Sigma models.py models.py. Using the reconstruction matrices in the decomposition would not match eq. 6 and 7. The paper's eq. 8 and 9 put Sigma on the reconstruction side eq. 8 eq. 9.
• Chebyshev Sigma values are rounded. The appendix gives Sigma and G to four decimals app. C.3. Compute them with the quadrature instead of copying.
• Numerical limits of the filters. The paper reports that filter matrices drift from the orthogonality constraint for k > 20 in floating point, and recommends k in [1, 6] app. C.4. Entries below 1e-8 are set to zero in the code utils.py.
• Odd k and the node at 0.5. For odd k, a Legendre or Chebyshev node sits at x = 0.5, where the wavelets are discontinuous app. C.4. The code avoids the quadrature for Legendre psi, and uses 2k nodes for Chebyshev. Keep this in mind if you port the quadrature to a different library.
• PyTorch versions. The NS notebook uses torch.rfft and torch.irfft NS.ipynb NS.ipynb, which are legacy APIs. The 1D, 2D and 3D models use torch.fft models.py models.py. The README says the NS work used PyTorch 1.7, and requirements.txt is not pinned README.md requirements.txt.
• Global reference in the NS notebook. MWT.forward calls model.Lk NS.ipynb instead of self.Lk. It works only while a global model with the same Lk exists. Use self.Lk in a port.
• Duplicate batch size in the NS notebook. batch_size is set to 20, then to 10 later NS.ipynb NS.ipynb. The run uses 10.
• Numbers that disagree within the paper. The abstract quotes 0.0033 for KdV and 0.0018 for Burgers, against 0.00338 to 0.00418 and 0.00178 to 0.00199 in tab. 1 and tab. 7. §3.5 gives 0.00189 for Burgers at s=2048, while tab. 7 gives 0.00186 §3.5 tab. 7.
• Equation typo, Euler-Bernoulli. The printed operator is ∂⁴u/∂x − ω²u = f §3.2. The fourth derivative is presumably ∂⁴u/∂x⁴; the paper does not say.
• Chebyshev vs Legendre. Chebyshev is weaker than Legendre at every resolution in the tables above. Against FNO it is mixed: it beats FNO at every KdV and Darcy resolution, but not at every Burgers or Navier-Stokes setting (for example Burgers at s=256 and NS at nu=1e-5).

Lineage

• Builds on: the multiwavelet basis of "Adaptive Solution of Partial Differential Equations in Multiwavelet Bases" (J. Comput. Phys., 2002) §2.2, and the non-standard form of Beylkin, Coifman and Rokhlin ("Fast wavelet transforms and numerical algorithms I") §2.2. The paper extends the multiwavelets to non-uniform measures, which lets it use Chebyshev bases §2.2 app. B.
• Compared against: GNO ("Neural Operator: Graph Kernel Network for Partial Differential Equations") §3; MGNO ("Multipole Graph Neural Operator for Parametric Partial Differential Equations") §3; FNO ("Fourier Neural Operator for Parametric Partial Differential Equations") §3; LNO, a low-rank kernel similar to unstacked DeepONet §3; FNO-2D, FNO-3D, U-Net, TF-Net and Res-Net on Navier-Stokes tab. 3.
• Superseded by: not known from content already in this wiki.