Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery
Authors: Ning Liu
Yue Yu
Venue: ICML 2025 poster (2025)
Keywords: Neural Operators
Physics Discovery
Attention Mechanism
Inverse PDE problems
Primary area: applications->chemistry_physics_and_earth_sciences

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.

Neural Interpretable PDEs:
Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery
Ning Liu 12 Yue Yu
Abstract
1. Introduction
Interpretability in machine learning (ML) models is of paramount importance, partieularly in the study of phys- ical systems, where grasping the underlying governing prin- ciples is as critical as achieving high predictive accuracy. Unlike conventional data-driven tasks, physics-based prob- lems necessitate models that not only yield reliable outputs but also elucidate the fundamental mechanisms driving ob- served phenomena. This need especially acute in high- stakes domains such as materials science, healthcare, and engineering, where ML-driven insights can directly influ- ence human safety and technological progress (Coorey et al., 2022; Ferrari & Willcox, 2024; Liu et al., 2025). However, extracting physical laws from data remains a formidable challenge, as many learning tasks-such as deducing mate- rial properties from deformation fields are inherently an ill-posed inverse problem, lacking direct supervision and admitting multiple plausible solutions (Hansen, 1998). That means, although MI models may serve as effective surro- gates for predictive tasks, their inferred internal representa- tions may diverge from true physical parameters, leading to spurious interpretations and potential failures practical ap- plications. Thus, advancing interpretability not merely a matter of transparency but essential endeavor in ensuring the development of robust, trustworthy, and scientifically principled ML models for physics-based learning.
Attention mechanisms have emerged trans- formative tools in core AI domains such 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 severely ill-posed inverse PDE problem. In this work, we introduce Neural Interpretable PDEs (NIPS), novel neural operator architecture that builds upon and enhances Nonlocal Attention Operators (NAO) in both predictive accuracy and computational efficiency. NIPS employs linear attention mechanism to enable scalable learning and integrates learnable kernel network that acts as a channel-ine -inenneednddeeeddenn convolution in Fourier space. As a consequence, NIPS eliminates the need to explicitly compute and store large pair- wise 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 laap in scalable, interpretable, and efficient physics learning. Our code and data accompanying this paper are available at https://githii ccm//iimmoon1234 Nonlocal -ttpral--tteteetttteooeeeeeee
Uncovering interpretable mechanisms for physical systems presents a fundamental challenge: inferring governing laws that are often high- or infiniteddimensional from data con- sisting of discrete measurements of continuous functions. data-driven surrogate model must therefore learn not only the mapping between input-output function pairs but also the latent representations that govern the system's behav- ior. From partial differential equation (PDE) perspective, constructing a surrogate model corresponds to solving a forward problem, whereas identifying the underlying gov- erning mechanism constitutes an inverse problem, which is often rank-deficient even ill-posed, especially when data is limited. Neural network models exacerbate this ill- posedness due to their intrinsic approximation biases (Xu et al., 2019), making the inference of governing laws par- ticularly challenging. To address this, recent deep learning
'Department of Mathematics, Lehigh University, Bethle- hem, PA 18015, USA *Global Engineering and Materials, Inc., Princeton, NJ 08540, USA. Correspondence Yue <yuy214@lchin lcdii..................
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
approaches have been developed as inverse problem solvers (Fan & Ying, 2023; Molinaro al., 2023; Jiang al., 2022; Chen et al., 2023; Lu & Yu, 2025), aiming to reconstruct unknown parameters from solution data. These approaches generally rely on effectively embedding prior information, in terms of governing equations (Yang et al., 2021; Li et al., 2021), regularization techniques (Dittmer et al., 2020; Ob- mann et al., 2020; Ding et al., 2022; Chen et al.. 2023), or additional operator structures (Uhlmann, 2009; Lai al., 2019; Yilmaz, 2001). However, in complex systems, such priors are often unavailable or highly problem-specific, lim- iting the generalizability of these approaches. Consequently, existing methods can only solve inverse problems for a fixed system, requiring a complete reconfiguration when the sys- tem evolves- as when material properties degrade in a material modeling task.
We conduct zero-shot learning experiments on unseen physical systems, demonstrating both the generaliz- ability of NIPS in predicting system responses and its capability in discovering hidden PDE parameters.
2. Background and Related Work
Neural operators for hidden physics learning. Learning complex physical systems from data has become increas- ingly prevalent in scientific and engincering fields (Karni- adakis et al., 2021; Liu et al., 2024b; Ghaboussi et al.. 1991; Carleo et al., 2019; Zhang al., 2018; Liu et al., 2023b; Cai et al., 2022; Pfau et al., 2020). In many cases, the underly- ing governing laws remain unknown, concealed within the data, and must be uncovered through physical models. Ide- ally, these models should be interpretable, enabling domain experts to derive meaningful insights and make further pre- dictions, thereby advancing the understanding of the target physical system (Jafarzadch et al., 2024; Wang et al., 2025). Furthermore, they should be resolution- intiriinnnnvvrannt, capable of handling data at different scales. Neural operators (NOs) are designed to map between spaces (Li et al., 2020a;e; You et al., 2022a; Ong et al., 2022; Cao, 2021; Lu et al., 2019; 2021: Goswami et al., 2022; Liu et al., 2024a), making them powerful tool for discover- ing continuum physical laws by capturing the relationships between spatial and spatio-temporal data. A major mile- stone in operator learning is the Fourier Neural Operator (FNO) (Li et al., 2020c; Liu et al., 2023a), which utilizes the Fast Fourier Transform (FFT) to efficiently perform convolu- tions in Fourier space. This approach reduces computational complexity, enabling scalable learning of nigh-ddmmnsionee physical systems while retaining the capacity to capture intricate physical interactions. By incorporating FFT, NOs can Iearn kemnlls with truncated modes in Fourier space, enhancing efficiency and making them a promising tool for complex modeling physical systems across various scales.
In this work, we introduce Neural Interpretable PDFs (NIPS), a novel attention-based neural operator architecture that extends and enhances Nonlocal Attention Operators (NAO) (Yu et al., 2024) in both predictive accuracy and com- putational efficiency. NAO introduces an attention-base kernel map to simultaneously learn the operator and its asso- ciated kernel, effectively extracting hidden knowledge from diverse physical systems. By implicitly enforcing prior in- formation and exploring the function space of identifiability for the kernels, the attention mechanism enables automatic inference of underlying system contexts in an unsupervised manner (Lu & Yu, 2025). Nevertheless, NAO relies on an attention mechanism of quadratic complexity and explic- itly models spatial interactions through square projection matrix, which are both time- and venmory-ntensive. In con- trast, NIPS adopts a linear attention mechanism for scalable learning and incorporates a learnable kernel network that functions as a space. This design eliminates the need to compute and store large pairwise interactions explicitly, effectively amortizing the cost of spatial interactions through the Fourier transform. Our key contributions are:
function
independent convolution in Fourier
Inverse PDE solving. Inverse PDE solving involves un- covering the underlying physical laws or parameters from observational data (Liu et al., 2024c). This task is partic- ularly challenging due to the high dimensionality of the problem and the need to infer unknowns from limited or noisy data. Recent progress in inverse PDE solving has led to the development of innovative approaches (Cho & Son, 2024), such as Neural Inverse Operators (NIOs) (Molinaro et al., 2023) and Nonlocal Attention Operators (NAOs) (Yu et al., 2024). NIOs are designed to map solution operators to the underlying PDE parameters as functions in super- vised setting, leveraging a composition of DeepONets and FNOs. This architecture has demonstrated superior per- formance in solving inverse PDE problems, outperforming traditional optimization methods. On the other hand, NAOs employ attention mechanisms to capture nonlocal interac- tions among spatial tokens in an unsupervised setting. They
We introduce Neural Interpretable PDEs (NIPS), a novel iitention-basee neural operator architecture for simultaneous physics modeling (as forward PDE solver) and governing physical mechanism discovery (as an inverse PDE solver).
NIPS integrates attention with a learnable kernel net- convolution in Fourier
work for channel-independen space, eliminating the need to compute and store large pairwise interactions and amortizing the cost of han- dling spatial interactions into the Fourier transform.
We reformulate the original NAO using linear attention, enabling scalable learning in large physical systems. Together with Fourier convolution, NIPS effectively harmonizes Fourier insights with attention for scalable and interpretable physics discovery.
2

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
effectively mitigate ill-posedness and rank deficiency by im- plicitly extracting global prior knowledge from training data of multiple systems (Lu & Yu, 2025). NAOs can generalize across unseen data resolutions and system states, demon- strating their potential as interpretable foundation models for physical systems.
tions across different frequency modes, while also provid- ing efficient evaluation of domain integrals. As a result, NIPS delivers an accurate, efficient, and interpretable PDE solution operator.
3.1. Problem Setting
Attention mechanism. Attention mechanisms (Vaswani et al., 2017) are renowned for their remarkable in-context learning capabilities in language models. In particular, they enable zero-shot generalizability, allowing models to adapt to new tasks from just a few examples in the prompt with- out retraining (Vladymyrov et al., 2024; Lu et al., 2024; Oko et al., 2024). Recently, attention mechanisms have become integral to learning hidden physics, especially in tasks involving PDE solving. Beyond generalizability, at- tention mechanisms enhance model capacity by capturing long-range dependencies and spatiotemporal interactions, making them highly effective for modeling complex phys- ical systems. To improve the accuracy of forward PDE solvers, linear attention has been employed in NOs to re- place softmax normalization and act as learnable kernel (Cao, 2021). This relaxes the quadratic complexity require- ment while maintaining the ability capture intricate physi- cal relationships. Further innovations include the integration of Galerkin-type linear attention within encoder-decoder ar- chitectures (Li et al., 2022), hierarchical transformers for multiscale learning (Liu et al., 2022), and heterogeneous nor- malized attention mechanisms for handling multiple input features (Hao et al., 2023). On related note, polynomial attention (Kacham al., 2023; Deng al., 2024) offers promising alternative to linear attention by replacing the softmax function with polynomial kernels, enabling efficient computation of attention scores, particularly for long-range dependencies and complex physical systems. Recent studies have applied attention mechanisms not only to solve individ- ual PDEs but also to tackle multiple types of PDEs within unified framework, demonstrating the potential of these methods in advancing interpretable and scalable models for physical mechanism discovery (Yang & Osher, 2024; Ye et al., 2024; Sun et al., 2024; Zhang, 2024).
We consider a series of PDEs with different hidden physical parameters:
= x EEQ
(I)
Here, 2 C CR' is the domain of interest, f(x) represents the loading function on Ω, and u(x) is the corresponding solution of the system. Ls denotes the unknown govern- ing law, such as balance laws, which is determined by the (possibly unknown and nigh-dimensiona)) parameter field b. For example, in material modeling, Ls typically repre- sents the constitutive law, and b can be a vector (b E Rd) representing the homogenized material parameter field, or a vector-valued function (b L ( 124))) )) representing the heterogeneous material properties.
Many physical modeling tasks can be formulated as either forward or inverse PDE-solving problems. In a forward problem setting, the goal is to find the PDE solution given the PDE information, such as coefficient functions, bound- ary/initial conditions, and loading sources. Concretely, given the governing operators X., the parameter (field) b. and loading field f(x) in (1), the objective is to solve for the corresponding solution field u(x:) using either classi- cal PDE solvers (Brenner & Scott, 2007) or data-driven approaches (Lu et al., 2019; Li et al., 2020c). As a result, a forward map is constructed:
G (b, f) - - u.
(2)
Here, b and f are input vectors/functions, and u is the output function.
Conversely, solving an inverse PDE problem involves re- constructing the underlying full or partial PDE information from solutions, where one seeks to construct an inverse map:
H (u,, f) -
Solving an inverse problem is typically more challenging due to the ill-posed nature of the PDE model. In general, limited number of function pairs (u, f) from single system is insufficient for inferring the underlying parameter field b, rendering the inverse problem generally non-identifiable (Molinaro et al., 2023).
3. Neural Interpretable PDEs
This section introduces NIPS, a Fourier-domain kernel op- erator that simultaneously solves both forward and inverse PDE problems. Compared to existing NOs, NIPS incor- porates two key innovations. First, it constructs a data- dependent neural operator through the linear attention mech- anism, parameterizing an inverse mapping from data to the underlying kernel and offering a generalizable kernel method for PDEs with varying hidden parameter fields. Sec- ond, the attention mechanism is combined with spectral convolution to enhance expressivity by enabling interac-
Formally, we consider S training datasets from different systems, each of which contains d function pairs (u, f) u x F:
Drr = (((a(( 1
(4)
The index 1) corresponds to different hidden PDE parameter field/set 6ad E B. Here, 4 C (2: c L2( R() represent the (infinite-dimensionn Banach spaces of the
3

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery
Training
Iterative Fourier & linear attention block
Output Fourier kemel block
Downstream application
Load-displaeeen displacemet pairs
Zero-shot generelization
outgut
outgut
matanul
nommallzze
masmual
komel
Data sugmentation via minoom permutation (rheer =3)
matmul
TiDeur
Tinsar
linear
linear
fast
K[:d]
V[:d]
kemel
concat
Interpoetation flaverse PDE) 4
Prediction (forward PDE)
input
inpult
input
Myst
Tokenize
Iterativee Fourier linear attention block
Output Fourier kernel block
Datai (contertt)- dependent kemel
nomallze
Figure 1. Illustration of the NTPS architecture. NIPS applies physics-infrmeed data augmentation by randomly permuting embedding dimensions to prevent features from being tied to specific sequence. 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 kemel only requires forward pass from contextual data, without requiring model retraining.
solution function u and loading function f. respectively, and B is finite-dimensional manifold where the hidden parameter b takes values on. Our goal is solve both the forward and inverse PDE problems simultaneously from multiple systems. That means, given Dir and measure- ments (()(e*(())) ( f(esst (x=)) corresponding to a new and hidden parameter field fest we aim to infer the key information about brest and provide a surrogate mapping glest frest us
Consequently, the forward solver G is tailored specifically for that particular PDE.
The first key feature of NIPS is the transition from a static kernel to data-dependent kernel. In particular, we propose to replace Ko with kernel map K(u11dd f1.a] in the form:
GK, [x,[f(<x)
K Klunsd fisd] (ar, y) f
(6)
test without additional training tuning.
where a E St and the nonlinear kernel map is constructed using L numbers of iterative attention blocks:
3.2. Attention-Batt Kernel Map
9, (0 (x) = (x), o (0) (x) = (2),
To develop generalizable forward PDE solver capable of handling multiple PDEs, our architecture is inspired by the kernel method for PDEs (Evans, 2002). In this method, the solution of a linear PDE is formulated as an integral involving a kernel function that captures the influence of one point space on another:
9 n 1 (x)
(1, | (x.
9 1- 12 (t)
11
dy
1515I, I, with
(7)
([g ( ( W. (lv g Wi
Klg w00 Wi]
k (x,
Klu!t), wil) Wil
(8)
(f) = / K2 y) fy)dd = u(x)) a ΩQ. (5)
Kip.q: =] Σ / w" (a. (x (z) I (l) ds.
Here, the kernel function depends on the parameter field b. In generic integral NOs (Li et al., 2020b;e; You al., 2022a;b) such FNOs and GNOs, the kernel Ko is pa- rameterized as either a convolutional kernel or shallow MLP and is then optimized to minimize the data loss Dd [ u:ll on function pairs from a single PDE.
Here. N is a chosen normalizer, w viL W. (W/)T characterizes the (trainable) interaction of d input function pair instances, W.P is projection func-
tion, and 0 ([W, (x, 2), W. W W are learnable

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
functions/parameters. By setting (x) == u,(2)) we notice that the kernel K [(ui,d, fial can be expressed as: K (uida, fi.d; (x, v)
where R is learnable Fourier kernel that substitutes F(WP), we and WK are learnable matrices. Here, F and F denote the Fourier transform and its inverse, respectively. This formulation reduces both parameter count and computational overhead. When taking WP as a point-wise evaluation on the grid, it requires O(N3) memory, whereas learning R directly Fourier space re- duces memory footprint to ((m). with m < . being the number of retained Fourier modes. Consequently, the computational cost for the first integral is reduced from ((Nad) to ONd N), leading to an overall complexity of O(Nd df + No N) for NIPS.
Σ w
WP (xx z) (g (2) ( (v)) dz
W (a. 10
01 ( W2 [|g (w) dz.
w,
Hence, the proposed architecture simultaneously serves as an inverse PDE solver, mapping (u, pairs the parameter-deeendent kemel Kb. which provides forward PDE solver in the form of a kernel method as in 9x.
3.4. NIPS: A Generalizable and Efficient PDE Solver
3.3. Efficient Fourier Kernel Learning
With the aforementioned modifications, the model is illus- trated in Figure I, along with the training algorithm summa- rized in Algorithm 1.
Although the attention block in (7) enables generalizable kernel for multiple PDEs, it is inherently more computation- ally expensive than its single-PDE counterpart. Comparing (7) and (5), one can see that the single-PDE solver involves a spatial integral per layer update, whereas the proposed multi- PDE solver requires evaluating two integrals over the entire domain Ω: the first integral in (8) maps data pairs to the kernel K, while the second performs the kernel-weighted integral in (7), updating the data pairs for the next block. When these integrals are computed on discretized domain with N grid points, both require d( flops, making the model computationally expensive on large discretizations.
4. Experiments
We assess the performance of NIPS across a broad spec- trum of physics modeling and discovery tasks, focusing on several key aspects. First, we highlight the advantages of using Fourier kernel with linear attention, and compare NIPS to various baseline models, including NAO with full LayerNorm (NAO-f), NAO with optimized normalization (NAO), NAO with full learnable square projection (NAO- WP), and a NO with a convolution-based attention mech- anism (AFNO (Guibas et al., 2021)). Our evaluation em- phasizes generalizability, particularly zero-shot prediction performance when modeling new physical systems with unseen governing equations, as well as performance across different resolutions. Additionally, we examine the data efficiency-accurae trade-off in inverse PDE learning tasks and assess the interpretability of the learned kernels. All ex- periments are optimized using the Adam optimizer. For fair comparison, we tune syperppraameter--nnlldiiii learning rates, decay rates, and regularization parameters- to mini- mize training loss. Experiments are conducted on a single NVIDIA A100 GPU with 40 GB of memory. Additional details data generation and implementation can be found in the Appendix.
The class of shift-equivariant kernels can be decomposed as a linear combination of cigenfunctions, enabling efficient multiplication in the eigen-transform domain and accelerat- ing integral evaluations. For single-PDE solvers, a notable example is the FNO model (Li et al., 2020c), which enforces translation invariance on the kernel (b y) := K3 ( y) in (5), reducing the computational complexity from ( to ONN log N). Inspired by this, the second core feature of NIPS enhances efficiency by enforcing convolutional structure and computing the integral as element-wise multi- plication in the frequency domain.
To alleviate the computational cost, we first notice that the proposed architecture can be reformulated from linear attention perspective:
S, Klp.q; W] y) (y)dy=
4.1. Darcy Flow
1 J. /. W (a, 2) (p(=)W9)em dz L. (K) T
q (v)dd
Vdu This reduces the computational complexity of the kernel- weighted integral in (7) from O(N2d) to O(Ne)). Addi- tionally, we impose convolutional structure on W., and evaluate the first integral in the spectral domain:
We begin by examining NIPS's ability to simultancously learn both forward and inverse solution surrogates on the 2D Darcy flow benchmark. Specifically, we consider the modeling of 2D subsurface flows through a porous medium characterized by a heterogeneous permeability field. Follow- ing the setup in Yu et al. (2024), the high-fidelity synthetic simulation data for this problem are governed by the Darcy's s flow. In this setting, the physical domain is Ω [0, 1J2, b(x) represents the permeability field, and the Darcy's is equation
K E.por q: (z y)q, (w)dy=
F- 1 F(pw / ( (u)dy,
dd
5

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
Algorithm 1 The overall training algorithm and architecture of NIPS.
1: Given Der == (x) u. (x)) - for training, Drest = (x), (x))
T n=1, for test sets. Each corresponds
to a different PDE parameter field b
2: Training Phase:
3: for ep = 1 epochmar do
4: Update = {[R, w wiK using Adam with respect the optimization problem:
T
I1 KI ui (z. y) -u,( u,(()) (12(0)
0 argmin
(9)
T 1=1,f=1
| (x) 12(0)
where K K[FT.d "yul is defined as:
9 (0) (x) = (ar); 0 (ar) = =)((x))
9 (x) -g (1-11 11 (x) +N
1
F F 1(g w() /
F
9e (y)dy
Vdk
1 F (Ri
F(g
(
(w) of (W)"))))
of (x) = (1
(x) +N
_F F(O () / (D)-)(
1
fj
(y)dy
Vdl
1
(R. F(v F(uli-1) (1)1(z))) / (9 (o 95 (1)))) 151<L-1,
dx
K|uiad fi:d: (x, v) = F F-1(RR (RL F(g LL (Z))()) ()
d
(1) ())-)())W)) "
F F
Rs
(10)
ddk
5:
end for
6: Test Phase:
7: Compute prediction error on the test dataset, an evaluation for the forward PDE solver:
Toter .d
Sa K K( ai (z) y) (1)dy (iaalyy-((aaa ūx(x)) [2(())
Elest forward
Trestd
i x
(2())
1=1,11
If the ground-truth microstrnctturekeemell is available, compute kernel error on the test dataset, as an evaluation for the inverse PDE solver:
K K if (1)((x y) - Kōr (x, v) (2((X))
Elseerse =
Trest
| (xx, (2x,)))))) (
q=1
is given by: -V- (V-((())Vpaa)) = 9xx), x x Q.
NIPS by comparing its performance against its variants (i.c., NAO, NAO-f, NAO-W"), using a fixed embedding dimen- sion of = 50, query-key projection dimension of dg = 40, and a sequence length of 21 x 21 441. The results are re- ported in Table All models are configured with two layers. Notably, NAO-f incorporates full LayerNorm in every layer, while NAO-W/P employs a layerwise learnable projection, both of which significantly increase the number of trainable parameters. To isolate the effects of these modifications, we
p(x) =0 0, x E a..
(11)
We aim learn the solution operator of Darcy's equation and simultancously compute the pressure field p(æ) and discover the underlying (hidden) permeability field b(x).
Ablation study. We first conduct an ablation study on
6

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery
keep the remaining architectural components unaltered. We first examine the impact of LayerNorm placement in NAO and NAO-f. In NAO, LayerNorm is applied across both the token and projection dimensions in the first layer, but only across the projection dimension in subsequent layers. In contrast, NAO-f normalizes across both dimensions in all layers. A comparison of test errors reveals that NAO outper- forms NAO-f, achieving 10.4% improvement in prediction accuracy while using 46.6% fewer trainable parameters. These results indicate that applying full LayerNorm across all layers is neither necessary nor beneficial, as it not only drastically increases the number of trainable parameters but also degrades model performance.
Table 1. Test errors, per-epoch runtime, and number of trainable parameters for the Darcy flow problem, where bold numbers high- light the best method. The per-epoch runtimes for both the original quadratic attention and the re-formulated linear attention are re- ported and separated by "f" where applicable.
Per-epoch time (s)
Test
Case
Model
#param
error
NIPS
98,396
1.6
2.28 %
2 layers, 441 tokens
NAO
101,134
5.4/1.7
3.09%
NAO-f
189,234
5.4/1.7
3.45%
(21 21)
NAO-W'
1,658,746
5.8/1.9
4.00%
AFNO
101,200
1.6
52.09%
Next, we adopt the optimal LayerNorm placement and in- vestigate whether incorporating layerwise learnable projec- tions improves performance. Comparing NAO with NAO- WP. we observe that the number of trainable parameters increases by more than 15 times, while accuracy deterio- rates by 29.4%. This decline stems from the large square projection weights, whose size scales quadratically with the sequence length, explicitly computing and storing exten- sive pairwise token interactions. The substantial increase in trainable parameters also exacerbates overfitting, further degrading model performance. In contrast, NIPS achieves the highest accuracy while using the fewest trainable param- eters among all models, outperforming the best NAO variant by 26.2% in accuracy. This improvement is attributed to NIPS's ability to implicitly model token dependencies by learning to interact with spectral modes, which encode more concise representation of the data's 's intrinsic behavior. As result, the model focuses more effectively on mean- ingful patterns while reducing redundancy. An additional advantage of learning spatial interactions in Fourier space is the substantial reduction in trainable parameters. Unlike NAO where weight size scales with the total number of spa- tial tokens, NIPS scales with the number of spectral modes considered, leading to significant computational savings. This efficiency gain is evident in per-epoch runtime compar- isons: while NIPS slightly outperforms the linear version of NAO when the sequence length is relatively small, it accel- erates training by a substantial 28.8% when the sequence length increases to 41 x 41 = 1681, accompanied by a 23.5% improvement in predictive accuracy.
layers,
NIPS
108,692
2.5
1.03%
141 tokens (21 x 21)
NAO
109,494
7.8/2.5
1,45%
AFNO
121,600
52.84% 1.9
layers. 1681 tokens (41 41)
NIPS
366,932
8.4
1.14%
NAO
357,494
38.8/11.8
1.49%
AFNO
364.000
24.5
98.70%
model depth increases.
trysics-informer data augmentation. The above studies are conducted with a fixed number of random permutations, Ryand = 100, and a projection dimension of de = 40. Here, we investigate the impact of these two parameters. Random permutations in the embedding dimension help encode physical knowledge in the data by ensuring that features in the embedding space are not tied to a specific sequence order. This allows the model to disregard sequence effects in the embedding dimension and instead focus on learning spatial dependencies across tokens, which enhances generalization. Meanwhile, the projection dimension dç determines the level of information compression. As shown in Table 2, model performance improves monotonically with increasing Arand and dy, eventually saturating when both parameters become sufficiently large.
Interpretable physics discovery. To demonstrate the phys- ical interpretability of the learned kernel, we present in the first row of Figure the ground -tuth microstructure b(x). a test loading field instance g( (m), and the corresponding solu- tion p(æ). ). By summing the kernel strength along each row, one can discover the interaction strength of each material point x with its neighbors. Since this strength correlates with the permeability field b(x). the underlying microstruc- ture can be recovered. The bottom row of Figure shows the discovered microstructure for this test case. Due the continuous nature of the learned kernel, the discovered mi- crostructure appears smoothed, necessitating a thresholding step to distinguish the two-phase structure. The discovered microstructure (bottom-midde plot) matches well with the hidden ground-trun microstructure (top-left plot), except near the domain boundary. This discrepancy arises from the applied Dirichlet-type boundary condition (p(x) = on
Comparison with additional baselines. We then com- pare the performance of NIPS against an additional baseline, particularly NO with convolution-base attention mecha- nism, i.c., AFNO. To ensure a fair comparison, we maintain a comparable number of trainable parameters across differ- ent layer configurations and input sequence lengths. While AFNO's runtime is similar to that of NIPS, it fails to achieve zero-shot generalization unseen PDE parameters, with test errors consistently exceeding 50% across all considered scenarios. Moreover, its performance deteriorates further as
7

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
from (-A + 5-)-2 and 2) OOD microstructure sampled from (-A + 52)-1 with ID loading field g. Intuitively, these OOD tasks are more challenging due to distributional shifts. Nevertheless, NIPS consistently achieves the low- est errors in both test performance and in the recovery of the kernel and microstructure across nearly all scenarios. Additionally, we consider a real-world data setting by in- troducing additional noise into the dataset. In particular, we perturb the training data with additive Gaussian noise: π(x) g(z) + e (N000)) where g(x) denotes the true source field and e(x)) represents zero-mean Gaus- sian noise with variance a². Experiments are conducted for noise levels a = 0.01 and 0.1. As shown in Table 3, while NIPS's predictive accuracy naturally degrades with increasing noise, it remains robust overall, demonstrating resilience observational perturbations.
Table Parametric studies the effect of the number of random permutations and the size of d2. A 2-layer NIPS with 441 tokens adopted for demonstration.
Case
Brand
Tétrain
Test error
1
90
42.04%
2
180
28.98%
5
450
18.37%
d, 40
10
900
9.02%
25
2250
3.52%
.50
4500
2.53%
100
9000
2.28%
Case
dx
Terain
Test
5
4500
14.67%
10
4500
6.93%
20
4500
4.03%
Arand 50
30
4500
3.01%
40
4500
2.53%
50
4500
2.45%
4.2. Mechanical MNIST Benchmark
In this experiment, we study the learning of heterogeneous and nonlinear material responses using the Mechanical MNIST (MMNIST) dataset (Lejeune, 2020). MMNIST consists of 70,000 heterogeneous material specimens un- dergoing large deformations, each governed by a Neo- Hookean material model with a varying modulus derived from MNIST bitmap images. These material specimens are subjected to random 2D uniaxial extension, shear, equibiax- ial extension, and confined compression loadings, with their corresponding material responses captured as prediction tar- gets. The total number of tokens is thus 29 x 29 x 2 1682, where "2" indicates the two in-plane directions.
Microstructure b(x)
Loading g(x)
Solution p(x)
Kernel
Recovered b(X)
Error b(x): 18.96%
In this and the following experiments, we directly compare NIPS with the best- pest-performing baseline, NAO. Since the resolution is fixed 29 29 in the MMNIST dataset, evaluate model performance across different depths. As shown in Table 4, NIPS consistently outperforms NAO in both the 2-layer and 4-layer configurations, achieving per- formance gains of 59.2% and 78.9%, respectively. This underscores the advantages of implicitly learning token in- teractions in Fourier space, where NIPS captures intrinsic features (i.e., spectral modes) in data rather than relying on extensive discrete token dependencies.
Figure 2. Interpretable microstructure discovery in experiment 1.
822) in samples, which prevents the measurement pairs (p(z). (z)) from containing information near 02, making it impossible to identify the kernel at the boundaries.
With additional physics knowledge, one can quantitatively assess the interpretability provided by the recovered ker- nel. Specifically, the learned kernel should correspond to K-!, where K is the stiffness matrix of (11). Given this, the underlying permeability field b(rr) can be re- covered by solving the following optimization problem: 8" = argming g Ex. | [i.j] K (i.li where B = ==bb))).... b(Ix)] denotes the pointwise value of at each grid point, and Kn ]n represents the corresponding stiffness matrix obtained using finite difference solver. In Table we report the relative emrrr of both K and on an 11 x 11 grid to quantitatively assess interpretability. Beyond the in-distribution (ID) setting where the microstructure is sampled from a Gaussian random field with covariance oper- ator (- - and the loading field from (-A -
4.3. Synthetic Tissue Learning
In the final experiment, demonstrate NIPS's capability in learning synthetic tissues with highly organized structures, where collagen fiber arrangements vary spatially. Under- standing the underlying hidden microstructural properties within the latent space of the complex, high high-dimensional tissues is crucial, as they are inferred from experimental mechanical measurements. In this setting, we loading data with a sequence length of 29 x 29 x 2 1682 as input and aim to simultaneously learn forward solution opera- tor to predict the corresponding displacement field while
we also consider two out-of distrribtitiii (OOD) scenarios: 1) microstructure with OOD loading field g, sampled
8

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery
Table
Test emror,, kernel error,, and microstracture errors for the Darey flow problem, where hold numbers highlight the hest method.
Test error
Kernel error
Microstructure error
Data Setting
Case
Model
#param
ID
No noise
NIPS
327,744
4.09%
9.24 %
7.92%
NAO
331,554
4.15%
10.35%
7.09 %
OOD, Scenario
No noise
NIPS
327,744
3.40%
10.63 %
11.23 %
NAO
331,554
6.86%
14.46%
20.37%
OOD, Scenario
No noise
NIPS
327,744
4.98%
9.68% %
16.06 %
NAO
331.554
5.57%
10.94%
20.30%
ID
Noise ~ 100,0.01%,
NIPS
327,744
4.40%
9.47%
7.97%
OOD, Scenario
Noise c ~ (N00 0.013)
NIPS
327,744
3.88%
11.14%
13.63%
OOD, Scenario
Noise ~ .N(0,
NIPS
327,744
5.65%
10.21%
17.36%
ID
Noise (V(0,
NIPS
327.744
9.98%
21.77%
14.84%
OOD, Scenario
Noise ~ (O000.13
NIPS
327,744
3.84%
21.08%
19.87%
0OD, Scenario
Noise ~ 0.11..1
NIPS
327,744
11.67%
22.77%
26.54%
somewhat marginal performance improvement relative to the increased model complexity, we opt not to pursue this direction further.
Table 4. Test errors and number of trainable parameters for the MMNIST benchmark. Bold numbers highlight the hest method.
Case
Model
#param
Test error
layers with
NIPS
720,600
2.11 %
5. Conclusion
1682 tokens
NAO
722,748
5.17%
We present NIPS, novel attention-based neural operator architecture for simultancous forward PDE solving and gov- erning physics discovery. NIPS integrates the attention mechanism with learnable kernel network for channel- independent convolution in Fourier space, effectively elimi- nating the need to compute and store large pairwise interac- tions and amortizing the cost of handling spatial interactions into the Fourier transform. Through reformulation in lin- ear attention perspective, NIPS harmonizes Fourier insights with attention for scalable and interpretable physics dis- covery. Our zero-shot learning experiments demonstrate that NIPS generalizes effectively across both forward and inverse PDE learning tasks, thereby enhancing physical in- terpretability and enabling more robust generalization across diverse physical systems.
layers with 1682 tokens
NIPS
768,600
1.11%
NAO
763,548
5.27%
Table 5.
Test errors and number of trainable parameters for experi- ment 3, where bold numbers highlight the best method.
Case
Model
#param
Test error
NIPS
768,600
4.95%
4 layers with
NIPS-mlp
2,183,582
4.52%
1682 tokens
NAO
763,548
5.62%
uncovering the underlying microstructures in the tissues.
We present our experimental results in Table 5, compar- ing the 4-layer NIPS with NAO of the same depth. This problem exhibits significant nonlinearity due to the inherent complexity of the tissue microstructure, Nevertheless, NIPS successfully learns an accurate solution operator, achieving a test error within 5% and outperforming NAO by 11.9%. Notably, NIPS's predictive accuracy can be further improved by passing the V vector through an MLP before its matrix multiplication with the K vector in the final layer, enabling nonlinear mixing of the loading data. To illustrate this idea, we introduce simple two-layer MIP with skip connection and report the results in Table 5 as NIPS-mlp. This modifica- tion improves accuracy by an additional 7.7%, with further gains possible through careful MLP design. However, also nearly triples the number of trainable parameters. Given the
Limitations: As discussed in NAO, the core idea of attention-batn operators is to extract prior knowledge across multiple PDE tasks, in the form of an identifiable kernel space. This approach assumes that the target kernel shares structural similarities with those seen during training. Consequently, performance may degrade when the target kernel significantly deviates from the training distribution. For example, if the operator is trained solely on diffusion problems, where all kernels (i.e., stiffness matrices) are sym- metric, the learned prior will inherently favor symmetric structures. As result, NIPS may struggle to predict the stiffness matrix for systems with non-symmetric kernels, problems.
such as those arising

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
Acknowledgments
tackle cardiovascular disease A review of an emerg- ing interdisciplinary field. NPJ Digital Medicine, 5:126, 2022.
The authors would like to thank Yiming Fan for the help- ful discussion. Y. Yu would like to acknowledge support by the National Science Foundation (NSF) under award DMS-2427915, the Department of the Air Force under award FA9550-22-1-01999 and the National Institute of Health under award research were conducted on Lehigh University's Research Computing infrastructure partially supported by NSF Award 2019035.
Deng, C., Yue, Z., and Zhang, Z. Polynormer: Polynomial- expressive graph transformer in linear time. arXiv preprint arXiv:2403.01232. 2024.
Portions of this
Ding, W., Ren, K., and Zhang, L. Coupling deep
learning with full waveform inversion. arXiv preprint arXiv:2033.0199. 2022.
Dittmer, S., Kluth, T., Maass, P., and Otero Baguer, D. Regularization by architecture: A deep prior approach for inverse problems. Journal of Mathematical Imaging and Vision, 62:456-470, 2020.
Impact Statement
This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Evans, C. Partial Differential Equations (Graduate Stud- ies in Mathematics, Vol. 19). American Mathematical Society, Providence, Rhode Island, 2002.
References
Fan, Y. and Ying, L. Solving traveltime tomography with deep learning. Communications in Mathematics and Stafistics, 11(1):3-19, 2023.
Brenner, S. and Scott, R. The mathematical theory of fi- nite element methods, volume 15. Springer Science & Business Media, 2007.
Ferrari, A. and Willcox, K. Digital twins mechanical and aerospace engineering. Nature Computational Science, (3):178-183, 2024.
Cai, S., Mao, Z.. Wang, Z., Yin, M., and Karniadakis, G. E. Physics-informed neural networks (PINNs) for fluid me- chanies: A review. Acta Mechanica Sinica, pp. 1-12, 2022.
Ghaboussi, J., Garrett Jr, J., and Wu, X. Knowledge-based modeling of material behavior with neural networks. Jour- nal of engineering mechanics, 117(1):132-1))) 1991.
Cao. S. Choose transformer: Fourier galerkin. Ad- vances in neural information processing systems, 34: 24924-24940, 2021.
Goswami, S., Bora, A., Yu, Y., and Karniadakis, G. E.
Physics-informee informed neural operators. 2022 arXiv preprint arXiv:2207.05743 2022.
Carleo, G., Cirac, I., Cranmer, K., Daudet, L., Schuld, M., Tishby, N., Vogt-Maranto, L., and Zdeborová, L. Machine learning and the physical sciences. Reviews of Modern Physics, 9440002 2019.
Guibas, J., Mardani. M., Li, Z., Tao, A., Anandkumar, A., and Catanzaro, B. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13588 2021.
Chen, K., Wang, C., and Yang, H. let data talk: data- regularized operator learning theory for inverse problems. arXiv preprini ar arXiv:2310.09855 2023.
Hansen, P. C. Rank-deficient and discrete ill-posed prob- lems: numerical aspects of linear inversion. SIAM, 1998.
Hao, Z., Wang. Z., Su, H., Ying, C., Dong, Y., Liu, S., Cheng, Z., Song, J., and Zhu, J. Gnot: A general neu- ral operator transformer for operator learning. In Inter- national Conference on Machine Learning, pp. 12556- 12569. PMLR, 2023.
Cho, S. W. and Son, H. Physics-informed deep inverse op- erator networks for solving pde inverse problems. arXiv preprint ar arXiv:2412.0316.1 2024.
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T.. Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. Rethinking attention with performers. arXiv preprint ar 2009 / /4794, 2020.
Jafarzadeh, S., Silling, S., Liu, N., Zhang, Z., and Yu, Y. Peridynamic neural operators: A data-driven nonlocal constitutive model for complex material responses. Com- puter Methods Applied Mechanics and Engineering. 425: 425:116914, 2024.
Coorey, G., Figtree, G. A., Fletcher, D. F., Snelson, V. J., Vernon, S. T., Winlaw, D., Grieve, S. M., McEwan, A., Yang, J. Y. H., Qian, P., O' O'Brien, K., Orchard, J., Kim, J., Patel, S., and Redfern, The health digital twin to
Jiang. H., Khoo, Y., and Yang, H. Reinforced inverse seat- tering. arXiv preprint arXiv:2206. 04 186 2022.
10

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
Kacham, P., Mirrokni, V., and Zhong, Polysketchformer: Fast transformers via sketches for polynomial kemelss. arXiv preprint arXi::2310.01655. 2023.
automatically encoded conservation laws. In Forty-first International Conference on Machine Learning, 2024a.
Liu, N., Li, X., Rajanna, M. R., Reutzel, E. W., Sawyer, B., Rao, P., Lua, J., Phan, N., and Yu, Y. Deep neural operator enabled digital twin modeling for additive manufacturing. arXiv preprint wrxiv:2405.095777 2024b.
Karniadakis, G. E., Kevrekidis, G., Lu, L., Perdikaris, P., Wang, S., and Yang, thysics-informed machine learning, Nature Reviews Physies, (66):422-440,, 2021.
Lai, R.-Y., Li, Q. and Uhlmann, G. Inverse problems for the stationary transport equation in the diffusion scaling. SIAM Journal on Applied Mathematics, 79(6):2340-2358, 2019.
Liu, N., Zhang, L., Gao, T., and Yu, Y. Disentangled repre- sentation learning for parametric partial differential equa- tions. arXiv preprint arXiv:2410.02130 2024c.
Liu, N., Jafarzadeh, S., Lattimer, B. Y., Ni, S., Lua, J., and
Lang, A. and Potthoff. J. Fast simulation of gaussian random fields. Monte Carlo Methods & Applications, 17(3). 2011.
Yu, Y. Harnessing large language models for data-scarce learning of polymer properties. Nature Computational Science, pp. 1-10, 2025.
Lejeune, E. Mechanical mnist: benchmark dataset for mechanical metamodels. Extreme Mechanics Letters, 36: 100659, 2020.
Liu, X., Xu, B., and Zhang, L. Ht-net: Hierarchical trans- former hased operator learning model for multiscale pdes. 2022.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03:: 03485, 2020a.
Lu, F. and Yu, Y. Transformer learns the cross-task prior and regularization for in-context learning. arXiv preprint arXiv:2505. 12138, 2025.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Stuart, A., Bhattacharya, K., and Anandkumar, A. Multipole graph neural operator for parametric partial differential equations. Advances in Neural Information Processing Systems, 33:NeurIPS 2020, 2020b.
Lu, L., Jin, P., and Karniadakis, G. E. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of opera- tors. arXiv preprint arXiv: 1910.03193.. 2019.
Li, Z., Kovachki, N. B., Azizzadenesheli. K., Bhattacharya, K., Stuart, A., and Anandkumar, A. Fourier NeuralOp- crator for Parametric Partial Differential Equations. In International Conference on Learning Representations, 2020c.
Lu, L.,, Jin, P., Pang, G., Zhang, Z.., and Karniadakis, G. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218-22) 2298 2021.
Lu, Y. M., Letey, M. I., Zavatone- Tavatone-Veth, J. A., Maiti, A., and Pehlevan, C. Asymptotic theory of in-context learning by linear attention. arXiv preprint 1rXiv:2405.1175/ 2024.
Li, Z., Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Azizzadenesheli, K., and Anandkumar, A. Physics- informed neural operator for learning partial differential equations. 2021 arXiv preprint arXiv:2111.03794 2021.
McGreivy, N. and Hakim, A. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid- related partial differential equations. Nature Machine Intelligence, (110):256-122 2024.
Li, Z., Meidani, K., and Farimani, A. B. Transformer for partial differential equations' operator learning. arXiv preprint ar XrXiv:2205. 13671, 2022.
Molinaro, R., Yang, Y., Engquist, B., and Mishra, S. Neural inverse operators for solving pde inverse problems. arXiv preprint arXiv:2301.11166: 2023.
Liu, N., Jafarzadeh, S., and Yu, Y. Domain agnostic fourier neural operators. Advances in Neural Information Pro- cessing Systems, 36, 2023a.
Liu, N., Yu, Y. You, H., and Tatikola, N. Ino: Invariant neu- ral operators for learning complex physical systems with momentum conservation. In International Conference Artificial Intelligence and Statistics, pp. 6822-6833 PMLR, 2023b.
Obmann, D., Schwab, J., and Haltmeier, M. Deep synthe- sis regularization of inverse problems. arXiv preprint arXiv:2002.00155, 2020.
Oko, K., Song, Y., Suzuki, T., and Wu, D. Pretrained trans- former efficiently lears cow-dimensional target functions in-context. Advances in Neural Information Processing Systems,
Liu, N., Fan, Y., Zeng. X., Klöwer, M., Zhang, L., and Yu, Harnessing the power of neural operators with
2024.

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
Ong, Y. Z.., Shen, Z., and Yang, H. IAE-Net: integral autoencoders for diseretizationninaaraae learning. arXiv preprint urXiv:2203.05142, 2022.
You, H., Zhang, Q., Ross, C. J., Lee, C.-H., and Yu, Y. Learning deep implicit fourier neural operators (IFNOs) with applications to heterogeneous material modeling. Computer Methods in Applied Mechanics and Engineer- ing, 398: 15296, 2022b. doi: https://doi.org/10.100000 cma.2022.11532 115296.
Pfau, D., Spencer, J. S., Matthews, A. G., and Foulkes, W. M. C. Ab initio solution of the many-electron schrödinger equation with deep neural networks. Physical Review Research, 2(3):033429, 2020.
Yu, Y., Liu, N., Lu, F., Gao, T., Jafarzadeh, S., and Silling, S. Nonlocal attention operator: Materializing hidden knowl- edge towards interpretable physies discovery. Advances in Neural Information Processing Systems, 37, 2024.
Sun, J., Liu, Y., Zhang, Z., and Schaeffer, H. Towards foundation model for partial differential equation: Multi- operator learning and extrapolation. arXiv preprint arXiv:2404. 12355, 2024.
Zhang, L., Han, J., Wang, H., Car, R., and Weinan, E. Deep potential molecular dynamics: scalable model with the accuracy of quantum mechanics. Physical Review Letters, (20(14):14300) 2018.
Uhlmann, G. Electrical impedance tomography and calderón's problem. Inverse problems, 25312); 2009.
Zhang, Z. Modno: Multi operator learning with distributed neural operators. arXiv preprint arXiv:2404 240,::2400402222,,, 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.., Gomez, A. N., Kaiser, t., and Polosukhin, I. At- tention all you need. Advances in neural information processing systems, 30, 2017.
Vladymyrov, M., Von Oswald, J., Sandler, M., and Ge, R. Linear transformers are versatile in-context learners. Advances in Neural Information Processing Systems, 37: 48784-48809, 2024.
Wang, J., Tian, X., Zhang, Z., Silling, S., Jafarzadeh, S., and Yu, Y. Monotone peridynamic neural operator for nonlinear material modeling with conditionally unique solutions. arXiv preprint arXiv:2505.0166) 2025.
Xu, Z.-Q.J, J., Zhang, Y.. Luo, T., Xiao, Y., and Ma, Fre- quency principle: Fourier analysis sheds light deep neural networks. arXiv preprint arXiv:1901.06523. 2019.
Yang, L. and Osher, J. Pde generalization of in-context operator networks: A study 1d scalar nonlinear con- servation laws. arXiv preprint arXiv:2401.07335 2024.
Yang, L., Meng, X., and Karniadakis, G. E. B-pinns: Bayesian physics-informet neural networks for forward and inverse pde problems with noisy data. Journal of Computational Physics, 425:109913, 2021.
Ye, Z., Huang, X., Chen, L., Liu, H., Wang, Z., and Dong, B. Pdeformer: Towards a foundation model for one- dimensional partial differential equations. arXiv preprint arXiv:2402.1 12652, 2024.
Yilmaz, Ö. Seismic data analysis, volume 1. Society of exploration geophysicists Tulsa, 2001.
You, H., Yu, Y., D'Elia, M., Gao, T., and Silling, S. Non- local kernel network (NKN): A stable and resolution- independent deep neural network. Journal of Compu- tational Physics, pp. arXiv preprint arXiv:2201.02217,, 2022a.
12

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
A. Data Generation and Additional Results
A.1. Experiment Darcy Flow
We generate synthetic data based on Darcy flow in a square domain Q = [0, 1]2 subjected to Dirichlet boundary conditions. The problem described by the equation: - -Vbb())Vp(=)) = g(x) with p(a)) = 0 on all boundaries. This equation models diffusion in heterogeneous fields, such as subsurface water flow in porous media, where the heterogeneity is represented by the location-dependent conductivity b(x). Here, p(æ) is the source term, and g(x) is the hydraulic height (the solution). For each sample, we solve the equation on both 21 x 21 and 41 x 41 grids using an in-house finite difference code. We consider 100 random microstructures, each consisting of two distinct phases. In particular, we sample from a Gaussian random field with zero mean and covariance operator (-A + 5 for the in-distribution (ID) scenario, and (-A + 55)-1 for the nuttof-distribution (OOD) scenario. For each sample ç. we divide the square domain into two subdomains, with conductivities b(x) = 12 when {(x) < 0 and b(x) = 3 when {(x) 2 O. Additionally, we generate 100 distinct source functions g(x) using Gaussian random field generator with zero mean and covariance operator (-A + 53)-1 for the solve the Darey problem with all 100 source terms, yielding a dataset of N 10, 000 function pairs in the form of {x:(2)), 9: (z,))) where j = 1, ... ,441 for the 21 x 21 mesh and 1, 2, 1681 for the 41 x 41 mesh, with 2; denoting the discretization points on the square domain.
ID case, and (- A + 5
4 for the OOD case. For each microstructure,
In operator learning, we note that the permutation of function pairs in each sample should not alter the learned kernel, i.e.,
K[u fi-a] = K [um(1:d) fo(1:))
(12)
where a is the permutation operator. To address this, we augment the training dataset by permuting the function pairs in each task. Specifically, with 100 microstructures (tasks) and 100 function pairs per task, we randomly permute the function pairs and take 100 different permutations per task. This process generates total of 10,000 samples (9,000 for training and 1,000 for testing) in the form of (==".00- f(.100) )/
We demonstrate the scalability of NIPS using larger discretizations (e.g., 121 x 121 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 Table 6, where "x" indicates that the method exceeds memory limits on a single NVIDIA A 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 6.5k tokens, NIPS can easily scale up 15k tokens on single A 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.
Table 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- formuulteed linear attention are reported and separated by "P". "x" denotes exceeding memory limits on single NVIDIA A10 GPU with 40 GB memory.
# tokens
441
1,681
3,721
6,561
10,201
14,641
NIPS
2.5
8,4
17.1
25.1
45.6
65.66
Per-epoch runtime (s)
NAO
7.8/2.5
58.8/11,8
x/20.9
x/x
x/x
X/x
NIPS 0.46
1.68
3.68
6.46
10.04
14.39
Peak memory usage (GB)
NAO
2.47/0.66
32.61/5.69
x/25.53
x/x
X/x
x/x
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 Table 7. Consistent with the previous findings, NIPS outperforms the best baseline by 28.9%.
Table 7. Test errors and number of trainable parameters for the Darcy flow problem. Bold numbers highlight the best method.
Model
NIPS
NAO
AFNO
Test error
2.31% % 10.03%
3.25%
52.92%%-10.72%
13

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
A.2. Experiment 2 Mechanical MNIST benchmark
Mechanical MNIST is benchmark dataset of heterogeneous materials undergoing large deformation, modeled by Neo- Hookean material with varying modulus derived from MNIST bitmap images (Lejeune, 2020). contains 70,000 material specimens, each governed by the Neo-Hookean model with modulus varying according to the MNIST images. Figure 3 illustrates samples from the MMNIST dataset, including the underlying microstructure, randomly selected loading fields, and the corresponding displacement fields in the two in-plane directions.
(a) MMNIST material microstructure
O
2
3
4
5
6
7
9
(b) Exemplar loading and the corresponding displacement field x direction
(c) Exemplar loading and the corresponding displacement field in y direction
Figure 3. Illustration of exemplar MMNIST samples in experiment (a): material parameter field corresponding different b. (b): displacement fields (second row) и. corresponding to the same loading field (first row) fa- (c): displacement fields (second row) corresponding to the same loading field (first row) fy-
A.3. Experiment Synthetic Tissue Learning
We generate synthetic tissue data by sampling the fiber orientation distribution for each sample, which consists of two segments with orientations a and o2 on each side, respectively, separated by line passing through the center of square domain. The values of æ1 and a2 are independently sampled from uniform distribution over 2/ (0, 2r). and the centerline's rotation is sampled from (0, x]. We generate 500 material sets, each containing 100 loading/displaeement pairs, and divide these into training and test sets with 450/50 split. The loading in this dataset is taken as the body load, f()). Each instance is generated the restriction of a 2D random field, ((z) 11 F(T)) where T(æ:) is a Gaussian white noise random field on ?R. 7 (wf + w represents correlation function, and w1 and we are the wave numbers directions, respectively. The operators F and F-1 denote the Fourier transform and its inverse, respectively. This random field is expected to have a zero mean and covariance operator C = (-A)-22) with A being the Laplacian under periodic boundary conditions on [0, 2)2 and we then restrict to Ω. For details on Gaussian random field sample generation, we refer readers to Lang & Potthoff (2011). Finally, for each sampled loading field f. and microstructure field b(x), we solve for the displacement field u, (x:) on the entire domain.
and y
14

---

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Diseovery
B. Implementation Details and Further Discussion on Baselines
The configuration of each baseline detailed below, where the parameter choice of each model is selected by tuning the number of layers and the width (channel dimension), ensuring that the total number of parameters remains within the same order of magnitude. We follow the two key principles in McGreivy & Hakim (2024) for fair comparisons: (1) evaluating models at either equal accuracy or equal runtime, and (2) comparing against an efficient numerical method. First, emphasize that our comparisons focus 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 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, experimental setup adheres the principles outlined in McGreivy & Hakim (2024).
NAO: We use a 4-layer NAO model with LayerNorm applied across both the token and projection dimensions in the first layer, but only across the projection dimension subsequent layers. Both the original quadratic attention and the reformulated linear attention tested. We parameterize the kernel network P.P. and W.P.J with 3-layer MLP with hidden dimensions (32, 64) and LeakyRel.U activation functions.
NAO-f: The NAO-f model follows the same configuration as NAO, except that LayerNorm is applied across both the token and projection dimensions in all layers.
NAO-WP:
The NAO-W' model follows the same configuration as NAO, except that a learnable projection weight of size [2 X Thokens: X Riokens] is added to each iterative layer.
We closely follow the setup in Guibas et al. (2021) and stack 4 AFNO layers together to form the final AFNO model. For each AFNO layer, the number of blocks is set to 1, and the channel dimension size is set to the embedding size with the hidden size factor equal to 3. All other parameters such as the sparsity threshold and the hard thresholding fraction are set to the default values of 0.01 and 1, respectively.
AFNO:
Note that NIPS is conceptually related to the Performer (Choromanski et al., 2020), which introduces kernel-based approx- imations for efficient self-attention. Performer replaces the standard softmax attention with positive orthogonal random features, enabling linear scaling with respect to sequence length. While NIPS similarly aims to mitigate the quadratic complexity of self-attention, NIPS achieves this through a Fourier-domain reformulation that leverages convolutional structure. Unlike Performer's 's randomized feature mappings, NIPS deterministically decomposes attention using spectral rep- resentations, making particularly well-suited for physics-based PDE learning. Additionally, NIPS is also related Graph Neural Operators (GNOs) (Li et al., 2020a), which provide graph-based formulation for learning nonlocal interactions in operator learning tasks. GNOs use graph message passing to propagate information over irregular domains, effectively capturing spatial dependencies. In contrast, NIPS operates in the Fourier domain, leveraging spectral representations to encode long-range dependencies in continuous and structured manner. While graph-based architectures such as GNOs and NAO can learn operators on discretized graph structures, our Fourier-based formulation provides a leveraged efficiency on structured grids.
15


Reviews:
[{"review_id": "J2hAeSWZed", "reviewer": "Reviewer_JpKF", "summary": "The paper introduces Neural Interpretable PDEs (NIPS), an attention-based neural operator architecture for solving forward and inverse PDE problems. NIPS utilizes a learnable kernel network to optimize efficiency in Fourier transform interactions by mitigating the computation and storage of the pairwise interactions. Improved upon Nonlocal Attention Operators (NAO), the author incorporates linear attention to improve the scalability and interpretability of large physical systems learning. The authors demonstrate that NIPS surpasses NAO and other baselines across multiple physics modeling benchmarks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments . Our response:\n\n**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.\n\nIn 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.\n\nTable 1: Test errors, kernel errors, and microstructure errors for the Darcy flow problem. Bold numbers highlight the best method.\n|Data Setting | Case | Model | \\#param | Test error | Kernel error | Microstructure error |\n| :------------- | :-----------: | :-----------: | :-----------: | :-----------: | :------------- | :-----------: |\n|ID | No noise | NIPS | 327,744 | **4.09\\%** | **9.24\\%** | 7.92\\% |\n| | | NAO | 331,554 | 4.15\\% | 10.35\\% | **7.09\\%** |\n|OOD, Scenario 1 | No noise | NIPS | 327,744 | **3.40\\%** | **10.63\\%** | **11.23\\%** |\n| | | NAO |331,554  | 6.86\\% | 14.46\\% | 20.37\\% |\n|OOD, Scenario 2|No noise | NIPS | 327,744 | **4.98\\%** | **9.68\\%** | **16.06\\%** |\n| | | NAO |331,554  | 5.57\\% | 10.94\\% | 20.30\\% |\n|ID|Noise $\\epsilon\\sim\\mathcal{N}(0,0.01^2)$ | NIPS | 327,744 | 4.40\\% | 9.47\\% | 7.97\\% |\n|OOD, Scenario 1|Noise $\\epsilon\\sim\\mathcal{N}(0,0.01^2)$ | NIPS | 327,744 | 3.88\\% | 11.14\\% | 13.63\\% |\n|OOD, Scenario 2|Noise $\\epsilon\\sim\\mathcal{N}(0,0.01^2)$ | NIPS | 327,744 | 5.65\\% | 10.21\\% | 17.36\\% |\n|ID|Noise $\\epsilon\\sim\\mathcal{N}(0,0.1^2)$ | NIPS | 327,744 | 9.98\\% | 21.77\\% | 14.84\\% |\n|OOD, Scenario 1| Noise $\\epsilon\\sim\\mathcal{N}(0,0.1^2)$ | NIPS | 327,744 | 3.84\\% | 21.08\\% | 19.87\\% |\n|OOD, Scenario 2| Noise $\\epsilon\\sim\\mathcal{N}(0,0.1^2)$ | NIPS | 327,744 | 11.67\\% | 22.77\\% | 26.54\\% |\n\n**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:\n\\begin{equation}\n    \\widetilde{g}(x) = g(x) + \\epsilon(x), \\quad \\epsilon \\sim \\mathcal{N}(0, \\sigma^2),\n\\end{equation}\nwhere $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.\n\n[1] Z Li. Fourier neural operator with learned deformations for PDEs on general geometries. JMLR, 2023\n\n[2] J Gupta. Towards multi-spatiotemporal-scale generalized PDE modeling. TMLR, 2022\n\n**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.\n\n**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.", "claims_and_evidence": "The author claims efficiency in Fourier Transform interaction, interpretability and scalability based on the improved attention structure, and generalizability on zero-shot unseen physical systems. The numerical results in Table 1 provide evidence that NIPS has lower test errors and faster training times in the Darcy Flow problem compared to the baselines. Interpretability is primarily demonstrated through visualizations in the Darcy Flow section. While these visualizations provide a qualitative explanation, it is suggested to extend them to other tested datasets or incorporate formalized metrics, as interpretability is a relatively subjective term.", "methods_and_evaluation_criteria": "The benchmark datasets utilized are 2D Darcy Flow, Mechanical MNIST, and Synthetic Tissue Learning. Unlimited to fluid dynamics, the dataset also incorporates a material and biomechanics dataset. The evaluation metric utilized is mainly relative mean squared error (rMSE), which is a standardized metric in evaluating prediction performances.", "theoretical_claims": "The theoretical claims in Section 3.2 and Section 3.3 introduce the expression of attention mechanism-based kernel map and the reformulation of Fourier-domain kernel to enhance the efficiency. Specifically, the reformulation and the assessment of Big O in Section 3.3 looks correct to me.", "experimental_designs_or_analyses": "The experiments confirm NIPS' capability in forward and inverse PDE solving as well as zero-shot generalization. Some vital baseline for this method, such as NAO and NAO-relatives, are included. Additionally, multiple cases of token and layer numbers are discussed. One suggestion is to include real-world PDE data to further demonstrate NIPS' practical applicability, such as airfoil flow and climate data. They could provide additional insights into how well NIPS generalizes on noisy data, which is beyond synthetic PDE problems.", "relation_to_broader_scientific_literature": "This paper proposes a useful method that is closely related to multiple facets in PDE learning, including Neural Operators (NO) and Fourier Neural Operators (FNO), its prior work Nonlocal Attention Operators (NAO). It is also related to both forward and inverse PDE solving in efficiency and effectiveness considerations.", "essential_references_not_discussed": "The author discusses the majority of essential references.", "other_strengths_and_weaknesses": "Clarity: The paper flows smoothly, making readers easy to follow.", "other_comments_or_suggestions": "Suggestions are mentioned in the earlier sections.", "questions_for_authors": "Question 1: How would interpretability vary across different zero-shot conditions?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZwKD8hhou1", "reviewer": "Reviewer_yXYz", "summary": "This article introduces an operator learning method that computes the solution to a PDE with a data-dependent kernel called NIPS. The kernel comes from a linear attention mechanism performs in Fourier space, which improves the efficiency of the method. The method can also be used for the inverse problem to discover the governing parameters of a PDE given the solution. This method is tested Darcy flow, mechanical MNIST, and synthetic tissue models. The methods is claimed to be accurate, efficient, and interpretable.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments and questions. Our response:\n\n**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.\n\nTable 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.\n|                         | \\# tokens |     441  | 1,681 | 3,721  | 6,561  | 10,201 | 14,641 |\n|-------------------------|-----------:|--------:|--------:|--------:|--------:|--------:|--------:|\n| **Per-epoch runtime (s)** | **NIPS**  |       2.5    | 8.4    | 17.1   | 25.1   | 45.6   | 65.66  |\n|                         | **NAO**   | 7.8/2.5| 58.8/11.8| x/20.9| x/x    | x/x    | x/x    |\n| **Peak memory usage (GB)** | **NIPS**  | 0.46   | 1.68   | 3.68   | 6.46   | 10.04  | 14.39  |\n|                         | **NAO**   | 2.47/0.66 | 32.61/5.69 | x/25.53 | x/x   | x/x    | x/x    |\n\n**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.", "claims_and_evidence": "The claim of efficiency is not substantiated clearly. The computational domains of the computer experiments remain fairly small. All discretization are under 2000 points which does not allow for large computational domains. Larger domains need to be studied to illustrate the efficiency and scalability of the method.", "methods_and_evaluation_criteria": "The chosen applications are also weak baselines [1]. More computational experiments are needed to benchmark the method. \n\nMcGreivy, Nick, and Ammar Hakim. \"Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations.\" Nature Machine Intelligence 6.10 (2024): 1256-1269.", "theoretical_claims": "This is an empirical study.", "experimental_designs_or_analyses": "See concern above.", "relation_to_broader_scientific_literature": "This work is an extension to Nonlocal Attention Operators with linear attention in the Fourier domain.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ng8aJZUezW", "reviewer": "Reviewer_JrVi", "summary": "This paper introduces Neural Interpretable PDEs (NIPS), a novel neural operator architecture that enhances both predictive accuracy and computational efficiency in modeling complex physical systems. NIPS builds upon Nonlocal Attention Operators (NAO) by employing a linear attention mechanism combined with a learnable kernel network that functions as a channel-independent convolution in Fourier space. This design eliminates the need to explicitly compute and store large pairwise interactions, effectively amortizing the cost of spatial interactions through the Fourier transform.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments. Our response:\n\n**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\\%.\n\nTable 1: Test errors and number of trainable parameters for the Darcy flow problem. Bold numbers highlight the best method.\n|Model | NIPS | NAO | AFNO |\n| :------------- | :-----------: | :-----------: | :-----------: |\n|Test error | **2.31\\%**$\\pm$**0.03\\%** | 3.25\\%$\\pm$0.14\\% | 52.92\\%$\\pm$0.72\\% |\n\n**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.\n\n**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.\n\n**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.\n\nTable 2: Test, kernel, and microstructure errors for Darcy flow. Bold numbers highlight the best method.\n|Data Setting | Case | Model | \\#param | Test error | Kernel error | Microstructure error |\n| :------------- | :-----------: | :-----------: | :-----------: | :-----------: | :------------- | :-----------: |\n|ID | No noise | NIPS | 327,744 | **4.09\\%** | **9.24\\%** | 7.92\\% |\n| | | NAO | 331,554 | 4.15\\% | 10.35\\% | **7.09\\%** |\n|OOD, Scenario 1 | No noise | NIPS | 327,744 | **3.40\\%** | **10.63\\%** | **11.23\\%** |\n| | | NAO |331,554  | 6.86\\% | 14.46\\% | 20.37\\% |\n|OOD, Scenario 2|No noise | NIPS | 327,744 | **4.98\\%** | **9.68\\%** | **16.06\\%** |\n| | | NAO |331,554  | 5.57\\% | 10.94\\% | 20.30\\% |\n\n**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.\n\n**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.", "claims_and_evidence": "The claims made in this submission are generally well-supported by empirical evidence and theoretical analysis. The authors provide comprehensive quantitative results across three distinct experimental setups (Darcy flow, MMNIST, and synthetic tissue learning), consistently demonstrating NIPS's superior performance compared to baseline models.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in the paper are well-aligned with the challenges of learning interpretable operators for physical systems. The NIPS architecture logically addresses the limitations of existing approaches by combining Fourier techniques with linear attention, which is appropriate for handling PDE operators efficiently while maintaining interpretability.", "theoretical_claims": "The paper does not present formal mathematical proofs in the traditional sense with explicitly stated theorems and detailed proof steps", "experimental_designs_or_analyses": "The experimental designs in this paper are generally sound and well-executed. The Darcy flow experiments (Section 4.1) use appropriate synthetic data generation procedures with clear controls across model variants while maintaining comparable parameter counts.\n\nThe parametric studies in Table 2 systematically explore the effects of random permutations and projection dimensions with sufficient sampling to establish trends. One limitation is that statistical significance of the performance differences is not explicitly established through multiple runs with different random seeds.", "relation_to_broader_scientific_literature": "The key contributions of this paper build upon several important strands of research in the computational physics and machine learning literature. NIPS integrates two significant developments in neural operators: the Fourier Neural Operator (FNO) framework introduced by Li et al. (2020c), which leverages spectral methods for efficient computation, and the Nonlocal Attention Operators (NAO) recently proposed by Yu et al. (2024), which focus on interpretability through attention mechanisms.", "essential_references_not_discussed": "While the authors cite linear attention work by Cao (2021), they do not discuss Performer (Choromanski et al., 2020, NeurIPS) which introduced kernel-based approximations for efficient attention that are conceptually related to their Fourier-domain approach. Second, the paper could benefit from citing Graph Neural Operators (GNOs) by Li et al. (2020, NeurIPS), as these provide an alternative perspective on handling non-local interactions in operator learning that would help contextualize NIPS's approach.", "other_strengths_and_weaknesses": "Pros:\n\nThe harmonization of Fourier insights with attention mechanisms represents a creative combination of existing ideas that yields tangible performance improvements. The dual capability of NIPS to solve both forward and inverse problems within a unified framework is especially significant, as it addresses a fundamental challenge in computational physics. \n\nCons:\n\nWhile empirical results are strong, a more rigorous mathematical foundation would strengthen the paper's contribution. Additionally, the experimental evaluation could benefit from a more detailed analysis of failure cases or boundary conditions where the method might underperform. The paper would also be strengthened by exploring more diverse or complex physical systems beyond the three test cases presented, particularly systems with higher dimensionality or stronger nonlinearities.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. Your experiments show AFNO consistently achieving error rates exceeding 50% across all scenarios, which seems unusually poor for an established method. Could you clarify whether this reflects a fundamental limitation of AFNO for these specific tasks, or if there might be implementation factors affecting its performance?\n\n2. While you demonstrate strong zero-shot generalization capabilities within similar PDE systems, how would NIPS perform when transferring to fundamentally different classes of PDEs? For instance, could a model trained on diffusion equations generalize to wave equations or advection-dominated systems?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: NIPS presents a meaningful advance in neural operators, balancing efficiency, interpretability, and generalization. The rebuttal convincingly resolved scalability concerns and contextualized limitations (e.g., symmetry assumptions). Given the strong empirical results, clear technical contribution, and alignment with ICML’s scope, acceptance is warranted. Future work could explore broader PDE classes and theoretical guarantees.
Decision: Accept (poster)