Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery

TL;DR

The paper addresses the challenge of learning forward PDE solutions and inverse PDE parameters in ill-posed settings while preserving interpretability. It introduces Neural Interpretable PDEs (NIPS), which combine linear attention with Fourier-domain spectral convolution to form a data-dependent kernel map that jointly solves forward and inverse PDE tasks across multiple systems, enabling zero-shot generalization. Key contributions include the data-dependent kernel map $K[\mathbf{u}_{1:d},\mathbf{f}_{1:d}]$, a linear-attention reformulation to reduce computation, and Fourier kernel learning that facilitates scalable, interpretable physics discovery with strong empirical gains over NAO and AFNO baselines in Darcy Flow, MMNIST, and synthetic tissue tasks. The approach enables scalable physics discovery with interpretable kernels and zero-shot generalization, and the authors provide open-source code to support reproducibility and broader adoption across physics domains.

Abstract

Attention mechanisms have emerged as transformative tools in core AI domains such as natural language processing and computer vision. Yet, their largely untapped potential for modeling intricate physical systems presents a compelling frontier. Learning such systems often entails discovering operators that map between functional spaces using limited instances of function pairs -- a task commonly framed as a severely ill-posed inverse PDE problem. In this work, we introduce Neural Interpretable PDEs (NIPS), a novel neural operator architecture that builds upon and enhances Nonlocal Attention Operators (NAO) in both predictive accuracy and computational efficiency. NIPS employs a linear attention mechanism to enable scalable learning and integrates a learnable kernel network that acts as a channel-independent convolution in Fourier space. As a consequence, NIPS eliminates the need to explicitly compute and store large pairwise interactions, effectively amortizing the cost of handling spatial interactions into the Fourier transform. Empirical evaluations demonstrate that NIPS consistently surpasses NAO and other baselines across diverse benchmarks, heralding a substantial leap in scalable, interpretable, and efficient physics learning. Our code and data accompanying this paper are available at https://github.com/fishmoon1234/Nonlocal-Attention-Operator.

Figures (3)

• Figure 1: Illustration of the NIPS architecture. NIPS applies physics-informed data augmentation by randomly permuting embedding dimensions to prevent features from being tied to a specific sequence. The data is then tokenized and normalized, followed by several Fourier and linear attention layers that implicitly extract hidden prior knowledge from multiple physical systems. For each system, the final layer maps the last iterative layer feature from contextual input-output function pairs to the corresponding kernel. In downstream tasks, discovering the hidden kernel only requires a forward pass from contextual data, without requiring model retraining.
• Figure 2: Interpretable microstructure discovery in experiment 1.
• Figure 3: Illustration of exemplar MMNIST samples in experiment 2. (a): material parameter field corresponding to different $\bm{b}$. (b): displacement fields (second row) $\bm{u}_x$ corresponding to the same loading field (first row) $\bm{f}_x$. (c): displacement fields (second row) $\bm{u}_y$ corresponding to the same loading field (first row) $\bm{f}_y$.