Coupled Multiwavelet Neural Operator Learning for Coupled Partial Differential Equations

TL;DR

This work addresses the challenge of solving coupled PDEs by learning decoupled integral kernels in the Wavelet domain. It introduces the Coupled Multiwavelet Neural Operator (CMWNO), which uses a nonstandard multiwavelet representation and a dice training strategy to separately learn and then reconstruct coupled operators, effectively decoupling their interactions. The method achieves state-of-the-art results on Gray-Scott reaction-diffusion systems and non-local mean field games, with relative $L^2$ error improvements of approximately $2\times$ to $4\times$ over strong baselines. The approach offers a scalable, data-driven framework for complex coupled dynamics, with potential broad impact on simulations of physics, biology, and economics where coupled PDEs arise.

Abstract

Coupled partial differential equations (PDEs) are key tasks in modeling the complex dynamics of many physical processes. Recently, neural operators have shown the ability to solve PDEs by learning the integral kernel directly in Fourier/Wavelet space, so the difficulty for solving the coupled PDEs depends on dealing with the coupled mappings between the functions. Towards this end, we propose a \textit{coupled multiwavelets neural operator} (CMWNO) learning scheme by decoupling the coupled integral kernels during the multiwavelet decomposition and reconstruction procedures in the Wavelet space. The proposed model achieves significantly higher accuracy compared to previous learning-based solvers in solving the coupled PDEs including Gray-Scott (GS) equations and the non-local mean field game (MFG) problem. According to our experimental results, the proposed model exhibits a $2\times \sim 4\times$ improvement relative $L$2 error compared to the best results from the state-of-the-art models.

Figures (9)

• Figure 1: Architecture of CMWNO. Note that there are two coupled operators, $T_u$ and $T_v$, in our system, which aligns the number of coupled variables. The network $\bar{T}$ is only applied for the coarsest scale $L$ (0 in this system). The dashed arrows correspond to the auxiliary information from the unused operator without gradient during training process. For the interaction between operators, when we update the operator $T_u$, the decomposed ingredients from $T_v$ will be equipped into the reconstruction module of $T_u$ in the Wavelet domain, vice versa.
• Figure 2: Dice strategy. For each sample, one only needs to go through a specific path (round diagonal corner rectangle). Inside each path, the order of updating is from left to right, where the darker block indicates the operator we want to update and the lighter blocks provide decomposition information from the fixed operator.
• Figure 3: Learning curve - Relative $L2$ error $vs$ epochs for neural operators.
• Figure 4: Comparing the models by varying the coupling coefficient $\lambda$ at the initial condition (U-GRF, V-GRF) with resolution $s=1024$.
• Figure 5: The output of GS couple equations at the initial condition (U-Rand, V-GRF). (Left) The predicted output of the models to $u(x,\tau=1)$. (Right) The predicted output of the models to $v(x,\tau=1)$.