Neural Spectral Methods: Self-supervised learning in the spectral domain

TL;DR

Neural Spectral Methods (NSM) address the challenge of solving parametric PDEs by learning solution mappings directly in a spectral basis and optimizing a spectral residual loss. By formulating the PDE residual in the spectral domain and applying Parseval's identity, NSM achieves exact residual computation and a fixed computational cost that does not scale with grid resolution, enabling constant-time inference. The framework introduces a spectral neural operator with activations at collocation points and a diagonalized kernel in the spectral domain, which together reduce aliasing and backpropagation complexity relative to grid-based PINN approaches. Empirically, NSM delivers 100x speedups in training and 500x in inference while surpassing grid-based PINN-based methods in accuracy across Poisson, Reaction-Diffusion, and Navier–Stokes problems, demonstrating substantial practical gains for parametric PDE solving in smooth regimes and moderate dimensions.

Abstract

We present Neural Spectral Methods, a technique to solve parametric Partial Differential Equations (PDEs), grounded in classical spectral methods. Our method uses orthogonal bases to learn PDE solutions as mappings between spectral coefficients. In contrast to current machine learning approaches which enforce PDE constraints by minimizing the numerical quadrature of the residuals in the spatiotemporal domain, we leverage Parseval's identity and introduce a new training strategy through a \textit{spectral loss}. Our spectral loss enables more efficient differentiation through the neural network, and substantially reduces training complexity. At inference time, the computational cost of our method remains constant, regardless of the spatiotemporal resolution of the domain. Our experimental results demonstrate that our method significantly outperforms previous machine learning approaches in terms of speed and accuracy by one to two orders of magnitude on multiple different problems. When compared to numerical solvers of the same accuracy, our method demonstrates a $10\times$ increase in performance speed.

Figures (7)

• Figure 1: Schematic of NSM. We refer to Neural Spectral Methods (NSM) as a general approach to learn PDE solution mappings in the spectral domain. Our method consists of two components: a) The parameters $\phi$ are converted to spectral coefficients, $\tilde{\phi}$. In each NN layer $l$, the spectral coefficients $\tilde{v}_\theta^{(l)}$ are transformed by a linear operator $\tilde{\mathcal{K}}$, with the activation $\sigma$ then applied on collocation points in the physical space. b) The prediction $\tilde{u}_\theta$ is transformed by $\tilde{\mathcal{F}_\phi}$, the spectral form of the differential operator, which gives the spectral coefficients $\tilde{R}$ of the residual function. The exact residual norm is obtained by Parseval's Identity, giving the spectral loss $||\tilde{R}||_2^2$. We contrast our method against the commonly employed grid-based neural operators with a PINN loss. c) General neural operators learn the PDE solutions as transformations of function values on $x_i$. We consider the kernel integral in a more general sense, with transformation $T$ not restricted to a Fourier basis. d) Autograd or finite difference methods are used to obtain the higher-order derivatives. The PINN loss is then obtained by approximating the norm of the residual function on the sampled points.
• Figure 2: Reaction-Diffusion equation with $\nu=0.01$. In a) and b), the $L_2$ relative error and PDE residual on the test set are plotted over training each model. The grid-based methods (FNO trained with a PINN loss) show improved accuracy as the grid resolution increases, but are significantly slower to train. When tested on different resolutions, significant aliasing errors occur on test grid resolutions that differ from the training grid resolution. In contrast, NSM has a much low error and PDE residual, and achieves this lower error 100$\times$ faster than the grid-based methods. In c), when compared with iterative numerical solvers on different resolutions, NSM achieves the same level of accuracy with a 10$\times$ speedup. Notably, both the accuracy and computational cost of NSM remains constant, regardless of grid resolution.
• Figure 3: Navier-Stokes equation with $\nu = 10^{-4}$. In a) and b), the relative error and PDE residual on the test set are plotted. NSM achieves low $L_2$ error and PDE residual $100\times$ faster than FNO + PINN methods, and is an order of magnitude more accurate. c) NSM captures fine features of the vorticity evolution accurately, while the grid-based approach fails to predict the overall shape.
• Figure 4: (a) Source term (b) Solution of periodic Poisson equation. An additional constraint $u(0, 0) = 0$ is imposed to ensure the uniqueness of the solution. (c) Relative error curves during training. NSM converges significantly faster and to a lower error than the other methods.
• Figure 5: Solutions of the 1D Reaction-Diffusion equation with differentdiffusion coefficients $\nu$. Following krishnapriyan2021characterizing, we use a reaction coefficient $\rho = 5$ and four different values for the diffusion coefficient. As $\nu$ increases, the problem becomes progressively harder to solve.

Theorems & Definitions (8)

• Example 1: Fourier basis
• Example 2: Chebyshev polynomials
• Theorem 1: Parseval's Identity
• Corollary 1
• Definition 1: Orthogonal basis
• Proposition 1
• Proposition 2
• proof