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

We thank the reviewer for the insightful comments. Our response:

Multiple runs with different random seeds: We repeat the experiments three times with randomly selected seeds in the first case of experiment 1, using NIPS and the major baseline models including NAO and AFNO. The results are reported in the table below. Consistent with the findings in the original manuscript, NIPS outperforms the best baseline by 28.9%.

Table 1: Test errors and number of trainable parameters for the Darcy flow problem. Bold numbers highlight the best method.

Additional discussion on Performer and GNO: We appreciate the suggestion and agree that additional discussion on these works will help contextualize NIPS. NIPS is conceptually related to Performer, which introduces kernel-based approximations for efficient self-attention by replacing softmax attention with positive orthogonal random features, achieving linear scaling with sequence length. While both aim to reduce the quadratic complexity of self-attention, NIPS employs a Fourier-domain reformulation that leverages convolutional structure. Unlike Performer’s randomized feature mappings, NIPS deterministically decomposes attention using spectral representations, making it well-suited for physics-based PDE learning. Our work is also related to GNOs, which use graph message passing to model nonlocal interactions in operator learning tasks. While GNOs capture spatial dependencies over irregular domains, NIPS operates in the Fourier domain, leveraging spectral representations to encode long-range dependencies in a structured manner. Whereas graph-based architectures like GNOs and NAO learn operators on discretized graph structures, our Fourier-based approach is particularly efficient on structured grids. We will incorporate this discussion in the revised manuscript.

More rigorous mathematical foundation: We appreciate the suggestion. NIPS inherits resolution invariance and improved identifiability from NAO, and we will add this discussion to the manuscript. While further mathematical analysis is an interesting future direction, our current focus is on developing a novel architecture that enhances predictive accuracy and computational efficiency for both forward and inverse PDE problems.

More detailed analysis of challenging and failure cases: To explore where NIPS might fail, we consider two out-of-distribution scenarios. The training dataset uses a microstructure $b$ generated by a Gaussian random field with covariance operator $(-\Delta + 5^2)^{-4}$ and a loading field $g$ generated by $(-\Delta + 5^2)^{-1}$. The out-of-distribution scenarios are: 1) in-distribution (ID) microstructure $b$ and out-of-distribution (OOD) $g$ using $(-\Delta + 5^2)^{-4}$; 2) OOD microstructure $b$ using $(-\Delta + 5^2)^{-1}$ and ID $g$. Results below show that both NIPS and NAO perform well across these scenarios in terms of forward operator errors and kernel errors, but microstructure errors increase, especially in scenario 2, where the out-of-distribution microstructure is rougher and more detailed. This complexity makes it harder to recover the ground truth, particularly in an unsupervised learning context.

Table 2: Test, kernel, and microstructure errors for Darcy flow. Bold numbers highlight the best method.

Explanation on AFNO performance: AFNO performs poorly because it is designed to learn solutions for a single PDE system and cannot handle inputs from multiple systems. In contrast, both NAO and NIPS leverage global prior knowledge from training data of multiple systems, allowing them to generalize across unseen system states.

Transferring NIPS to fundamentally different PDEs: As mentioned in NAO, attention operators extract prior knowledge in the form of identifiable kernel spaces from multiple PDE tasks. This approach may not work well if the kernel for a new problem differs significantly from those in the training data. For example, if NIPS is trained on diffusion problems with symmetric kernels (stiffness matrices), it will struggle to predict stiffness matrices for non-symmetric systems, like those in advection problems. We will include this discussion in the revised paper.

We thank the reviewer for the insightful comments . Our response:

Formalized interpretability metric: We thank the reviewer's valuable suggestion. A quantitative evaluation of the interpretability can be provided by the recovered kernel. Taking the Darcy flow example for instance, the discovered kernel should correspond to the stiffness matrix $K$, and the underlying permeability field $b(x)$ can be obtained by solving an optimization problem: $B^*=argmin_B \sum_{i,j}||K_B[i,j]-K[i,j]||^2$. Here $B=[b(x_1),\cdots,b(x_N)]$ denotes the pointwise value of $b$ at each grid point, and $K_B[i,j]$ is the corresponding stiffness matrix obtained using a finite difference solver.

In the table below, we take a $11\times11$ grid, and calculate the relative errors of $K$ and $b$ to provide a quantitative evaluation of the interpretability. Beyond the in-distribution (ID) scenario where the microstructure $b$ is generated using a Gaussian random field with covariance operator $(-\Delta + 5^2)^{-4}$ and the loading field $g$ is generated with $(-\Delta + 5^2)^{-1}$, two out-of-distribution (OOD) scenarios are included: 1) data with ID microstructure $b$ and OOD $g$ using covariance operator $(-\Delta + 5^2)^{-4}$; 2) data with OOD microstructure $b$ using covariance operator $(-\Delta + 5^2)^{-1}$ and ID $g$. Intuitively, these two OOD tasks are more challenging. One can see that NIPS achieves the smallest error in both test errors and kernel/microstructure errors almost in all scenarios.

Table 1: Test errors, kernel errors, and microstructure errors for the Darcy flow problem. Bold numbers highlight the best method.

Real-world data to demonstrate practical applicability: We appreciate the reviewer's valuable suggestion. To our best knowledge, most available airfoil flow and climate datasets [1-2] are also generated synthetically by solving PDEs. To address the reviewer's request, we instead consider a real-world data setting, by including additional noise in the dataset. In particular, we perturb the training data with additive Gaussian noise:

$

\widetilde{g}(x) = g(x) + \epsilon(x), \quad \epsilon \sim \mathcal{N}(0, \sigma^2),

$

where $g(x)$ is the true source field, and $\epsilon(x)$ represents zero-mean Gaussian noise with variance $\sigma^2$. The results are presented using the Darcy flow experiment, with $\sigma=0.01$ and $0.1$. From the results, we can see that NIPS's predictions unavoidably deteriorates under the increased level of observational noises, but they remain robust.

[1] Z Li. Fourier neural operator with learned deformations for PDEs on general geometries. JMLR, 2023

[2] J Gupta. Towards multi-spatiotemporal-scale generalized PDE modeling. TMLR, 2022

Add model architecture to appendix: We thank the reviewer's valuable suggestions. Besides Fig 1 and Algorithm 1, we will add an additional section to provide details on the model architecture. We will also release source codes and datasets upon paper acceptance to guarantee clarity and reproducibility of all experiments.

Interpretability variation across different zero-shot conditions: As discussed above, we consider three zero-shot generalization scenarios (ID, OOD for loading $g$, and OOD for microstructure $b$), with results listed in the above table. Both NIPS and NAO generalize well across these scenarios in forward operator errors and kernel errors. However, microstructure errors may deteriorate, particularly in scenario 2 (OOD microstructure $b$). This is expected, as the microstructure in this scenario is not only OOD but also rougher, incorporating more fine-grained details. This added complexity makes recovering the ground truth more challenging, especially in an unsupervised learning setting.

We thank the reviewer for the insightful comments and questions. Our response:

Larger discretization: We demonstrate the scalability of NIPS using larger discretizations (e.g., $121\times121$ with 14,641 tokens in total), and compare its per-epoch runtime and peak GPU memory usage with the best-performing baseline, NAO. The results are summarized in the table below, where ``x'' indicates that the method exceeds memory limits on a single NVIDIA A100 GPU with 40 GB memory due to the explicit computation of spatial token interactions. Compared to the quadratic NAO that fails at 3.7k tokens and the linear NAO that fails at 6.5k tokens, NIPS can easily scale up to 15k tokens on a single A100 GPU. Note that the analysis is performed on a single GPU. To further accelerate training, one can utilize multiple GPUs and leverage distributed training frameworks, such as PyTorch’s Distributed Data Parallel (DDP) module, which will be an interesting direction for future work.

Table 1: Scalability demonstration via per-epoch runtime and peak memory footprint for the Darcy flow problem. All models are 4-layer. The per-epoch runtimes for both the original quadratic attention and the re-formulated linear attention are reported and separated by /. ``x'' denotes exceeding memory limits on a single NVIDIA A100 GPU with 40 GB memory.

Weak-baseline applications: Thank you for pointing us to this insightful work. The referenced paper highlights two key principles for fair comparisons: (1) evaluating models at either equal accuracy or equal runtime, and (2) comparing against an efficient numerical method. First, we clarify that our comparisons focus on state-of-the-art (SOTA) machine learning methods rather than standard numerical solvers used in data generation. Specifically, we evaluate our proposed model, NIPS, against NAO, its variants, and AFNO, which are established baseline models in the field. To ensure fairness (rule 1), we maintain a comparable total number of trainable parameters across models, as large discrepancies could lead to misleading conclusions. Within this constraint, we assess (1) per-epoch runtime to measure computational efficiency and (2) test errors to evaluate prediction accuracy. Regarding rule 2, we select SOTA neural operator models as baselines to demonstrate the advantages of NIPS in accuracy and efficiency. These models represent the strongest available baselines for learning-based PDE solvers. Therefore, our experimental setup adheres to the principles outlined in the referenced paper. We will incorporate this discussion into our revised manuscript to further clarify our evaluation methodology.