Nonlocal Attention Operator: Materializing Hidden Knowledge Towards Interpretable Physics Discovery
Authors: Yue Yu
Ning Liu
Fei Lu
Tian Gao
Siavash Jafarzadeh
Stewart A Silling
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Foundation Model
Neural Operators
Inverse PDE Problems
Physical Modeling
Primary area: machine_learning_for_physical_sciences

Abstract: Despite recent popularity of attention-based neural architectures in core AI fields like natural language processing (NLP) and computer vision (CV), their potential in modeling complex physical systems remains under-explored. Learning problems in physical systems are often characterized as discovering operators that map between function spaces based on a few instances of function pairs. This task frequently presents a severely ill-posed PDE inverse problem. In this work, we propose a novel neural operator architecture based on the attention mechanism, which we coin Nonlocal Attention Operator (NAO), and explore its capability towards developing a foundation physical model. In particular, we show that the attention mechanism is equivalent to a double integral operator that enables nonlocal interactions among spatial tokens, with a data-dependent kernel characterizing the inverse mapping from data to the hidden parameter field of the underlying operator. As such, the attention mechanism extracts global prior information from training data generated by multiple systems, and suggests the exploratory space in the form of a nonlinear kernel map. Consequently, NAO can address ill-posedness and rank deficiency in inverse PDE problems by encoding regularization and achieving generalizability. Lastly, we empirically demonstrate the advantages of NAO over baseline neural models in terms of the generalizability to unseen data resolutions and system states. Our work not only suggests a novel neural operator architecture for learning an interpretable foundation model of physical systems, but also offers a new perspective towards understanding the attention mechanism.

Nonlocal Attention Operator: Materializing Hidden Knowledge Towards Interpretable Physics Discovery
Yue Yu" Department of Mathematics, Lehigh University, Bethlehem, PA 18015, USA yuy2140lehigh.
Ning Liu Global Engineering and Materials, Inc., Princeton, NJ 08540, USA
Fei Lu Department of Mathematics, Johns Hopkins University, Baltimore, MD 21218, USA
Tian Gao IBM Research, Yorktown Heights, NY 10598, USA
Siavash Jafarzadeh Department of Mathematics, Lehigh University, Bethlehem, PA 18015, USA
Stewart Silling Center for Computing Research, Sandia National Laboratories, Albuquerque, NM 87123, USA
Abstract
Despite the recent popularity of attention-basee based neural architectures in Al fields like natural language processing (NLP) and computer vision (CV), their po- tential in modeling complex physical systems remains under-explored. Learning problems in physical systems are often characterized as discovering operators that map between function spaces based on a few instances of function pairs. This task frequently presents a severely ill-posed PDE inverse problem. In this work, we propose a novel neural operator architecture based on the attention mechanism, which we coin Nonlocal Attention Operator (NAO), and explore its capability to- wards developing a foundation physical model. In particular, we show that the tention mechanism is equivalent to double integral operator that enables nonlocal interactions among spatial tokens, with data- data-dependent kernel characterizing the inverse mapping from data to the hidden parameter field of the underlying operator. As such, the attention mechanism extracts global prior information from training data generated by multiple systems, and suggests the exploratory space in the form a nonlinear kernel map. Consequently, NAO can address ill-posedness and rank deficiency in inverse PDE problems by encoding regularization and achieving gen- cralizability. We empirically demonstrate the advantages NAO over baseline neural models in terms of generalizability to unseen data resolutions and system states. Our work not only suggests a novel neural operator architecture for learning interpretable foundation models of physical systems, but also offers new perspec- tive towards understanding the attention mechanism. Our code and data accompa- nying this paper available at
124/NAA.
Introduction
The interpretability of machine learning (ML.) models has become increasingly important from the security and robustness standpoints [Rudin et all 2022, Molnar, 2020]. This is particularly true in physics modeling problems that can affect human lives, where not only the accuracy but also the transparency of data-driven models are essential in making decisions [Coorey et ar, 2022, herrni and Willcox, [0ze)). Nevertheless, it remains challenging to discover the underlying physical system and the governing mechanism from data. Taking the material modeling task for instance, given that only the deformation field is observable, the goal of discovering the underlying material parame-
Corresponding Author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

ter field and mechanism presents an ill-posed unsupervised learning task. That means, even if an ML model can serve as good surrogate to predict the corresponding loading field from given deformation field, its inference of the material parameters can still drastically deteriorate.
To discover interpretable mechanism for phys- ical systems, a major challenge is to infer the governing laws of these systems that are often
Training
(U, f)) (u, E2)
Wos
Kernel space of identiflabiiiy
high- or infinite-dimensional. from data that are
Integrate
comprised of discrete measurements of continu- ous functions. Therefore, a data-driven surrogate model needs learn not only the mapping be- tween input and oulput function pairs, but also the mapping from given function pairs to the hidden state. From the PDE-based modeling standpoint, learning a surrogate model corresponds to a for- ward problem, whereas inferring the underlying mechanism is an inverse problem. The latter is generally an enduring ill-posed problem, especially when the measurements are scarce. Unfortunately, such an ill-posedness issue may become even more severe in neural network models, due to the inherent bias of neural network approximations [Xu 2019) To tackle this challenge, many deep learning methods have recently been proposed as inverse PDE solvers [Pan and Ying, 2023, Molinaro et all 2026, grang et all, 2022, Chen cl all 2023]. The central idea is to incorporate prior information into the learning scheme, in the form of governing PDEs Yang et all 2021. Li all 2021I,, regularizers [Dittmer et all, 20200, Obmann et all 2020. Ding all 2022 Chen et all [ additional operator structures [Uhlmanni 2009, Lai atal. 2019, Yilmaz 2ND1I However, such prior information is often either unavailable or problem-
2 4
(u. f))
W
W
Yuu
Kemee (mechanism)
Downstream application
K K[4 4aa
Figure I: Illustration of NAO's architecture.
specific in complex systems. As a result, these methods can only solve the inverse problem for particular system, and one has to start from scratch when the system varies (e.g., when the material of the specimen undergoes degradation in material modeling task).
In this work, we propose Nonlocal Attention Operator (NAO), a novel attention based neural op- erator architecture to simultaneously solve both forward and inverse modeling problems. Neural operators (NOs) [L et all 2020%lc] learn mappings between infinite-dimensional function spaces in the form of integral operators, hence they provide promising tools for the discovery of continuum physical laws by manifesting the mapping between spatial and/or spatiotemporal data; see You all
[2022], Liu all [20221 2023). Ong et all [2022]] Cao [2021]] Lu all [20119 2021], Goswami al [202]. Gupta et al, [||2 and references therein. However, most NOs focus on providing an efficient surrogate for the underlying physical system as forward solver. They are often em- ployed as black-box universal approximators but lack interpretability of the underlying physical laws. In contrast, the key innovation of NAO that it introduces a kernel map based on the atten- tion mechanism for simultaneous learning of the operator and the kernel map. As such. the kernel map automatically infers the context of the underlying physical system an unsupervised manner. Intuitively, the attention mechanism extracts hidden knowledge from multiple systems by providing a function space of identifiability for the kernels, which acts as an automatic data-driven regularizer and endows the learned model's generalizability to new and unseen system states.
In this context, NAO learns kernel map using the attention mechanism and simultaneously solves both the forward and inverse problems. The kernel map, whose parameters extract the global infor- mation about the kernel from multiple systems, efficiently infers aesouttioninvarriat kernels from new datasets. As a consequence, NAO can achieve interpretability of the nonlocal operator and enable the discovery of hidden physical laws. Our key contributions include:
We bridge the divide between inverse PDE modeling and physics discovery tasks, and present a method to simultaneously perform physics modeling (forward PDE) and mecha- nism discovery (inverse PDE).
We propose a novel neural operator architecture NAO, based on the principle of contextual discovery from input/output function pairs through kernel map constructed from multiple physical systems. As such, NAO is generalizable to new and unseen physical systems, and offers meaningful physical interpretation through the discovered kernel.
We provide theoretical analysis to show that the attention mechanism in NAO acts to pro- vide the space of identifiability for the kernels from the training data, which reveals its ability to resolve ill- ill-posed inverse PDE problems.

---

We conduct experiments on pero-shot learning to new and unseen physical systems, demon- strating the generalizability of NAO in both forward and inverse PDE problems.
2 Background and related work
Our work resides at the intersection of operator learning, attention-based models, and forward and inverse problems of PDEs. The ultimate goal is to model multiple physical systems from data while simultaneously discovering the hidden mechanism.
Neural operator for hidden physics. Learning complex physical systems directly from data is ubiq-
uitous in scientific and engincering applications [Ghaboussi et all [1998. al 20224 Ghaboussi
aLa 1991 Carleo et all, 2019, Karmadakis et al, 2021, Zhang et all 2018. Ca cf all 2022, Plan
et all 20XL, He et all 2021, Besnard et all 2006]. In many applications, the underlying govern- ing laws are unknown, hidden in data to be revealed by physical models. Ideally, these models should be interpretable for domain experts, who then use these models to make further predic- tions and expand the understanding of the target physical system. Also, these models should be resolution-invariant. Neural operators are designed learn mappings between infinite-dimensionall function spaces [[ el all 2020a Ric, You et al. 2022 Ong et al, 20222 Cad, 202I1 Lu et all 2OLL. 2021, Goswami et all 2022. Gupta et all, 2021]. As result, NOs provide promising tool for the discovery of physical laws by manifesting the mapping between spatial and/or spatio-temporal data.
Forward and inverse PDE problems. Most current NOs focus providing an efficient surrogate for the underlying physical system as a forward PDE solver. They are often employed as black- box universal approximators without interpretabilitty of the underlying physical laws. Conversely, several deep learning methods have been proposed as inverse PDE solvers [Fan and Ying, 2023. Molinaro all 2023. Jiang et all 2022, Chen et all, 2023], aiming reconstruct the parameters in the PDE from solution data. Compared to the forward problem, the inverse problem is typically more challenging due its ill-posed nature. To tackle the ill-posedness, many NOs incorporate prior information, in the form of governing PDEs [Yang et all 2021. Li el all, 2021I,, regularizers
2IPK' or additional operator structures [Uhlmann 2009. Carctal 2011. Yilmaz 200011 To our knowledge, our NO architecture the first that solves both the forward (prediction) and inverse (discovery) problems simultaneously.
Dittmeretal. 2020, Obmann 2020, Ding all 2022, Chen
Attention mechanism. Since 2017, the attention mechanism has become the backbone of state-of- the-art deep learning models on many core AI tasks like NLP and CV. By calculating the similarity among tokens, the attention mechanism captures long-range dependencies between tokens [[aswann etall 20 Then, the tokens are spatially mixed to obtain the layer output. Based on the choice of mixers, attention-base based models can be divided into three main categories: discrete graph-based atten- tions [Child 201,, Ho et all 2019, Wang et all, 2020, Katharopoulos all M020], MLP-based
attentions [lolstrkhin et all 20241 Louvron et all 2022, Liu ei all 2021], and convolution bbsed atten-
tions [Lee-Thorp et all 2I1IL ctal 2021 Gunbas et all 20211 Nekoozadch all 2123]. While
most attention models focus on discrete mixers, it is proposed in Ci all [2021] Nckoozadai
etall [2023]] [as et all [201%] Cad [202[]] Wei and Zhang [2023]] to frame token mixing as a kernel
integration, with the goal of obtaining predictions independent of the input resolution.
Along the line of PDE-solving tasks, various attention mechanisms have been used to enlarge model capacity. To improve the accuracy of forward PDE solvers, Cao [2021]] removes the softmax nor- malization in the attention mechanism and employs linear attention as a learnable kernel in NOs. Further developments include the Galerkin-type linear attention in an encoder-decoder architecture in OFormer [Lretal ROZA)) a hierarchical transformer for learning multiscale problems [L et all 2022] and heterogeneous normalized attention with a geometric gating mechanism |HHo ct all 2023 to handle multiple input features. In particular, going beyond solving single PDE, the foun- dation model feature of attention mechanisms has been applied towards solving multiple types of PDEs within a specified context in Yang and Osher [0(24)) J. Ye et all [2022]] Sun et all [21)21]] Zhang [2024]] However, none of the existing work discovers hidden physics from data, nor do they discuss the connections between the attention mechanism and the inverse PDE problem.
3 Nonlocal Attention Operator
Consider multiple physical systems that described by class of operators mapping from input functions u E X to output functions f E Y. Our goal is to learn the common physical law, in the form of operators Lx X + Y with system-dependent kernels K:
L =
(I)
3

---

Here X and Y are Banach spaces, denotes an additive noise describing the discrepancy between the ground-trun operator and the optimal surrogate operator, and K is a kernel function representing the nonlocal spatial interaction. As such, the kemel provides the knowledge of its corresponding system, while (I) offers zero-shot prediction model for new and unseen systems.
To formulate the learning, we consider Rerain training datasets from different systems, with each dataset containing d, function pairs (u, f):
1itroin Der = (") 1==
In practice, the data of the input and output functions are spatial mesh )= C Ω C R
The Rerain models with kernels {K" Jrp1 Strain correspond to different material micro-structures or different parametric settings. As a demonstration, we consider models for heterogeneous materials with operators in the form
Cx (x(u](x)= - K ((l.:))((((()dd) x E 2
(3)
where gy) is functional of determined by the operator: for example glu((y) = u(y) in Section 5.3 Our approach extends naturally other forms of operators, such as those with ra- dial interaction kernels in Section 55 and heterogeneous interaction in the form of Lx [=)( = So K (I. y)g (ar, y)dy. Additionally, for simplicity, we consider scalar-valued functions u and and note that the extension to vector-valued functions is trivial.
Remark: Such operator learning problem arises in many applications in forward and inverse PDE-solving problems. The inference of the kernel K is an inverse problem, and the learning of the nonlocal operator is forward problem. When considering a single physical system and taking K in (G) as an input-ing input-nndeeendentt kernel, classical NOs can be obtained for forward PDE-solving tasks [Li MONG Ginbas aall [02T] and governing law learning tasks [Lol 2021. eatarzadch
2024].
Different from existing work, we consider the operator learning across multiple systems.
3.1 Kernel map with attention mechanism
The key ingredient in NAO is kernel map constructed using the attention mechanism. It maps from data pairs to an estimation of the underlying kernel. The kernel map
((us: [( - K[u K[unaa; f1:4; 0]
has parameters 0 estimated from the training dataset (α). As such, it maps from the token (u f1.d)
of the dataset {(u,, f4) kernel estimator, acting as an inverse PDE solver.
major innovation of this kernel map is its dependence on both u and f through their tokens. Thus, approach distinguishes itself from the forward problem-solving NOs in the related work section, where the attention depends only on u.
We first transfer the data {(u,, f.) to tokens (u fi according to the operator in (B) by
Und = (u,...., uz) = (olus((a)))sssaass E Axd
(5)
fia = (f (zx ((((
Nxd.
where d = dy and N = flz, assuming that g|u] has spatial mesh {ir = =xx)/ J
Then, our discrete (L + 1)-lyyer attention model for the inverse PDE problem writes:
Xin x (0) (U F())) (unia; fiaa) 2Nd
x (i) Attn (1-1) (1--));0)x0 (1 (-1 (U F()) 22 15 L
(6)
X x- K[urd, fi:d; 0 Jnr find Nxd
Here, 0: = (w. e dxd,, W.K E dxdd)) and the attention function is
Aun [
Xw, w/ KT x x")
12 2Nx2N
dx
The kernel map is defined as
( 1 ( ( ((ll-1)
K K|u fine; 0 = P.n.
Vdla
+ Pif
(1ol-133 ( (((l))) (((--))
Vile

---

where 0 := {WPx E RNXN, W E RNNN (W, W=i are learnable parameters. Here, we note that de only controls the rank bound when lifting each point-wise feature via (W In the following, we denote WOO = 1 W (WKK), which characterizes the (trainable) inter- action of d input function pair instances.
3.2 Nonlocal Attention Operator in continuum limit
As suggested by Cad [2021]] we take the activation function o as linear operator. Then, noting that the matrix multiplication in (6) is a Riemann sum approximation of an integral (with a full derivation Appendix A), we propose the nonlocal attention operator as the continuum limit of (B)
(z) 00 (z) g(us)(r), (x) f5(x); 9
(a)) ) -/a (E.N)
1 (x)
(v) ( ) dy +
2 (
) ISICL
Lituna (ul())) / KuLaa flua: (2, ())
(8)
in which the integration is approximated by the Riemann sum in our implementation. Here, Ei
g (--) (z) we (v) Ed.em (c) w. (v)) OK
kddh (z.
(x)
TW 115 Du (18 QK Ww, w)
Kluia, fl:d; (x. = Σ P.Pi
1) lws -(gg dz
/ w,v=1
(Ps((, (1
(., vlo! (2)2" de.
We learn the parameters 0 by minimizing the following mean squared error loss function:
E(O)
Terain
Σ EE
Ex Sxtola ( ff) (1( dr.
(9)
Aerain
The performance of the model is evaluated test tasks with new and unseen kernels Drest (1lcol (x)())) based on the following two criteria:
C Cucussr sggres [xx) l [(e0(()) 22
1 d
Operator (forward PDE) Error: Ey
(10)
[|fteer (x)) 12(0)
Kernel (inverse PDE) Error: Ex : K [u dest ftas ;0] - Krest
(2) 120)
(11)
Here, (p) is the empirical measure in the kernel space defined in Lucetal [2023]] Note that are trained using multiple datasets. Intuitively speaking, the attention mechanism helps encode
the prior information from other tasks for learning the nonlocal kernel, and leads estimators significantly better than those using a single dataset. To formally understand this mechanism, we analyze shallow 2-layer NAO in the next section.
Understanding the attention mechanism
To facilitate further understanding of the attention mechanism, we analyze the limit of the two- layer parmmttrepppr kernel in ( and the range of the kernel map, which falls in the space of identifiability for the kernels from the training data. We also connect the kernel map with the regularized estimators. For simplicity, we consider operators of the form
1' ((r)glu((r;2)d z)dr, x eED Ω,
Lx (x(u](x)=
(12)
which is the radial nonlocal kernel in Sec.5.1l
kernel We show that, as the number N and the spatial mesh n,, approach infinity, the limit of the two-layer attention-paree parameterized kernel is double integral. Its proof is in Appendix BTl
4.1 Limit of the two-layer iidentioo-paramettern
5

---

For simplicity, we assume that the dataset in (2) has du = 1 and Perain with uniform mesh {2,]ī=r- We define the tokens by
U1dd = -U1,..... ud) = (r(u)(rk, 2 22)) IisSsSIIsSN E ERNx fiid = (f(r))ssss E R Rixd
(13)
where d nr and {rk Ni is the spatial mesh for K's independent variable r E [0,
Lemma 4.1. Consider the two-layer atlention model in (G)-() with bounded parameters. For each d and N, let {æz}}=1 and {r&] be uniform meshes of the compact sets Q and [0, 5), and let
{Ai} be the resulting uniform partition Sf Assume that glu] in (12) is continuous on [0, 6) x St Then,
lim
lim
K[uiad fiidi on -rk-1)
= f(()) = 1 wPo((rDD ( [olullr', x)) (x) (z. d'' (14)
+ +WP. ( [rO)wak (a, v)olu((T: w) drdy)
where WQK (I,y) = limo-+oo Ei wQw [i,5"]1 A, A, y) is the scaled L x 9)) limit of the parameter matrix Wok = ES Wabi WK and (P =
Nim Ex WP.(k]]] (re
4.2 Space of identifiability for the kernels
For given training dataset, we show that the function space in which the kernels can be identified is the closure of a data-adaptive reproducing kernel Hilbert space (RKHS). This space contains the range of the kernel map and hence provides the ground for analyzing the inverse problem.
Lemma 4.2 (Space of Identifiability). Assume that the training data pairs in ( are sampled from continuous functions {u {u])d du Niravn with a compact support. Then, the function space the loss func- fim=l
tion in (5) has a unique minimizer K(s) = K[ [u]edd Plai (s) is the closure of data-adaptive RKHS HG with a reproducing kernel G determined by the training data:
G(r,s) = lp' (r)) s),
where of is the density of the empirical measure p defined by
1 Divnin
p((r)
Σ
(r,
(15)
Z
1=1
and the function G is defined by G(r, ;= Revaind
Sor glu) x)g[u [w]
dzz
The above space is data-adaptive since the integral kernel & depends on data. It characterizes the information in the training data for estimating the nonlocal kernel K(s) K [uu. f" f (s). In general, the more data, the larger the space is. On the other hand, note that the loss function's minimizer with respect to ((s) is not the kernel map. The minimizer is fixed estimator for the training dataset and does not provide any information for estimating the kernel from another dataset.
Comparison with regularized estimators. The kernel map solves the ill-posed inverse problem using prior information from the training dataset of multiple systems, which is not used in classi- cal inverse problem solvers. To illustrate this mechanism, consider the extreme case of estimating the kernel in the nonlocal operator from a dataset consisting of only a single function pair (u, f). This inverse problem is severely ill-posed because of the small dataset and the need for deconvolu- tion to estimate the kernel. Thus, regularization is necessary, where two main challenges present: (i) the selection of proper regularization with limited prior information, and (ii) the prohibitive computational cost of solving the resulting large linear systems many times.
In contrast, our kernel map K [urd; fiza; (s), with the parameter 0 estimated from the training datasets, acts on the token (uiad; fi:d) of (u, f) to provide an estimator. It passes the prior infor- mation about the kernel from the training dataset to the estimation for new datasets. Importantly,

---

it captures the nonlinear dependence of the estimator on the data (u, f). Computationally, it can be applied directly to multiple new datasets without solving the linear systems. Section B.Z we further show that regularized estimator depends nonlinearly on the data pair (u, f). In particular, similar to Lemma I.I. there is an RKHS determined by the data pair (u, f). The regularized estima- tor suggests that the kernel map can involve a component quadratic in the feature glu), similar to the limit form of the attention model in Lemma
5 Experiments
We assess the performance of NAO on a wide range of physics modeling and discovery datasets. Our evaluation focuses several key aspects: 1) we demonstrate the merits of the continuous and linear attention mechanism, compare the performance with the baseline discrete attention model (de- noted as Discrete-NAO), the softmax attention mechanism (denoted Softmax-NAO), NAO with input on u only (denoted as NAO-u), the convolution- attention mechanism (denoted as AFNO [Guibas et all [021]), and an MI.P-based encoder architecture that maps the datum [ui:d fi:dl di- rectly to a latent kernel (denoted as Autoencoder); 2) we measure the generalizability, in particular, the zero-shot prediction performance modeling a new physical system with unseen governing equations, and across different resolutions; 3) we evaluate the data efficiency-accuracy trade-off in ill-posed inverse PDE learning tasks, as well as the interpretability of the learned kernels. In all experiments, the optimization is performed with the Adam optimizer. To conduct fair comparison across different methods. we tune the hyperparameters, including the learning rates, the decay rates, and the regularization parameters, to minimize the training loss. In all examples, we use 3-layer models, and parameterize the kernel network P.P.a and WPP with a 3-layer ML. with hidden dimensions (32, 64) and LeakyReI.U activation. Experiments are conducted on single NVIDIA GeForce RTX 3090 GPU with 24 GB memory. Additional results and details on data generation and training strategies are provided in Appendix C
5.1 Radial kernel learning
Table 1:
Test errors and the number of trainable parameters for the radial kernel problem, where bold numbers highlight the best methods. The small operator errors and large kernel errors of discrete- NAO highlight the ill-posedness of the inverse problem. NAO overcomes the ill ill-posedness and yields desolution-invarinn estimators.
Case
model
#param
Operator test
Kernel test
OODI
ID
OODI
Discrete-NAO
16526
1.33%
25.81 %
29.02%
28.80%
Softmax-NAO
18843
13.45%
66.06%
67.55%
85.80%
d 302, d,=10
AFNO
19605
22.62%
68.76%
NAO
18843
1.48%
8.10%
5.40% %
10.02 %
NAO-u
18842
13.68%
66.68%
20.46%
74.03%
Autoencoder
16424
12.97%
1041.49%
22.56%
136.79%
d020 de=5
Discrete-NAO
10465
1.63%
15.80%
33.21%
30.39%
NAO
12783
2.34%
9.23%
6.87%
14.62%
d =022, d;=20
Discrete-NAO
28645
1.35%
18.70%
35.49%
30.81%
NAO
30963
1.33% %
9.12%
4.63%
9.14% %
d= 100, d,=10
Discrete-NAO
8446
1.73%
14.92%
34.52%
35.20%
NAO
10763
1.07 %
6.35%
7.41 %
17.02%
50 d,=10
Discrete-NAO
6446
2.29%
10.31%
41.80%
45.30%
NAO
8763
1.56%
7.19%
15.95%
29.47%
d= 30, dg=10
Discrete-NAO
5646
5.60%
11.31%
58.24%
64.23%
NAO
7963
2.94% %
8.04%
22.65%
33.77%
In this example, we consider the learning of nonlocal diffusion operators, in the form:
(x] = [ 7v - x (u(()) u(>)]dy f(x), Vx e Q.
(16)
Unlike a (local) differential operator, this operator depends on the function u nonlocally through the convolution of u(y) u(x). and the operator is characterized by radial kemne Th. It finds broad physical applications describing fracture mechanics [Silling. 2000], anomalous diffusion
behaviors [Bucur et anl 2016), and the homogenization of multiscale systems aill 2020].
7

---

0.4
Ground-truch in-distribution test kernel
1.2
Ground-trath out-of-distrbbttion test kernel
Discrete-NAO: kernel error =55,24%
Discrete-NAO: kemel error 664223%
NAO, same resolution: kernel enror-22.65%
NAO; kemel crror-33.77%
K([x-y))
0.3
NAO, cross resolution: kernel emor- 20%
K((x-())
0.8
0.2
0.6
0.4
0.2
8
X--V
X--yl
Figure 2: Results on radial kernel learning, when learning the test kernel from a small (d 30) number of data pairs: test an ID task (left), and test OOD task (right).
In this context, our goal is to learn the operator L as well as to discover the hidden mechanism, namely the kernel K ju fi.z: 0 (x.9) = Yn
zz). In the form of the operator in (l22), we have (rr) (r) and (r,](r, = u(x + r) + u(x r) (z(x)for E [0.
To generate the training data, we consider 7 sine-type kernels
(o(1y-z) ==
=))) sin - x (o.11 - =l), 17
(17)
Here, 1 denotes task index. We generate 4530 data pairs (g (g" with a fixed resolution Ar 0.0125 for each task, where the loading function Ly [u*] f" is computed by the adaptive Gauss- Kronrod quadrature method. Then, we form a training sample of each task by taking d pairs from this task. When taking the token size d = 302, each task contains 4530
302 15 samples. We consider
two test kernels: one following the same rule of (IZ) with 7j 5 (denoted as the "in-distribution (ID) test" system), and the other following different rule (denoted as the -futtsi-distiiuutiio (OOD) testl" system):
Tood (ly =|) =- -y -111=
xx) =z)) sin xI) po.1 (1y--l) -x)) (18)
Both the operator error (LEC) and the kernel error (II) are provided in Table M While the former measures the error of the learned forward PDE solver (i.c.. learning a physical model), the latter
demonstrates the capability of serving as an inverse PDE solver (i.e., physics discovery).
Ablation study. We first perform ablation study NAO, by comparing its performance with its variants (Discrete-NAO, Softmax-NAO, and NAO-u), AFNO, and Autoencoder, with a fixed token dimension d = 302, query-key feature size dę = 10, and data resolution Ar 0.0125. When comparing the operator errors, both Discrete-NAG and NAO serve as good surrogate models for the ID task with relative errors of 1.33% and 1.48%, respectively, while the other three baselines show > 10% errors. Therefore, we focus more on the comparison between Discrete-NAO and NAO. This gap becomes more pronounced in the OOD task: only NAO is able provide a surrogate of Cyood with < 10% eιror, who outperforms its discrete mixer counterpart by 68. 08.62%% indicating that NAO learns more generalizable mechanism. This argument is further affirmed when comparing the kernel errors, where NAO substantially outperforms all baselines by at least $1.39% in the ID test and 65. 25%1 in the OOD test. This study verifies our analysis Section E NAO learns the kernel map in the space of identifiability, and hence possesses advantages in solving the challenging ill-posed inverse problem. Additionally, we vary the query-key feature size from dę- = 10 da and da 20. Note that de determines the rank bound of WQK, the matrix that characterizes the interaction between different data pairs. Niscrete-NAO again performs well only in approximating the operator for the ID test, while NAO achieves consistent results in both tests and criteria, showing that it has successfully discovered the intrinsic dimension in the kernel space.
Alleviating ill ill-posedness. To further understand NAO's capability as an inverse PDE solver, we reduce the number of data pairs for each sample from d 302 d 30, making more ill- posed as an inverse PDE problem. NAO again outperforms its discrete mixer counterpart in all aspects. Interestingly, the errors in NAO increase almost monotonically, showing its robustness. For Discrete-NAO, the eттor also increases monotonically in the ID operator test, but there exists no consistent pattern other test criteria. Figure shows the learned test kernels both the ID and OOD tasks. It shows that Discrete-NAO learns highly oscillatory kernels, while our continuous NAO only has a discrepancy near I y O. Note that when Jar y/ = 0, we have u(y) u(x) = in the ground-trutl operator (16)) and hence the kernel value at this point does not change the operator
8

---

Figure 3: Kernel visualization in experiment 2, where the kernels correspond the inverse of stiff- ness matrix: ground truth (left), test kernel from Discrete-NAO (middle), kernel from NAO (right).
value. That means, our data provides almost no information at this point. This again verifies our analysis: continuous NAO learns the kernel map structure from small data based on prior knowledge from other task datasets.
Cross-resolution. We test the NAO model trained with Ax 0.0125 on dataset corresponding to Ar = 0.025, and plot the results in Figure 0 Left. The predicted kernel is very similar to the one learned from the same resolution, and the error is also on-par (22.65% versus 20.79%).
5.2 Solution operator learning
We consider the modeling of 2D sub-surface flows through a porous medium with a heterogeneous permeability field. Following the settings in Li [[020]] the high-fidelity synthetic simulation data for this example are described by the Darey flow. Here, the physical domain f2 [0, 1]2 b(xz) is the permeability field, and the Darcy's equation has the form:
V (b(x) g(x), I E Qi
p(z) = 0, æ E 82
(19)
In this context, we aim to learn the solution operator of Darcy's 's equation and compute the pressure field p(æ).. We consider two study scenarios. 1) g + p: each task has fixed microstructure b(x). and our goal is to learn the (linear) solution operator mapping from each loading field to the corresponding solution field p. In this case, the kernel K acts as the Green's function of (I9), and can be approximated by the inverse of the stiffness matrix. 2) b - p. each task has a fixed loading field g(æ), and our goal is to learn the (nonlinear) solution operator mapping from the permeability field b to the corresponding solution field p.
Table 2: Test errors and the number of trainable parameters in solution operator learning.
Case
model Discrete-NAO
#param
Linear Operator: - Nonlinear Operator: b +
d= 20, dg=20 900 samples
161991
8.61% %
10.84%
NAO
89778
8.33%
11.40%
d= 50, d,=40
Discrete-NAO
662163
3.28%
5.61%
9000 samples
NAO
189234
3.19% %
5.28 %
We report the operator learning results in Table D. where NAO slightly outperforms Discrete-NAO in most cases, using only 1/2 or 1/3 the number of trainable parameters. On the other hand, we also verify the kernel learning results by comparing the learned kernels in a test case with the ground- truth inverse of stiffness matrix in Figure G Although both Discrete-NAO and NAO capture the major pattern, the kernel from Discrete-NAO again shows spurious oscillation in regions where the ground-truth kernel has zero value. On the other hand, by exploring integrated knowledge space, the learned kernel from NAO does not exhibit such spurious modes.
kernel map in the
To demonstrate the physical interpretability of the learned kernel, in the first row of Figure E we show the ground-truth microstructure b(x). test loading field instance g(x), and the corresponding solution p(æ). By taking the summation of the kernel strength on each row, one can discover the interaction strength of each material point r with its neighbors. As this strength is related to the per- meability field b(xz), the underlying microstructure can be recovered accordingly. In the bottom row of Figure EL we demonstrate the discovered microstructure of this test task. We note that the discov- ered microstructure is smoothed out due to the continuous setting our learned kernel (as shown in the bottom left plot), and a thresholding step is performed to discover the two-phase microstruc- ture. The discovered microstructure (bottom right plot) matches well with the hidden ground-truth microstructure (left plot), except for regions near the domain boundary. This mismatch is due to

---

b(x)
g(k)
pix)
203
100
0.4
0.4
0.4
0.2
0.2
0.2
0.4
4.6
0
0.5
0.5
0.5
/ [K]w.f]0xy]] dy
Recovered b(x)
0.5
0.8
0.4
0.4
0.2
0.2
0.5
Figure 4: Demonstration of the generated data and the recovered microstructure from the learned kernel in Example 2. Top row: the ground-truth two-phase material microstructure from test task (left), an exemplar loading field instance (middle), and the corresponding solution field instance (right). Bottom row: summation of the learned kernel for each line, corresponding to the total interaction of all material points (left), and the discovered two-phase material microstructure after thresholding (right).
the applied Dirichlet-type boundary condition (p(x) = 0 0?) in all samples, which leads to the measurement pairs (p(æ), g(x)) containing no information near the domain boundary O and makes impossible to identify the kernel on boundaries.
5.3 Heterogeneous material learning
In this example, we investigate the learning of heterogeneous and nonlinear material responses us- ing the Mechanical MNIST benchmark [[Lejeune, 2020]. For training and testing, we take 500 heterogeneous material specimens, where each specimen is governed by a Neo-Hookean material with a varying modulus converted from the MNIST bitmap images. On each specimen, 200 load- ing/response data pairs are provided. Two generalization scenarios are considered. 1) We mix data from all numbers and randomly take 10% of specimens for testing. This scenario corresponds to an ID test. 2) We leave all specimens corresponding to the number '9' for testing, and use the rest for training. This scenario corresponds an OOD test. The corresponding results are listed in Table & where NAO again outperforms its discrete counterpart.
Table 3: Test errors and the number of trainable parameters in heterogeneous material learning.
Case
model NiO
#param
ID test
OOD test
d = 40, d==40
5,469,528
7.21% %
7.95%
22500 samples
NAO
142,534
6.57% %
6.26%
d == 100, do=100
Discrete-NAO NAO
7,353,768
6.34%
6.01%
45000 samples
303,814
4.75 %
5,58 %
6 Conclusion
We propose Nonlocal Attention Operator (NAO), a novel NO architecture to simultaneously learn both the forward (modeling) and inverse (discovery) solvers in physical systems from data. In par- ticular, NAO learns the function-to-functioo mapping based on an integral NO architecture and pro- vides surrogate forward solution predictor. In the meantime, the attention mechanism is crafted in building a kernel map from input-output function pairs to the system's function parameters, offering zero-shot generalizability
new and unseen physical systems. As such, the kernel map explores in the function space of identifiability, resolving the enduring ill- ill-posedness in inverse PDE problems. In our empirical demonstrations, NAO outperforms all selected baselines on multiple datasets of inverse PDE problems and out-of-distributin generalizability tasks.
Broader Impacts: Beyond its merits forward/inverse PDE modeling, our work represents an initial exploration in understanding the attention mechanism in physics modeling, and paves a theo- retical path towards building a foundation model in scientific MI.
Limitations: Due to limited computational resource, our experiments focus on learning from small medium number& (< 500) of similar physical systems. It would be beneficial to expand the coverage and enable learning across different types of physical systems.
10

---

Acknowledgments and Disclosure of Funding
S. Jafarzadeh would like acknowledge support by the AFOSR grant FA9550-22-1-0197 and Y. Yu would like to acknowledge support by the National Science Foundation (NSF) under award DMS- 1753031 and the National Institute of Health under award IR0IGM157589-01 Portions of this research were conducted on Lehigh University's s Research Computing infrastructure partially sup- ported by NSF Award 2019035, Lu would like acknowledge support by NSF DMS-223842
This article has been authored by an employee of National Technology and Engineering Solutions of Sandia, LLC under Contract No. DE-NA0003525 with the U.S. Department of Energy (DOE). The employee owns all right, title and interest and to the article and is solely responsible for its contents. The United States Government retains and the publisher, by accepting the article for publication, acknowledges that the United States Government retains a non-exclusive, paid-up, ir- revocable, worldwide license to publish or reproduce the published form of this article or allow others to do so, for United States Government purposes. The DOE will provide public access to these results of federally sponsored research in accordance with the DOE Public Access Plan http://wwwwwnnnngggggg/o/ollll// wwwp///wwwwws///oeeeeeeoeoeeeeee https///ddo-uul--c-cceeeee public-access-plar
References
Gilles Besnard, François Hild, and Stéphane Roux. finite-element displacement fields analysis from digital images: application to portevin- le châtelier bands. Experimental mechanics, 46(6):789- 803, 2006.
Claudia Bucur, Enrico Valdinoci, et al. Nonlocal diffusion and applications, volume 20. Springer, 2016.
Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis. Physics- informed neural networks (PINNs) for fluid mechanics: A review. Acta Mechanica Sinica, pages 1-12, 2022.
Shuhao Cao. Choose a transformer: Fourier or galerkin. Advances neural information processing systems, 34:24924-24944 2021.
Giuseppe Carleo, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie Vogt-Maranto, and Lenka Zdeborová. Machine learning and the physical sciences. Re- views of Modern Physics, 91(4):045002, 2019.
Ke Chen, Chunmei Wang, and Haizhao Yang. let data talk: data-regularized operator learning theory for inverse problems. arXiv preprint arXiv:2310.088
2023.
Rewon Child, Scott Gray, Alee Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Genevieve Coorey, Gemma A Figtree, David F Fletcher, Victoria J Snelson, Stephen Thomas Ver- non. David Winlaw. Stuart M Grieve, Alistair McEwan, Jean Yee Hwa Yang, Pierre Qian, Kieran O' O'Brien, Jessica Orchard, Jinman Kim, Sanjay Patel, and Julie Redfern. The health digital twin to tackle cardiovascular diseaseA review of an emerging interdisciplinary field. NPJ Digital Medicine, 5:126, 2022.
Wen Ding, Kui Ren, and Lu Zhang. Coupling deep learning with full waveform inversion. arXiv preprint arXiv:2203.0/ 007999 2022.
Sören Dittmer, Tobias Kluth, Peter Maass, and Daniel Otero Baguer. Regularization by architecture: A deep prior approach for inverse problems. Journal of Mathematical Imaging and Vision, 62: 456-470, 2020.
Qiang Du, Bjorn Engquist, and Xiaochuan Tian. Multiscale modeling, homogenization and nonlo- cal effects: Mathematical and computational issues. Contemporary mathematics, 754:115-1444 2020.
Yuwei Fan and Lexing Ying. Solving traveltime tomography with deep learning. Communications in Mathematics and Statistics, 11(1):3-19, 2023.

---

Alberto Ferrari and Karen Willcox. Digital twins in mechanical and aerospace engineering. Nature Computational Science, 4(3):178-183, 2024.
J Ghaboussi, JH Garrett Jr, and Xiping Wu. Knowledge-based modeling of material behavior with neural networks. Journal of engineering mechanics, 117(1):132-133 1991.
Jamshid Ghaboussi, David A Pecknold, Mingfu Zhang, and Rami M Haj-Ali. Autoprogressive training of neural network constitutive models. International Journal for Numerical Methods in Engineering, 4211):105-, 1998.
Somdatta Goswami, Aniruddha Bora, Yue Yu, and George Em Karniadakis, Physics-informed neu- ral operators. 2022 arXiv preprint vrxivv2207.05745 2022.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catan-
zaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.1358/ 2021.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet--aseed operator learning for differen- tial equations. Advances in neural information processing systems, 34:24048-24066 2021.
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. Gnot: A general neural operator transformer for operator learning. International Conference Machine Learning, pages 12556-12569, PMLR, 2023.
Qizhi He, Devin W Laurence, Chung-Hao Lee, and Jiun-Shyan Chen. Manifold learning based data-driven modeling for soft biological tissues. Journal of Biomechanies, 117:110124, 2021.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidi- mensional transformers. arXiv preprint 19Xiv: 2019.
Siavash Jafarzadeh, Stewart Silling, Ning Liu, Zhongqiang Zhang, and Yue Yu. Peridynamic neural operators: A data-driven nonlocal constitutive model for complex material responses. Computer Methods in Applied Mechanics and Engineering, 12911, 2024.
Hanyang Jiang, Yuehaw Khoo, and Haizhao Yang. Reinforced inverse scattering. arXiv preprint arXiv:2206.04 046. 2022.
George Em Karniadakis, loannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6):422-440, 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference chine learning, pages 5156-5165. PMLR, 2020.
Ru- Yu Lai. Qin Li, and Gunther Uhlmann. Inverse problems for the stationary transport equation in the diffusion scaling. SIAM Journal on Applied Mathematics, 79(6):2340-2355) 2019.
James Lee- Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint arXiv:2105.0335 2021.
Emma Lejeune. Mechanical mnist: A benchmark dataset for mechanical metamodels. Extreme
Mechanics Letters, 36: 36:100659, 2020.
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. arXiv preprint wxXiv:2205./367/ 2022.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadeneshelli Burigede Liu, Kaushik Bhattacharya, An- drew Stuart. and Anima Anandkumar. Neural operator: Graph kernel network for partial differen- tial equations. arXiv preprint WXX.::20033334885 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differ- ential equations. Advances in Neural Information Processing Systems, 33:NeurIPS 2020, 2020b.
12

---

Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier NeuralOperator for Parametric Partial Differential Equa- tions. In International Conference on Learning Representations, 2020c.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informae neural operator for learning partial differential equations. 2021 arXiv preprint arXiv:2111.03794, 2021.
Hanxiao Liu, Zihang Dai, David So, and Quoe V Le. Pay attention to mips. Advances in neural information processing systems, 34:9204-922 2021.
Ning Liu, Yue Yu, Huaiqian You, and Neeraj Tatikola. Ino: Invariant neural operators for learning complex physical systems with momentum conservation. In International Conference on Artifi- cial Intelligence and Statistics, pages 6822-6838. PMLR, 2023.
Ning Liu, Yiming Fan, Xianyi Zeng, Milan Klöwer, LU ZHANG, and Yue Yu. Harnessing the power of neural operators with automatically encoded conservation laws. In Forty-first International Conference on Machine Learning, 2024a.
Ning Liu, Siavash Jafarzadeh, and Yue Yu. Domain agnostic fourier neural operators. Advances in Neural Information P Processing Systems, 36, 2024b.
Ning Liu, Xuxiao Li, Manoj R Rajanna, Edward W Reutzel, Brady Sawyer, Prahalada Rao, Jim Lua, Nam Phan, and Yue Yu. Deep neural operator enabled digital twin modeling for additive manufacturing. arXiv preprint arXiv:245.009572, 2024c.
Xinliang Liu, Bo Xu, and Lei Zhang. Ht-net: Hierarchical transformer based operator learning model for multiscale pdes. 2022.
Fei Lu, Quanjun Lang, and Qingci An. Data adaptive rkhs tikhonov regularization for learning kernels in operators. In Mathematical and Scientific Machine Learning, pages 158-172. PMLR, 2022.
Fei Lu, Qingci An, and Yue Yu. Nonparametric learning of kernels in nonlocal operators. Journal of Peridynamics and Nonlocal Modeling, pages 1-24, 2023.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for iden- tifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv: 1910.099900019990 2019.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):118 22229,9,,, 2021.
Roberto Molinaro, Yunan Yang. Björ Engquist, and Siddhartha Mishra. Neural inverse operators for solving pde inverse problems. arXiv preprint a7Xiv:2301.11167, 2023.
Christoph Molnar. Interpretable machine learning. 2020.
Anahita Nekoozadeh, Mohammad Reza Ahmadzadeh, and Zahra Mardani. Multiscale attention via wavelet neural operators for vision transformers. arXiv preprint arXiv:2303. /2398, 2023.
Daniel Obmann, Johannes Schwab, and Markus Halmeeer. Deep regulazzionn of inverse problems. arXiv preprint (=XXi::002200/55 2020.
Yong Zheng Ong. Zuowei Shen, and Haizhao Yang. IAE-Net: integral autoencoders for discretization-invariaa learning. arXiv preprint 203.5:2003 2022.
David Pfau, James S Spencer, Alexander GDG Matthews, and W Matthew C Foulkes. Ab initio solution of the many-electron schrödinger equation with deep neural networks. Physical Review Research, 2(3)20334 2020.
Yongming Rao, Wenliang Zhao Zheng Zhu, Jiwen Lu, and Jie Zhou. Global filter networks for image classification. Advances in neural information processing systems, 34:980-993, 2021.
13

---

Cynthia Rudin, Chaofan Chen, Zhi Chen, Haiyang Huang, Lesia Semenova, and Chudi Zhong. In- terpretable machine learning: Fundamental principles and 10 grand challenges. Statistic Surveys, 16:1-85, 2022.
Stewart A Silling. Reformulation of elasticity theory for discontinuities and long-range forces. Jour- nal of the Mechanics and Physics of Solids, 48(1):175-22) 209,-2999911 2000.
Jingmin Sun, Yuxuan Liu, Zecheng Zhang, and Hayden Schaeffer. Towards foundation model for partial differential equation: Multi-operator learning and extrapolation. arXiv preprint arXiv:2404. 12355, 2024.
Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp-mixer: An all-mlp architecture for vision. Advances in neural information processing systems, 34:24261- 24272, 2021.
Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Gautier Izacard, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, et al. Resmlp: Feed- forward networks for image classification with data-efficient training. IEEE Transactions Pattern Analysis and Machine Intelligence, 45(4):5314 5321, 2022.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: unified understanding of transformer's 's attention via the lens of kernel. arXiv preprint arXiv:/908. 190. 2019.
Gunther Uhlmann. Electrical impedance tomography and calderón's problem. Inverse problems, 25
(12):123011, 2009.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768. 2020.
Min Wei and Xuesong Zhang. Super-resolution neural operator. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18247-18256, 2023.
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle:
Fourier analysis sheds light on deep neural networks. arXiv preprint arXiv: 1901.0639, 2019.
Liu Yang and Stanley J Osher. Pde generalization of in-context operator networks: A study ld scalar nonlinear conservation laws. arXiv preprint arXiv:2401.07364 2024.
Liu Yang, Xuhui Meng, and George Em Karniadakis. B-pinns: Bayesian physics-informed neu- ral networks for forward and inverse pde problems with noisy data. Journal of Computational Physies, 425:109913, 2021.
Zhanhong Ye, Xiang Huang, Leheng Chen, Hongsheng Liu, Zidong Wang, and Bin Dong. Pde- former: Towards foundation model for one-dimensional partial differential equations. arXiv preprint arXiv:2402.1263 2024.
Özdogan Yilmaz. Seismic data analysis, volume Society of exploration geophysicists Tulsa, 2001.
Huaiqian You, Yue Yu, Nathaniel Trask, Mamikon Gulian, and Marta DElia. Data-driven learning of nonlocal physics from high-fidelity synthetic data. Computer Methods in Applied Mechanics and Engineering. 174:133553, 2021.
Huaiqian You, Yue Yu, Marta D'Elia, Tian Gao, and Stewart Silling. Nonlocal kernel network (NKN): A stable and ssolution-inndeeedder deep neural network. Journal of Computational Physics, page arXiv preprint 2XXv:2201.0021111 2022.
14

---

Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and E Weinan. Deep potential molecular dynamies: a scalable model with the accuracy of quantum mechanics. Physical Review Letters, 120(14):14300, 2018.
Zecheng Zhang. Modno: Multi operator learning with distributed neural operators. arXiv preprint arXiv:2404.02892, 2024.
15

---

A Riemann sum approximation derivations
In this section, we show that the discrete attention operator (g) can be seen as the Riemann sum approximation of the nonlocal attention operator in (8), in the continuous limit. Without loss of generality, we consider a uniform discretization with grid size Az. Denoting the (-th layer out- put as U = -00)(zz)) ISFSdISKEN"
F = (F) ISSSaISAEN'
and w, w/K K as
vdl
woR E Raxd, for the 1-th (1 < L) layer update in (6) writes:
x (a,
3| = Σ
2N
d
de
1
x
(a,w) A/W (M, X|I bx x
01 + x
x1
vdx
>=1 w.vml
A=1
2N
Σ
(a
x| (y. v1) (1,00 (a.
=
That mcans, de only controls the rank bound when lifting each point-wise feature via W.KK while woR characterizes the (trainable) interaction of d input function pair instances. Moreover, when S N, X (0, 5] = 9 (Ia). When a N. X (1) [a, 31 = 1 (Fa-N). Then,
98 (za)
x (a w. OK [ (xx vi) x x11-11 (1 .3]
<= wi==l
2N
wr 11 12 (7: vi)
3] o (2a)
1,
(aa) (w: u/o! (m,)) 05 (1a)
2 (0 w, 11 13 (o,)) + 26 -02 (xta)
5= w
Similarly,
N (xa)
Σ
Σ
(ws (f
1 w,am1
2N
(a
w. x2 (7) vi) x (1 (7
3] + (x.)
N
Σ
w|9 ) "3 (z,))
7=l m,
( (z.)) (w, (, (uz)) 1/t (1x) + 1 (x.))
With the Riemann sum approximation: Sa p())d R A Ei p(xx), one can further reformu- late above derivations as:
9 (x) = A
g (x) (1 ((n) 93¹ (y)dy
w,uml
11 (20 ( w. 5 (w)dy + 9h (x),,
+ D w,u=l
(x) A
Σ
(,, (w)) 9/ (y)dy
Σ (x)
[, -)) (v)) f "0)dd + (z))
D
1t w.um1
16

---

Therefore, the attention mechanism of each layer is in fact an integral operator after rescaling:
( [ ) -/11amm0 x ( ((--l(y) ) hp +
)
( [
(20)
with the kernel K (2, v) defined as:
Ed س (x) (x) be visl Ewee ( (x) QK [w, p/ge (e)) Eup=el (z) Qk ( 1 (1
D (z) wi (1) (1)
(21)
For the L-th layer update, we denote the approximated value of fz( as fa(Ia) = Xoula, 3). then
N Σ
fa(xa)
=
K [undd frati OI rla(up](a,)
N
WP
( i (. glua)(ar))
Σ
Σ
(a,
(A (e b7s vl) slua)(e))
An=1w,v=1 m,
N
Σ
W 1PP:la
(w)) glua)(En)
Σ WPi
(a, (1 (F!t (w, p|gاt ((x)) g g(usl(er)
A w,tml
(Fil( 2) (ol- (=) (w, ugg ())) d= glusl()dyy
DD
32 w,u=l
Σ (PPl(xaa2) (2 (w, Ww - de (l)o)
Ax
Hence, (rescaled) continuous limit of the kernel writes:
d
K [[uad fisd (z, v)
Σ
(P:
(x) (w, vlg! ())) dz
Q w,n=1
W (x, (I)-" ww, dz.
B Proofs and connection with regularized estimator
B.1 Proofs
Proof of Lemma g
With (rx JV we can write the attention-baue based kernel in (
N
de
K (Uid L (TR)) =
u (rw 1 1 (u)
==1
(22)
du
d
+ WPS.
[
1=1 ==
17

---

Denoting Wo*(J',J) Ei= 1 W" .1, we write the kernel in (222) as
K[ fi ()(r) = Σ
W (P'((
glu](r', xi, ) u. U.
r==1
+
WPJ
Σ
[ WWok (. (l m
Then, as d + oc is achieved by refining the spatial mesh, viewing the summation in j as Riemann sum,
glul(r', ai ,,,,) (, / glul(r, (x, v)g[u] y)drdy.
lim 10000
==1
where the integral exists since g => (z, y) is bounded. Sending also the number of tokens, N, to infinity, we obtain the limit attention model in (l)).
Proof of Lemma 2 The proof is adapted from Lu all 2023
2022
Write K(r)
K [u h:d: (r). Notice that the loss function in (9) can be expanded as
E(K) 1 Berain du
1. K()))u((( (7)(0) da (2((0)
Betriin i=1
Terain du
[ l' (())K()) / glu7](s, (l) drddrss
Btrain
11 1=1
Thenii du
2
+ Const. K(r) [ a) of)
Ptrain
7=
= KLoK,K)13 2(K, KP) 13 Const.,
where Lo 12 + 1 the integral operator
(s)(s)== [ K(r)G(r,s)dr
and KD is the Riesz representation of the bounded linear functional
(K, (,KK))
1 Atrain du
Σ EEJ
[ K(r) / fl
Tttrain
Thus, the quadratic loss function has a unique minimizer in Null(Ca)
Meanwhile, since the data pairs are continuous with compact support, the function G is square- integrable reproducing kernel. Thus, the operator Lg is compact and He L1/ L2 Then, Null(Ce)' = HG, where the closure is with respect to L
B.2 Connection with regularized estimator
Consider the inverse problem of estimating the nonlocal kemel K given data pair (u, f). In the classical variational approach, one seeks the minimizer the following loss function
8(K) - 0 [ K(r)glu((;; - [(x) dr
(23)
The inverse problem is ill-posed in the sense that the minimizer can often be non-unique or sensitive to the noise or measurement error in data (u, f). Thus, regularization crucial to produce stable solution.
18

---

To connect with the attention-base model, we consider regularizing using an RKHS Hw with square-integrable reproducing kernel W. One seeks an estimator in Hw by regularizing the loss with the K il and minimizes the regularized loss function
(x,( ((())=(())) + J / (r)) s)rrds.
(24)
The next lemma shows that the regularized estimator is nonlinear function of the data pair (u, f), where the nonlinearity comes from the kernel G. and the parameter A
Lemma B.I. The regularized loss function in Ex.w (K) in ((24) is
(K) - 30 00
((r)K())Gu((r,)) + (r, (1,8)) s)) 10² K(r) (r)dr + Const.,
(25)
where G is defined in (2 Iss minimizer is
K= (Ca, + X.LW)
(26)
where LG., and Lw are integral operators with integral kernels Gy defined in (27) and W
LG. (() J 00
G(r),ulr,s)dr, (w)(s)== == [ (rr)W(r,
X is the optimal hyper-parameter controlling the strength of regularization, and K Sa glu](r,x) x) (x)dd is a function determined by the data f).
When there no prior information on the regularization, which happens often for the learning of the kernel, one can the data-adaptive RKHS HG. with the reproducing kernel G. determined by data:
Gu(r, - glu)(r,,))(()))) (u] x))d.
(27)
Ln 2022 shows that this regularizer can lead to consistent convergent estimators.
Remark B.2 (Discrete data and discrete inverse problem). In practice, the datasets are discrete. One can view the discrete inverse problem discretization of the continuous inverse problem. Assume that the integrands are compactly supported and when the integrals are approximated by Riemann sums, can write loss function for discrete K (K(r))...., K(NN)) Nx1
N.N Σ
(X) ~
+ (T ((rt))(Tr))GG.. - Σ K(rr)Kii n. (rx))rr Const..
R,K'1
=K +GG+AWJJK 2 K Const.,
(28)
where, recalling the definition of the token Una in (B),
G (Gu = 1/f glul(z, ma. un.auidd
W (W(rr,e)))sse and K (Ja glw (f) ISKEN
The minimizer of this discrete loss function with the optimal
X
K = G. + X. W) WK'... -K...
In partieular, when taking W
G., and using the Neumann series (A=+Gz 1)-1
( x G we have
k= (A G
- = G. X; 2G + --1)**> G K.
In particular, X, depends on both G. and Ku.f, Hence, the estimator K is nonlinear in G., and K and it is important make the attention depend nonlinearly on the token undd as in (
19

---

Proof of Lemma B Since K is radial and noticing that
/
(u(a + r5) u()))dc = uxx + r) + u(x u(x-r) - -u((x) (
|9|=1
since IR,,
write
/ -)(())))(((((+++ + - u(x)))dxdddd
<(()))<((())
12
K(r)K(s)) / I
16'=11 1KE== J (u(x + sE' u(x))) + +r) - =(xr))dxdEdd"
00
((T)K(s) / gluu(rr
z)dzdrds
J0
10
(r,
Meanwhile, by the Riesz represeetttiooiioooon theorem, there exists a function KudJ E T,2 oo) such that
K((r)) Ja (u(x+r) - -
/ (r) I 18|=1
J.- K(r)Kal (r)dr.
((r)
Combining these two equations, we can write the loss function as
E(K)-
[(Is))u(x +)) u(x))ds - f(x)) ((xdss)())) dx
K((()))(HH)) / (u(or + u)())+ u(x))dddrds
((ss)) (u(x + u( (x) Const.
! (())K())Gu(r, s)drds 2 { K(r) Const..
Then, write the regularized loss function εx,w (K)
(x) = (/,, I. ((r) AW (r, Ar, - 2 f. K(r)K" (r)dr Const.
=(LGo, + +XWw) K, K) 2(K, K 112000)))00))) + Const.
Selecting the optimal hyper-parameter X., which depends both (u, and W, and setting the Fréchet derivative of Ex,w over 1,2(0, oo) to be zero, we obtain the regularized estimator (26). D
C Data generation and additional discussion
C.1 Example 1: radial kernel learning
In all settings except the "single task" one, all kernels act on the same set of functions {u, 1u1,2 with #1 ( (z) and ug(x) sin(2.z) (2xx)1 | -,)] (z). the "single task" setting, create more diverse samples, the single kernel acts set of 14 functions: 4y (os(k=))(<=nn.)]]
k 1.... 7 and (xx) sin(kz) (z) k ..... 14. In the ground-truth model, the integral L,,, [w:] is computed by the adaptive Gauss-Kronnn
quadrature method, which is much more accurate than the Riemann sum integrator that we use in the learning stage. To create discrete datasets with different resolutions, for each Ax E 0.0125 x (1,2), we take values (u, fa J
20

---

0.8
1.5
Groundenuh odt-s-diccibation kendl
Ground truth distribttion lest2 kemeel
(sin only: -12
VAE burtcostpyy)) kemel 128.00%
0.6
NAO caa--co-ppyy kemel etror -6.92%
teintcostpoly. less panmotery kemel -94
NAO (sta oos poly. less samples): kemel crroe -4.04%
NAO (sirtcoctpolg): kend erce -10.48%
K((x-1))
0.4
K(x-yl
NAO (sin cosly.ly. less sample): kornel emror -25.58%
0.2
0.5
0
Xx-y|
XX-X|
Figure 5: OOD test results on radial kernel learning, with diverse training tasks and 302. OOD1 (left): true kemel 7(r) r(11 - 7) exp(-5r) sin (r(;;;; OOD2 (right): true kernel ((r) exp( exp(-0.3r")
Gaussian kernel which is very different from all training tasks.
V2x
{ui(Tz), J((1)): :y E of size Ar. We form training sample of each task by taking d pairs from this task. When taking the token size d and 1o function pairs, each task contains 2265 x & samples.
J. where x. is the grid point on the uniform mesh
More diverse tasks. To further evaluate the generalization capability of NAO as foundation model, we add another two types of kernels into the training dataset. The training dataset is now constructed based on 21 kernels of three groups, with 15 samples on each kernel:
sine-type kernels: 7/7 3 (r) = exp(-1r) sin (o.1 1,2,,,,,,,,
cosine-type kernels: 7, (r)
10-r (000)))(10
-
20
polynomial-type kernels: "m (r) = exp(-0.1r)p) 1,2,3, 4, 5, 6, 7, where Pt, is the degree-n Legendre polynomial.
('T00)
r=10)
(r),,
Based on this enriched dataset, besides the original "sine only" setting, three additional settings are considered to compare the generalizability across different settings. The results are reported in Table g and Figure HL In addition to the original setting corresponding to all "sine" kernels, in part of Table B (denoted as "sine+cosine+polyn"" we consider a "diverse task" setting, where all 315 samples are employed in training. In part III, we consider a "single task" setting, where only the first sine-type kernel, 7(r) exp(-7) (o((r) sin (6r) (o.)))). is considered as the training task, with 105 samples on this task. Lastly, in part IV we demonstrate a "fewer samples" setting, where the training dataset still consists of all 21 tasks but with only 5 samples on each task. For testing, besides the ID and OOD tasks in the ablation study, we add an additional OOD task with Gaussian-type kernel
Tood2 (r) exp(-0.5r") which is substantially different from all training tasks. V2n
As shown in Table BL considering diverse datasets helps in both OOD tests (the kernel error is reduced from 9.14% to 6.92% in OOD test 1, and from 329.662 69% to 10.48% in OOD test 2), but not for the ID test (the kernel error is slightly increased from 4.03% to 5.04%). We also note that since the "sine only" setting only sees systems with the same kernel frequency in training, it does not generalize to OOD test 2, where it becomes necessary to include more diverse training tasks. As the training tasks become more diverse, the intrinsic dimension of kernel space increases, requiring a larger rank size du of the weight matrices. When comparing the "fewer samples" setting with the "sine only" setting, the former exhibits better task diversity but fewer samples per task. One can see that the performance deteriorates on the ID test but improves on OOD test 2. This observation emphasizes the balance between the diversity of tasks and the number of training samples per task.
C.2 Example 2: solution operator learning
For one example, we generate the synthetic data based on the Darcy flow in a square domain of size 1 x subjected to Dirichlet boundary conditions. The problem setting is: - V(b<>)Vp(x) Vy(x)) = g()) subjected to p(x) = 0 on all boundaries. This equation describes the diffusion in heterogeneous fields, such as the subsurface flow of underground water in porous media. The heterogeneity represented by the location-dependent conductivity b(x). p(x) is the source term, and the hydraulic
21

---

Table 4: Training and test for the radial kernel problem with more diverse tasks with fixed d 302, where bold numbers highlight the best methods across different data settings. These results emphasize the balance between task diversity and the number of samples per task.
dx
model
#param
Operator test error
Kernel test error
ID
OODI
OOD2
ID
OODI
OOD2
sine only:
Train on sine dataset, 105 samples in total
10
Discrete-NAO
16526
I.33%
25.81%
138.00%
29.02%
28.80%
97.14%
NAO
18843
1.48%
8.10%
221.04%
5.40%
10.02%
420.60%
20
Eiscrete-NOO
28645
1.35%
18.70%
99.50%
35,49%
30.81%
101.08%
NAO
30963
1.33%
9.12%
211.78%
4.63%
9.14%
329.66%
40
Discrete-NAO
52886
1.30%
31.37%
49.83%
38.89%
30.02%
129.84%
NAO
55203
0.67 %
7.34%
234.43%
4.03%
12.16%
1062.3%
Autoencoder
16424
12.97%
1041.49%
698.72%
22.56%
136.79%
304.37%
sine+cosine+polln: Train on diverse (sine, cosine and polynomial)
dataset, samples total
10
Discrete-NAO
16526
2.27%
13.02%
11.50%
10.41%
30.79%
77.80%
NAO
18843
2.34%
14.37%
10.05%
7.26%
28.23%
94.38%
20
Discrete- NAO
28645
1.60%
6.03%
19.73%
21.83%
21.29%
18.97%
NAO
30963
1.64%
3.25%
3,58%
5.45%
8.87%
15.82%
40
Discrete- NAO
52886
1.45%
5,49%
18.26%
20.07%
19,46%
18.44%
NAO
55203
1.54%
3.09%
7.69%
5.04%
6.92%
10.48 %
Autoencoder
56486
12.67%
341.96%
211.61%
27.06%
52.43%
128.08%
Single task:
Train on single sine dataset, 105 samples total
10
NAO
18843
104.49%
104.37%
56.64%
100.31%
100.00%
94.98%
20
NAO
30963
116.89%
105.52%
85.40%
99.02%
99.55%
98.37%
40
NAO
55203
111.33%
104.12%
76.78%
101.61%%
100.33%
95.90%
Fewer samples: Train on diverse (sine, cosine and polynomial) dataset, 105 samples in total
10
NAO
18843
4.23%
15.34%
11.13%
10.11%
25, 25.23%
97.63%
20
NAO
30963
4.15%
10.15%
9.59%
8.84%
23.08%
20.05%
40
NAO
55203
3.69%
11.67%
6.08%
9.19%
24.04%
25.58%
height (x:) is the solution. For each data instance, we solve the equation on 21 21 grid using an in-house finite difference code. We consider 500 random microstructures consisting of two distinct phases. For each microstructure, the square domain is randomly divided into two subdomains with different conductivity of either 12 or 3. Additionally, we consider 100 different g(x) functions ob- tained via Gaussian random field generator. For each microstructure, we solve the Darcy problem considering all 100 source terms, resulting in dataset of N 10, 000 function pairs in the form of (x,( g: (x,)) and j 1,2, 441 where xj are the discretization points on the square domain.
In operator learning settings, we note that the permutation of function pairs in each sample should not change the learned kernel, i.e., one should have K [[uad faaa] = K ([uo(1:a); (o(1:))). where o is the permutation operator. Hence, 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 permut the function pairs and take 100 function pairs for 100 times per task. As result, we can generate total of 10000 samples (9000 for training and 1000 for testing) in the form of (u".100- oo 10000
C.3 Example 3: heterogeneous material learning
In the Mechanical MNIST dataset, we generate a large set of heterogeneous material responses subjected to mechanical forces. This similar to the approach Lejeune [2020] The material property (stiffness) of the heterogeneous medium is constructed assuming linear scaling between and 100, according to the gray-scale bitmap of the MNIST images, which results a set of 2D) square domains with properties that vary according MNIST digit patterns. The problem setting for this data set is the equilibrium equation:-Vy P(x) f(x) subjected to Dirichlet boundary
condition of zero displacement and nonzero variable external forces:f(x). P(x)) P(T + Vu(æ:) the stress tensor and a nonlinear function of displacement u(x)) The choice of material models determines the stress function, and here we employ the Neo-Hookean material model as in Lejeune
22

---

[
From MNIST, we took 50 samples of each digit resulting in total of 500 microstructures.
We consider 200 different external forces f(z) obtained via a Gaussian random field, and solve the problem for each pair of microstructure and external force, resulting in N 100, 000 samples. To solve each sample, we use FEniCS finite element package considering a 140 x 140 uniform mesh. We then downsample from the finite element nodes to get values for the solution u(x)) and external force f(z) on a the coarser 29 x 29 equi-spaced grid. The resulting dataset is of the form: {u,(r;), Ji and j 2, ..., 841 where z, 's are the equi-spaced sampled spatial points.
Similar to Example 2, we perform the permutation trick to augment the training data. particular, with 500 microstructures (tasks) and 200 function pairs per task, we randomly permute the function pairs and take 100 function pairs for 100 times per task. As a result, we can generate a total of 50000 samples, where 45000 are used for training and 5000 for testing.
D Computational complexity
Denote (N, d, dy, L, Reraia) as the size of the spatial mesh, the number of data pairs in each training model, the column width of the query/key weight matrices W and W.k the number of layers, and the number of training models. The number of trainable parameters in a discrete NAO is of size OLL d x dę + 22). For continuous-kemmel NAO, taking three-layer MLP dense net of size (d, h2, 1) for the trainable kernels W and WPF P.J for example, its the corresponding number of trainable parameters is Old x ln + h1 x h2). Thus, the total number of trainable parameters for continuous-kernel NAO is OLL x d du + d x h1 h1 x h2).
The computational complexity of NAO is quadratic in the length of the input and linear in the data size, with ((d2((N + dx) + computed as follows for each layer and the data of each training model: the attention function takes Oldd 2N + 4N2 d) flops, and the kernel map takes O(d de + Nd + dN2 flops; thus, the total is (d + d&) + 6N2d) d) flops. In inverse PDE problems, we generally have d < N, and hence the complexity of NAO is O(N'd) d) per layer per sample.
flops in each epoch in the optimization. It is
Compared with other methods, it is important note that NAO solves both forward and ill-posed inverse problems using multi-model data. Thus, we don't 't compare it with methods that solve the problems for a single model data, for instance, the regularized estimator in Appendix B. Methods solving similar problems are the original attention model [[Vasann et all 2017 convolution neural network (CNN), and graph neural network (GNN). As discussed in Maswani al [Xي]], these models have a similar computational complexity, if not any higher. In particular, the complexity of the original attention model (( and the complexity of CNN is (kNdd) with being the kernel size, and a full GNN is of complexity 2
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have clearly stated claims in the abstract and introduction, includ- ing the key contributions that our proposed model NAO provides a novel input-dependent kernel for context extraction, integrates forward and inverse PDE predictions in one frame- work, and offers the capability in solving ill ill-posed inverse PDE problems.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: The limitations of the work is discussed in the Conclusion section.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We have provided the full set of assumptions and corresponding complete proofs under each proposed Lemma with detailed derivations included in the Appendix A and B.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data provided not)?
Answer: [Yes]
Justification: We have elaborated the detailed architectures of the proposed model in the main body of the paper well data generation details in Appendix C. The code and accompanying data have been released at
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The code and accompanying data have been released at E33/Eishhmoon1337/MM..
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We briefly state the training and test strategies at the beginning of the Experi- ments section, with the full details provided in Appendix
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [No]
Justification: Error bars are not reported due to the limited computational resource. Nev- ertheless, we have tuned the hyperparameters of the considered models and reported extensive results in a large variety of settings, and they all support our claim.

---

8. Experiments Compute Resourees
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
The computing resources are reported in the first paragraph of the Exper- iments section. In particular, all the experiments are conducted on a single NVIDIA GeForce RTX 3090 GPU with 24 GB memory.
Justification:
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httpp////aarrss..cooa//////hclcl
Answer: [Yes]
Justification: The research conducted in the paper conform with the NeurIPS Code of Ethics.
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The discussion the societal impacts can be found in the Conclusion section.
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have cited the original papers that produced the code package or dataset where appropriate.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer:
[Yes]
Justification: The code and accompanying data have been released at https://github. 1234/NAO, with instructions and pretrained models provided.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects. 15. Institutional Review Board (IRB) Approvals Equivalent for Research with Human
Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
25

---

Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
26


Reviews:
[{"review_id": "hYcpDfnh2O", "reviewer": "Reviewer_eXH4", "summary": "This paper proposes a novel neural operator architecture called Nonlocal Attention Operator (NAO) for learning both forward and inverse problems in physical systems from data. The key contributions are:\n\n1. A new attention-based neural operator that can simultaneously learn the forward mapping (physics modeling) and inverse mapping (physics discovery) for PDE systems.\n2. Theoretical analysis showing how the attention mechanism provides a kernel map that explores the space of identifiability for kernels, helping to resolve ill-posedness in inverse problems.\n3. Empirical demonstration of NAO's advantages over baseline methods, especially for generalization to unseen systems and data-efficient learning in ill-posed inverse problems.\n4. Interpretability of the learned kernels, allowing insight into the discovered physical mechanisms.\n\nThe authors evaluate NAO on several PDE learning tasks, including radial kernel learning, solution operator learning for Darcy flow, and heterogeneous nonlinear material modeling.", "questions": "Can you provide more insight into how the learned kernels can be interpreted to gain physical understanding? Perhaps a case study showing how NAO discovers meaningful structure in a physical system would be illuminating.", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a novel neural operator architecture that leverages attention in a unique way for PDE learning. The idea of using attention to build a kernel map that can generalize across multiple physical systems is innovative.\n2. The theoretical analysis in Section 4 provides rigorous justification for how the attention mechanism helps address ill-posedness. The empirical evaluations are comprehensive, testing multiple aspects like generalization, data efficiency, and interpretability.\n3. The paper is generally well-written and clearly structured. The motivation, methodology, and results are presented logically.\n4. The ability to simultaneously learn forward and inverse mappings for PDE systems, while providing interpretable kernels, could have broad impact in scientific machine learning and physics-informed AI. The zero-shot generalization capability is particularly notable.", "weaknesses": "1. The experimental evaluation is limited to relatively simple PDE systems. It's unclear how well NAO would scale to more complex, high-dimensional PDEs encountered in real-world applications.\n2. The interpretability claims could be further substantiated. While learned kernels are visualized, there's limited discussion on how these provide physical insights beyond matching ground truth.", "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**Interpretability for physical understanding**: The physical meaning of the learned kernel is application-dependent: in the context of material modeling, the learned kernel characterizes the interaction between material points; in turbulence modeling problems, the interaction range of the kernel suggests the characteristic length of the eddies; in fracture mechanics, the vanishing of kernel values between two neighboring regions may indicate a crack between these two regions; just to name a few. Generally speaking, identifying the kernel is instrumental in identifying influential factors, characterizing spatial dependence, and analyzing the forward operator. \n\n(i) The learned kernels themselves are fundamental for understanding the physical model. They indicate the range and strength of interactions between components. For example, a learned interaction kernel can reveal interactions between particles or agents. In the Darcy's problem $g\\rightarrow p$ setting, when taking a summation of the kernel strength on each row, one can discover the interaction strength of each material point $x$ with its neighbors. Then, since such a strength is related to the permeability field $b(x)$, the underlying microstructure can be recovered. In Figure 2 of the one-page PDF, we demonstrate the discovered microstructure on a test specimen. We note that the discovered microstructure was smoothed-out due to the continuous setting of our learned kernel (as shown in the middle plot), a thresholding step was performed to discover the two-phase microstructure. The discovered microstructure (right plot) matches well with the hidden ground-truth microstructure (left plot), except on the regions near the domain boundary. This is because of the Dirichlet-type boundary condition setting ($p(x)=0$ on $\\partial\\Omega$) in all samples. As a result, the measurement pairs $(p(x),g(x))$ contain no information near the domain boundary $\\partial\\Omega$, making it impossible to identify the kernel on boundaries. \n\n(ii) The learned kernel characterizes the output's spatial dependence on the input. The learned kernel in the nonlocal diffusion model of our Example 1 demonstrates this spatial dependence. Such a range of spatial dependence characterizes the characteristic lengths, which is especially of interests in multiscale problems (see, e.g., [1]).\n\n(iii) The learned kernels facilitate the analysis of the forward operator, including its spectral properties, stiffness for computation, and error bounds for output prediction under perturbations in the input.\n\nWe will add these discussions and the additional figure to the revised manuscript.\n\n[1] You, Huaiqian, et al. \"A data-driven peridynamic continuum model for upscaling molecular dynamics.\" Computer Methods in Applied Mechanics and Engineering 389 (2022): 114400.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "sf0rfS9bR1", "reviewer": "Reviewer_J2Mv", "summary": "This paper introduces a novel approach called the Nonlocal Attention Operator (NAO) for simultaneously solving forward and inverse PDE problems across multiple physical systems. The key innovation is using an attention-based kernel map that learns to infer resolution-invariant kernels from datasets, enabling both physics modeling and mechanism discovery.\n\nThank you to the authors for their responses. It is great to see a detailed complexity analysis; still, the quadratic scaling is definitely a cause for concern -- hence the reasons for relaxing kernel-based formulations of self-attention to linear scaling versions (e.g., Mamba, informer, etc.). However, I can see how similar relaxations can be applied to NAO. So this is very helpful.", "questions": "1. How does the computational complexity of NAO scale with the number of data points and dimensionality of the problem? Provide a detailed analysis of computational costs and compare them with existing methods for both forward and inverse PDE solving.\n\n2. What are the minimum data requirements for NAO to perform effectively, particularly for the inverse problem? It may be beneficial to provide an analysis of performance vs. data quantity to guide potential users in determining data collection needs.\n\n3. How can NAO be integrated with existing physics-based models or simulations to enhance their capabilities? Maybe discuss potential hybrid approaches combining NAO with existing or traditional methods.", "limitations": "I read the limitation sections, and the authors provide, I think, an inkling of the answers to the questions that I have raised. I hope to see a more comprehensive discussion on the same, with the necessary addendums made to the limitations section.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main strengths I find are as follows:\n\n(i) Bridging forward and inverse PDE modeling: The approach uniquely addresses forward PDE solving (prediction) and inverse PDE solving (discovering hidden mechanisms) in a unified framework.\n(ii) Interpretability: The discovered kernel provides meaningful physical interpretation, addressing a common criticism of black-box neural network approaches.\n(iii) Theoretical foundation: The authors provide a theoretical analysis showing how the attention mechanism provides a space of identifiability for kernels, helping to resolve ill-posed inverse problems. The authors also provide, as part of their approach, extensibility to various forms of operators, including those with radial interaction kernels and heterogeneous interactions.\n(iv) Impact: Novel mechanisms for modeling and learning about physical phenomena with a focus on interpretability is a significant step in the right direction for me. I am unsure that this will directly lead to contributions in the foundation model space (I think it's a bit of a stretch -- see weaknesses for details). However, I can see that this work can be gainfully leveraged to drive insights for foundation model building.\n\nMy reading of the analysis provided also yielded the following strengths (specifically regarding the formal rigor of the proposed approach):\n\nThe authors present a general form of the operator in Eq. (3), which encompasses many physical systems. This generality is a strength of the approach. The kernel map (Eq. 4-7) uses a multi-layer attention model to learn the mapping from data pairs to kernel estimations. The use of both input (u) and output (f) in the attention mechanism is novel and crucial for addressing the inverse problem. The derivation of a continuum limit of the discrete attention model (Eq. 8) provides insight into the operator's behavior as the number of data points increases.\nLemmas 4.1 and 4.2 provide valuable insights into the limit behavior of the attention-parameterized kernel and the space of identifiability for kernels. This analysis helps to explain why the method can effectively solve ill-posed inverse problems.", "weaknesses": "I find only one glaring weakness, which is likely to surface during the implementation of this approach for physical systems of appreciable size and complexity: The method involves multiple integrations and matrix operations, which may be computationally expensive for large-scale problems. A discussion of computational efficiency and scalability would be beneficial. On a related note, It is also not clear to me if this method is data-efficient.\n\nA minor weakness I see is that I do not find that this method is easily integrated with current approaches to physics-informed machine learning, e.g., through simple extensions to pre-existing libraries. A brief discussion on this very practical aspect might strengthen the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments and suggestions. Our response:\n\n**Number of trainable parameters and computational complexity**: Denoting $(N,d,d_k,L,n_{train})$ as the size of the spatial mesh, the number of data pairs in each training model, the column width of the query/key weight matrices $W^Q$ and $W^K$, the number of layers, and the number of training models, the number of trainable parameters in a discrete NAO is of size $O(L\\times d\\times d_k+N^2)$. For continuous-kernel NAO, taking a three-layer MLP as a dense net of size $(d, h_1, h_2, 1)$ for the trainable kernels $W^{P,u}$ and $W^{P,f}$ for example, its the corresponding number of trainable parameters is $O(d\\times h_1+h_1\\times h_2)$. Thus, the total number of trainable parameters for a continuous-kernel NAO is $O(L\\times d\\times d_k+d\\times h_1+h_1\\times h_2)$.\n\nThe computational complexity of NAO is quadratic in the length of the input and linear in the data size, with $ O([d^2(3N+d_k) + 6N^2d] Ln_{train}) $ flops in each epoch in the optimization. It is computed as follows for each layer and the data of each training model: the attention function takes $ O(d^2d_k+  2Nd^2+4N^2d )$ flops, and the kernel map takes $O(d^2d_k+  Nd^2+2N^2d)$ flops; thus, the total is $O(d^2(3N+d_k) + 6N^2d)$ flops. In inverse PDE problems, we generally have $d\\ll N$, and hence the complexity of NAO is $O(N^2d)$ per layer per sample.\n\n\nCompared with other methods, it is important to note that NAO solves both forward and ill-posed inverse problems using multi-model data. Thus, we don't compare it with methods that solve the problems for a single model data, for instance, the regularized estimator in Appendix B. Methods solving similar problems are the original attention model [1], convolution neural network (CNN), and graph neural network (GNN). As discussed in [1], these models have a similar computational complexity, if not any higher. In particular, the complexity of the original attention model is $O(N^2d)$, and the complexity of CNN is $O(kNd^2)$ with $k$ being the kernel size, and a full GNN is of complexity $O(N^2d^2)$.\n\n[1] Vaswani, Ashish, et al. \"Attention is all you need.\" Advances in neural information processing systems 30 (2017).\n\n**Data efficiency**: The minimal data for NAO to perform well requires a balance between multiple factors, including the intrinsic dimension (smoothness) of the kernel, the token size $d$ for each sample, the rank $d_k$ of the weight matrices, the number of samples in each task, and the diversity of training tasks. Thus, it is difficult to have a theoretical guarantee. In the original manuscript, with all other factors fixed, we have tested the dependence on the token dimension and the rank of the weight matrices in Table 1. To further elaborate the importance of  sample number and task diversity, in Table 1 of the one-page PDF we have further perform tests on three settings: 1) diverse tasks; 2) single task with diverse function pairs; and 3) diverse tasks with a reduced number of samples. In particular, in the diverse task setting, the training dataset is constructed based on 21 kernels of three groups, with $15$ samples on each kernel:\n\nsine-type kernels: $\\gamma_\\eta(r)=\\exp(-\\eta r)\\sin(6r)\\mathbf{1}_{[0,11]}(r)$, $\\eta=1,2,3,4,6,7,8$.\n\ncosine-type kernels: $\\gamma_\\eta(r)=\\frac{10-r}{20}\\cos(\\eta r)(10-r)\\mathbf{1}_{[0,10]}(r)$, $\\eta=0,1,2,3,4,5,6$.\n\npolynomial-type kernels: $\\gamma_\\eta(r)=\\exp(-0.1 r)p_\\eta$( $\\frac{r-10}{10}$ ), $\\eta=1,2,3,4,5,6,7$, where $p_\\eta$ is the degree-$\\eta$ Legendre polynomial.\n\nIn the ''diverse task'' setting, all $315$ samples are employed in training. In the ''single task'' setting, only the first sine-type kernel:\n\n$\\gamma_1(r)=\\exp(-r)\\sin(6r)\\mathbf{1}_{[0,11]}(r)$\n\nis considered as the training task, with all $105$ samples on this task. Then, in the ''fewer samples'' setting, the training dataset still consists of all $21$ tasks but with only $5$ samples on each task. For the purpose of testing, besides the in-distribution (ID) and out-of-distribution (OOD) tasks in the original manuscript, we have added an additional (OOD) task with Gaussian-type kernel $\\gamma_{ood2}(r)=\\frac{\\exp(-0.5r^2)}{\\sqrt{2\\pi}}$, which is substantially different from all training tasks.\n\nAs shown in the results in Table 1 of the one-page PDF, for the ID test task, a small number of training tasks are sufficient, and the model performs well in all except the ''single task'' setting. When the test task has a larger discrepancy from the training tasks, as in the OOD test2 setting, it becomes necessary to increase the training task diversity, and the required rank size $d_k$ should also increase to obtain better model expressivity. When comparing the  ''fewer samples'' setting with the ''sin only'' setting,  the former has better task diversity but fewer samples per task. One can see that the performance on ID test deteriorates but the performance on OOD test2 improves. This observation highlights the balance between the diversity of tasks and the number of training samples.\n\n**Combination with physics-based modeling approaches**: When the physics structure is known, NAO can directly encodes such information in its architecture, such as the kernel structure in the nonlocal diffusion model in Example 1. By doing so, we anticipate to explore the kernel in a lower-dimensional manifold, which reduces the requirement on the size of training data. \n\nOn the other hand, NAO could also help physics-based modeling approaches by providing a surrogate foundation model when there is a large class of physical models to be estimated and simulated. In this setting, NAO can be first trained using data generated by these models, and then used as an efficient surrogate model. It will be useful for sampling and uncertainty quantification tasks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "Z0TcmhiNc3", "reviewer": "Reviewer_eNuS", "summary": "This paper explores using large-language deep-neural network to build a foundation physical model that can solve inverse PDE problem. It proposes Nonlocal Attention Operator model with a kernel map that is dependent on both input and output functions. The intermediate proofs help people understand the underlying math principles behind the attention mechanism, especially the relationship between attention and integration. The paper also does experiments on radial kernel learning, solution operator learning and heterogeneous material learning, all with decent results.", "questions": "- is it possible for you publish your source codes on Github/Gitlab so that other people can use your framework for practical applications and researchers can expand on top of your work?", "limitations": "- since the equations are quite complex, replication of the piece of work might not be consistent among different people, thus, lack of open-source code will limit its use/adoption in practical applications.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- mathematical proof is comprehensive\n- ablation study has compared performance among several variants\n- experiments are adequate to show the effectiveness of the model in different domains", "weaknesses": "- there are very few comparisons with other papers on experimental results", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time and constructive suggestions. Our source code, trained models, and datasets will be made publicly available on Github upon acceptance of the paper, so as to guarantee reproducibility of all results. Our other responses:\n\n**Limited baselines**: We did not compare much with other work because very few existing models are capable of discovering hidden physics or the governing physical parameters directly from data. Beyond the direct application of the attention model as in the discrete NAO setting, the most relevant work to our best knowledge is the adaptive Fourier neural operator that adopted a convolution-based attention mechanism, which we compared to in the ablation study. However, it is unable to learn a good solution across tasks (as indicated by the large errors in the in-distribution and out-of-distribution tests), nor can it discover physics interpretation for the inverse problem.\n\nTo further address the reviewer's concern, we have added an autoencoder model as an additional baseline in Example 1, see Table 1 of the one-page PDF. In this setting, an encoder is constructed to map the concatenated data pairs [$u_{1:d}$,$f_{1:d}$] to the hidden kernel K[$u_{1:d}$,$f_{1:d}$] as the latent representation, and then in the decoder the kernel is multiplied with $u$ following the last step of Eq.(6): K[$u_{1:d}$,$f_{1:d}$]$u_{1:d}$-> $f_{1:d}$. We note that the concatenated data pairs [$u_{1:d}$,$f_{1:d}$] is of size $O(Nd)$ as the input of the encoder, and the kernel $K$ is of size $N^2$ as the output. Hence, even a linear encoder contains around $N^3d$ trainable parameters. When the autoencoder model is of a similar size to NAO, its expressivity is very limited.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ot6DTiO89n", "reviewer": "Reviewer_BrBK", "summary": "This work presents a new neural network for tackling PDEs' forward and backward problems, especially for small physics systems. The forward problem simulates the outcome of a physics system given the inputs, while the backward problem estimates the parameters of a physics system given the inputs and outputs. The key idea of the presented method is to utilize the underlying common rules of multiple physics systems. This work found that using an attention layer to model multiple physics systems can significantly help the model solve the backward problem, also known as the discovery problem. The idea is intuitive: looking at multiple similar rules is helpful for discovering a new rule.", "questions": "For Section 5.1, consider including a broader range of tasks and different types of kernels. It would also be beneficial to incorporate more OOD testing. Another related question is: Given that multiple training tasks come from the same family, could this be equivalent to a weaker version of learning more data pairs from a single system?\n\nI am unsure whether the reported number of data pairs refers to the total data pairs across different tasks or for each individual task; I suspect it is the latter. If we provide more data pairs from a single task, could it achieve a similar level of performance? Conversely, if we increase the number of tasks and reduce the data pairs for each task, does this improve the generalization ability?", "limitations": "Agree with the limitation section.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The idea is interesting and clever. The nature of sharing knowledge among multiple systems is somewhat related to in-context learning in LLMs.", "weaknesses": "The experiments could be more extensive and solid, especially in supporting the main contribution. Although the experiment on the backward problem (discovery) reports the OOD results, the OOD setting is not far from the in-domain data, which raises concerns about the generalization ability.\n\nThe paper's writing is not very friendly to those with a deep learning background. The content around background knowledge and the connections between the attention mechanism and kernel methods is somewhat tangential to the main storyline. The idea is straightforward, so I would prefer a clean and compact narrative. However, this is not a major concern since I only have a deep learning background.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time and constructive suggestions. We have generated more families of kernels, and performed additional tests on three different training settings (diverse tasks, single task with diverse function pairs, and diverse tasks with a reduced number of samples). Additionally, we have created an OOD test task with a substantially different kernel from the trained ones (i.e., Gaussian-type kernel versus sine, cosine and polynomial-type kernels). All additional test results can be found in the table of the one-page PDF file. We will also modify the manuscript to make the narrative cleaner and more compact. Our other responses:\n\n\n**Broader range of tasks for different types of kernels**: Per the reviewer's suggestion, we have added another two types of kernels into the training dataset. The training dataset is now constructed based on 21 kernels of three groups, with $15$ samples on each kernel:\n\nsine-type kernels: $\\gamma_\\eta(r)=\\exp(-\\eta r)\\sin(6r)\\mathbf{1}_{[0,11]}(r)$, $\\eta=1,2,3,4,6,7,8$.\n\ncosine-type kernels: $\\gamma_\\eta(r)=\\frac{10-r}{20}\\cos(\\eta r)(10-r)\\mathbf{1}_{[0,10]}(r)$, $\\eta=0,1,2,3,4,5,6$.\n\npolynomial-type kernels: $\\gamma_\\eta(r)=\\exp(-0.1 r)p_\\eta$($\\frac{r-10}{10}$), $\\eta=1,2,3,4,5,6,7$, where $p_\\eta$ is the degree-$\\eta$ Legendre polynomial.\n\nBased on this enriched dataset, beyond the original ''sin only'' setting, three additional settings are considered in Table 1 of the attached PDF. In part II of Table 1 (denoted as ''sin+cos+poly''), we consider a ''diverse task'' setting, where all $315$ samples are employed in training. Then, in part III we consider a ''single task'' setting, where only the first sine-type kernel\n\n$\\gamma(r)=\\exp(-r)\\sin(6r)\\mathbf{1}_{[0,11]}(r)$\n\nis considered as the training task, with all $105$ samples on this task. Lastly, in part IV we demonstrate a ''fewer samples'' setting, where the training dataset still consists of all $21$ tasks but with only $5$ samples on each task. For the purpose of testing, besides the in-distribution (ID) and out-of-distribution (OOD) tasks in the original manuscript, we have added an additional (OOD) task with Gaussian-type kernel $\\gamma_{ood2}(r)=\\frac{\\exp(-0.5r^2)}{\\sqrt{2\\pi}}$, which is substantially different from all training tasks.\n\nAs shown in Table 1, considering diverse datasets helps in both OOD tests (kernel error reduced from 9.14% to 6.92% in OOD test 1, and from 329.66% to 10.48% in OOD test 2), but not for the ID test (kernel error slightly increased from 4.03% to 5.04\\%). We also note that since the ''sin only'' setting only sees systems with the same kernel frequency in training, it does not generalize to the OOD test 2. In this test, including more diverse training tasks becomes necessary. Then, as the training tasks become more diverse, the intrinsic dimension of kernel space increases, requiring a larger weight matrices rank size $d_k$. When comparing the  ''fewer samples'' setting with the ''sin only'' setting,  the former has better task diversity but fewer samples per task. One can see that the performance on ID test deteriorates but the performance on OOD test2 improves. This observation highlights the balance between the diversity of tasks and the number of training samples per task.\n\n**Difference from learning more data pairs from a single system**: Although all 7 sine-type kernels are corresponding to the same frequency, they are very different from each other. If we train the model from a single kernel, the model does not generalize at all. Then, the test error on validation task does not improve during training, and the model always predicts an all-zero kernel. As shown in the ''Single task'' setting of attached Table 1, the kernel and operator error are always around 100% for all test tasks.\n\n\n**Reported number of data pairs**: When we talk about the data pairs for each task, as in Example 1, the reported number of data pairs refers to the total number of data pairs of each task. Then, we form a training sample of each task by taking $d$ pairs from this task. When taking the token size $d=302$, each task contains $\\lfloor\\frac{4530}{302}\\rfloor=15$ samples. When having 7 training tasks corresponding to 7 different sine-type kernels, there are $105$ samples in total. Then, when discussing about the total number of data pairs, as in Examples 2-3, we are reporting it across all tasks. In operator learning settings, we note that the permutation of function pairs in each sample should not change the learned kernel, i.e., one should have K[$u_{1:d}$,$f_{1:d}$]=K[$u_{\\sigma(1:d)}$,$f_{\\sigma(1:d)}$], where $\\sigma$ is the permutation operator. Hence, we have augmented the training dataset by permuting the function pairs in each task. Taking example 3 for instance, with $500$ microstructures (tasks) and $200$ function pairs per task, we have randomly permuted the function pairs and taken $100$ function pairs for $100$ times per task. As such, we have created the total 50000 samples (45000 for training and 5000 for test) in Example 3. We will clarify this in the revised manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We thank the reviewers for the constructive comments and for recognizing the novelty and broad impact of our work, the soundness and solid theoretical foundation of our formulation, the effectiveness and meaningful physical interpretability of our model, as well as the comprehensive experimental studies. In our response, we have answered all the questions raised by the reviewers and clarified potential misunderstandings. Moreover, we have performed additional tests on: an additional baseline (autoencoder), three different training settings (diverse tasks, single task with diverse function pairs, and diverse tasks with a reduced number of samples), and an OOD test task with a substantially different kernel from the trained ones (i.e., Gaussian-type kernel versus sine, cosine and polynomial-type kernels). All additional test results can be found in the table of the one-page PDF file. Lastly, to provide further physical interpretability of the discovered kernel, we demonstrate NAO's capability in recovering the underlying microstructure in the Darcy's problem (Example 2). A comparison of the true microstructure and the discovered one can be found in Figure 2 of the one-page PDF file.

We sincerely appreciate the time and efforts of the AC and all reviewers. Your comments have helped us improve the quality of our manuscript, and we will incorporate them into the revised manuscript. We will also release the source codes, the trained models, and the datasets upon paper acceptance to guarantee reproducibility of all experiments. We are more than happy to answer any follow-up questions during the discussion period.

Meta-review: This paper presents an interesting study that utilizes a novel attention-based neural architecture for modeling complex physical systems. The authors introduce an attention mechanism specifically tailored to address such physical problems involving PED systems. The paper includes both theoretical discussions and strong empirical results. The reviewers unanimously agree that this is a solid contribution with the potential to bring new insights to the field.
Decision: Accept (spotlight)