SAOT: An Enhanced Locality-Aware Spectral Transformer for Solving PDEs

TL;DR

This work tackles neural operator learning for PDEs, where Fourier-based methods struggle with local detail and high-frequency components. It introduces Wavelet Attention (WA) to capture locality with linear-like complexity and a Spectral Attention Operator Transformer (SAOT) that fuses WA with Fourier Attention (FA) via a gated fusion block. The encoder–processor–decoder SAOT architecture achieves state-of-the-art performance across six operator benchmarks and exhibits strong discretization-invariant generalization. By combining locality-aware wavelet representations with global spectral modeling, the approach offers a robust, scalable pathway for accurate PDE solution mappings with practical impact in simulations and surrogate modeling.

Abstract

Neural operators have shown great potential in solving a family of Partial Differential Equations (PDEs) by modeling the mappings between input and output functions. Fourier Neural Operator (FNO) implements global convolutions via parameterizing the integral operators in Fourier space. However, it often results in over-smoothing solutions and fails to capture local details and high-frequency components. To address these limitations, we investigate incorporating the spatial-frequency localization property of Wavelet transforms into the Transformer architecture. We propose a novel Wavelet Attention (WA) module with linear computational complexity to efficiently learn locality-aware features. Building upon WA, we further develop the Spectral Attention Operator Transformer (SAOT), a hybrid spectral Transformer framework that integrates WA's localized focus with the global receptive field of Fourier-based Attention (FA) through a gated fusion block. Experimental results demonstrate that WA significantly mitigates the limitations of FA and outperforms existing Wavelet-based neural operators by a large margin. By integrating the locality-aware and global spectral representations, SAOT achieves state-of-the-art performance on six operator learning benchmarks and exhibits strong discretization-invariant ability.

Figures (3)

• Figure 1: Architecture of the SAOT model. It contains three main components: an encoder, a processor with $L$ pre-norm Transformer blocks, and a decoder. The core of our Transformer block is the proposed spectral attention, which is composed of a Fourier Attention (FA) and a Wavelet Attention (WA), as well as a gated fusion block.
• Figure 2: (a) Ground truth (GT); (b) Prediction errors of Fourier Attention (FA); (c) Prediction errors of Wavelet Attention (WA); (d) The energy spectrum of GT and the predictions of FA and WA with respect to the wavenumber. WA yields better-aligned energy spectrum predictions with true distributions than FA, especially at higher wavenumbers.
• Figure 3: Model generalization performance on the Darcy dataset with different mesh resolutions.