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

TL;DR

This paper introduces Neural Interpretable PDEs (NIPS), a neural operator architecture that combines Fourier-based kernels with attention mechanisms to improve the scalability, interpretability, and accuracy of physics discovery in complex systems.

Contribution

NIPS advances neural operator design by integrating Fourier space kernels with linear attention, enabling scalable and interpretable PDE learning with superior performance.

Findings

NIPS outperforms existing methods on multiple benchmarks.

NIPS achieves higher predictive accuracy and efficiency.

The approach effectively handles complex physical systems with limited data.

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.