Towards General Neural Surrogate Solvers with Specialized Neural Accelerators

TL;DR

This work introduces SNAP-DDM, a domain-decomposition framework that couples overlapping Schwarz methods with specialized neural operators to solve PDEs on arbitrarily large and complex domains. Central innovations include the Self-Modulating Fourier Neural Operator (SM-FNO) and a hybrid data-physics training loss, enabling near-unity subdomain accuracy and stable global convergence for 2D electromagnetics and steady-state fluids. The approach demonstrates strong performance on large-scale EM simulations with diverse boundary conditions and material layouts, while also revealing trade-offs between subdomain size, model capacity, and training data requirements. The results suggest SNAP-DDM can outperform end-to-end neural surrogates in scalability and flexibility, with practical implications for metamaterial design and high-contrast media problems, though challenges remain for high-index materials and fluid applications.

Abstract

Surrogate neural network-based partial differential equation (PDE) solvers have the potential to solve PDEs in an accelerated manner, but they are largely limited to systems featuring fixed domain sizes, geometric layouts, and boundary conditions. We propose Specialized Neural Accelerator-Powered Domain Decomposition Methods (SNAP-DDM), a DDM-based approach to PDE solving in which subdomain problems containing arbitrary boundary conditions and geometric parameters are accurately solved using an ensemble of specialized neural operators. We tailor SNAP-DDM to 2D electromagnetics and fluidic flow problems and show how innovations in network architecture and loss function engineering can produce specialized surrogate subdomain solvers with near unity accuracy. We utilize these solvers with standard DDM algorithms to accurately solve freeform electromagnetics and fluids problems featuring a wide range of domain sizes.

Figures (11)

• Figure 1: SNAP-DDM framework, with Electromagnetics as a demonstration. a) Global simulation domain and corresponding H-field solution for a 2D electromagnetics problem featuring arbitrary sources, global boundary conditions, and freeform grayscale dielectric structures. The global domain is subdivided into overlapping subdomains parameterized by position $(i,j)$. Three types of specialized Neural Operator models are trained to solve for three types of subdomain problems. b) Expression for the Robin type boundary condition used in the specialized Neural Operator subdomain models. $k(\mathbf{r}) = 2\pi \varepsilon(\mathbf{r})/\lambda$ is the wave vector in a medium with dielectric constant $\varepsilon$ and $\mathbf{n}$ is the outward normal direction. c) Flow chart of the iterative overlapping Schwarz method. In iteration $k$, electromagnetic fields in each subdomain are solved using the specialized Neural Operators, and the resulting fields are used to update the subdomain boundary value inputs for iteration $k+1$. The "Boundary value update" box shows how solved fields in the $(i,j)$ subdomain are used to update the boundary value fields in nearest neighbor subdomains for subsequent iterations. B.V.: boundary value. PML: perfectly matched layers.
• Figure 2: Self-Modulating Fourier Neural Operator architecture for DDM subdomain solvers. Modifications to the standard Fourier Neural Operator (FNO) are highlighted in red and include: 1) the addition of self-modulation connections that encode the input into a tensor, which is then multiplied with the linear transformation matrix R* in each Fourier layer; and 2) the addition of residual connections inside each Fourier layer. a: network input comprising a stack of images specifying the specialized subdomain layout and Robin boundary values. u: network output comprising a stack of images specifying the real and imaginary H-field maps. P: fully connected layer that increases the number of channels. Q: fully connected layer that decreases the number of channels. $\mathcal{F}$: Fourier transform. $\sigma$: leaky-ReLU activation function xu2015empirical. W: channel mixer via 1-by-1 kernel convolution. R: original linear transform on the lower Fourier modes. R*: modulated linear transform through an element-wise multiplication with the modulation tensor.
• Figure 3: Benchmarking of material boundary value subdomain solvers on unseen test data. The model inputs are a grayscale material dielectric distribution image ($\varepsilon=1$ to $\varepsilon=16$) and Robin boundary conditions, and the outputs are images of the real and imaginary H-fields. The real parts of outputted H-fields are shown. The L1 data loss is normalized to the mean absolute ground truth field value and the physics residue map is the summed expression in Equation \ref{['PDEloss']}.
• Figure 4: SNAP-DDM evaluated on different electromagnetics systems. These systems include: a) a titanium dioxide microlens, b) a thin film silicon-based metasurface, c) a volumetric grayscale metamaterial scatterer, and d) an optimized grayscale thin film stack featuring high reflectivity. The simulation domain contains grid resolution with physical dimensions of $6.25 nm$ and the wavelength $\lambda=1.05 \mu m$. Ground truth fields, model output fields, and field error are plotted on the same scale for each device. The largest simulation domain is in (a) and comprises an array of $40 \times 40$ subdomains.
• Figure 5: Time complexity comparison with SNAP-DDM and a conventional FDFD solver for two different dielectric materials.