Title: Nonlocal Attention Operator: Materializing Hidden Knowledge Towards Interpretable Physics Discovery
Venue: NeurIPS 2024
Accepted: yes
Keywords: Foundation Model

Neural Operators

Inverse PDE Problems

Physical Modeling
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

Abstract

Despite the 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. We empirically demonstrate the advantages of 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 a new perspective towards understanding the attention mechanism. Our code and data accompanying this paper are available at https://github.com/fishmoon1234/NAO .

1 Introduction

The interpretability of machine learning (ML) models has become increasingly important from the security and robustness standpoints [Rudin et al., 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 al., 2022, Ferrari and Willcox, 2024]. 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-

ter field and mechanism presents an ill-posed unsupervised learning task. That means, even if an ML model can serve as a good surrogate to predict the corresponding loading field from a given deformation field, its inference of the material parameters can still drastically deteriorate.

To discover an interpretable mechanism for physical systems, a major challenge is to infer the governing laws of these systems that are often high- or infinite-dimensional, from data that are comprised of discrete measurements of continuous functions. Therefore, a data-driven surrogate model needs to learn not only the mapping between input and output 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 forward problem, whereas inferring the underlying

(

u

1

,

f

1

)

(

u

2

,

f

2

)

.

.

.

(

u

n

,

f

n

)

W QK

Kernel space of

identifiability

Integrate

Training

Downstream application

u

test

Kernel (mechanism)

K[

u

1:n

,

f

1:n

]

f

test

W P,u

W P,f

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 et al., 2019]. To tackle this challenge, many deep learning methods have recently been proposed as inverse PDE solvers [Fan and Ying, 2023, Molinaro et al., 2023, Jiang et al., 2022, Chen et al., 2023]. The central idea is to incorporate prior information into the learning scheme, in the form of governing PDEs [Yang et al., 2021, Li et al., 2021], regularizers [Dittmer et al., 2020, Obmann et al., 2020, Ding et al., 2022, Chen et al., 2023], or additional operator structures [Uhlmann, 2009, Lai et al., 2019, Yilmaz, 2001]. However, such prior information is often either unavailable or problemspecific in complex systems. As a result, these methods can only solve the inverse problem for a particular system, and one has to start from scratch when the system varies (e.g., when the material of the specimen undergoes degradation in a material modeling task).

In this work, we propose Nonlocal Attention Operator (NAO) , a novel attention-based neural operator architecture to simultaneously solve both forward and inverse modeling problems. Neural operators (NOs) [Li et al., 2020a,c] 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 et al. [2022], Liu et al. [2024a,b, 2023], Ong et al. [2022], Cao [2021], Lu et al. [2019, 2021], Goswami et al. [2022], Gupta et al. [2021] and references therein. However, most NOs focus on providing an efficient surrogate for the underlying physical system as a forward solver. They are often employed as black-box universal approximators but lack interpretability of the underlying physical laws. In contrast, the key innovation of NAO is that it introduces a kernel map based on the attention 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 in 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 a kernel map using the attention mechanism and simultaneously solves both the forward and inverse problems. The kernel map, whose parameters extract the global information about the kernel from multiple systems, efficiently infers resolution-invariant 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 mechanism discovery (inverse PDE).

We propose a novel neural operator architecture NAO, based on the principle of contextual discovery from input/output function pairs through a 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 provide the space of identifiability for the kernels from the training data, which reveals its ability to resolve ill-posed inverse PDE problems.

Weconduct experiments on zero-shot learning to new and unseen physical systems, demonstrating 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 ubiquitous in scientific and engineering applications [Ghaboussi et al., 1998, Liu et al., 2024c, Ghaboussi et al., 1991, Carleo et al., 2019, Karniadakis et al., 2021, Zhang et al., 2018, Cai et al., 2022, Pfau et al., 2020, He et al., 2021, Besnard et al., 2006]. In many applications, the underlying governing laws are unknown, hidden in data to be revealed by physical models. Ideally, these models should be interpretable for domain experts, who can then use these models to make further predictions and expand the understanding of the target physical system. Also, these models should be resolution-invariant . Neural operators are designed to learn mappings between infinite-dimensional function spaces [Li et al., 2020a,b,c, You et al., 2022, Ong et al., 2022, Cao, 2021, Lu et al., 2019, 2021, Goswami et al., 2022, Gupta et al., 2021]. As a result, NOs provide a 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 on providing an efficient surrogate for the underlying physical system as a forward PDE solver. They are often employed as blackbox universal approximators without interpretability of the underlying physical laws. Conversely, several deep learning methods have been proposed as inverse PDE solvers [Fan and Ying, 2023, Molinaro et al., 2023, Jiang et al., 2022, Chen et al., 2023], aiming to reconstruct the parameters in the PDE from solution data. Compared to the forward problem, the inverse problem is typically more challenging due to its ill-posed nature. To tackle the ill-posedness, many NOs incorporate prior information, in the form of governing PDEs [Yang et al., 2021, Li et al., 2021], regularizers [Dittmer et al., 2020, Obmann et al., 2020, Ding et al., 2022, Chen et al., 2023], or additional operator structures [Uhlmann, 2009, Lai et al., 2019, Yilmaz, 2001]. To our knowledge, our NO architecture is the first that solves both the forward (prediction) and inverse (discovery) problems simultaneously.

Attention mechanism. Since 2017, the attention mechanism has become the backbone of state-ofthe-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 [Vaswani et al., 2017]. Then, the tokens are spatially mixed to obtain the layer output. Based on the choice of mixers, attention-based models can be divided into three main categories: discrete graph-based attentions [Child et al., 2019, Ho et al., 2019, Wang et al., 2020, Katharopoulos et al., 2020], MLP-based attentions [Tolstikhin et al., 2021, Touvron et al., 2022, Liu et al., 2021], and convolution-based attentions [Lee-Thorp et al., 2021, Rao et al., 2021, Guibas et al., 2021, Nekoozadeh et al., 2023]. While most attention models focus on discrete mixers, it is proposed in Guibas et al. [2021], Nekoozadeh et al. [2023], Tsai et al. [2019], Cao [2021], 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 normalization 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 [Li et al., 2022], a hierarchical transformer for learning multiscale problems [Liu et al., 2022], and a heterogeneous normalized attention with a geometric gating mechanism [Hao et al., 2023] to handle multiple input features. In particular, going beyond solving a single PDE, the foundation model feature of attention mechanisms has been applied towards solving multiple types of PDEs within a specified context in Yang and Osher [2024], Ye et al. [2024], Sun et al. [2024], 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 are described by a class of operators mapping from input functions u ∈ X to output functions f ∈ Y . Our goal is to learn the common physical law, in the form of operators L K : X → Y with system-dependent kernels K :

Here X and Y are Banach spaces, ϵ denotes an additive noise describing the discrepancy between the ground-truth operator and the optimal surrogate operator, and K is a kernel function representing the nonlocal spatial interaction. As such, the kernel provides the knowledge of its corresponding system, while (1) offers a zero-shot prediction model for new and unseen systems.

To formulate the learning, we consider n train training datasets from different systems, with each dataset containing d u function pairs ( u, f ) :

In practice, the data of the input and output functions are on a spatial mesh { x k } n x k =1 ⊂ Ω ⊂ R d x . The n train models with kernels { K η } n train η =1 correspond to different material micro-structures or different parametric settings. As a demonstration, we consider models for heterogeneous materials with operators in the form where g [ u ]( y ) is a functional of u determined by the operator; for example g [ u ]( y ) = u ( y ) in Section 5.3. Our approach extends naturally to other forms of operators, such as those with radial interaction kernels in Section 5.1 and heterogeneous interaction in the form of L K [ u ]( x ) = ∫ Ω K ( x, y ) g [ u ]( x, y ) dy . Additionally, for simplicity, we consider scalar-valued functions u and f and note that the extension to vector-valued functions is trivial.

Remark: Such an 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 a forward problem. When considering a single physical system and taking K in (3) as an input-independent kernel, classical NOs can be obtained for forward PDE-solving tasks [Li et al., 2020c, Guibas et al., 2021] and governing law learning tasks [You et al., 2021, Jafarzadeh et al., 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 a kernel map constructed using the attention mechanism. It maps from data pairs to an estimation of the underlying kernel. The kernel map

has parameters θ estimated from the training dataset (2). As such, it maps from the token ( u 1: d , f 1: d ) of the dataset { ( u i , f i ) } d u i =1 to a kernel estimator, acting as an inverse PDE solver.

Amajor innovation of this kernel map is its dependence on both u and f through their tokens. Thus, our 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 i , f i ) } d u i =1 to tokens ( u 1: d , f 1: d ) according to the operator in (3) by where d = d u and N = n x , assuming that g [ u ] has a spatial mesh { y k = x k } N k =1 .

Then, our discrete ( L +1) -layer attention model for the inverse PDE problem writes:

Here, θ l = ( W Q l ∈ R d × d k , W l K ∈ R d × d k ) and the attention function is

The kernel map is defined as

where θ := { W P,u ∈ R N × N , W P,f ∈ R N × N , { W Q l , W l K } L l =1 } are learnable parameters. Here, we note that d k only controls the rank bound when lifting each point-wise feature via W Q l ( W l K ) ⊤ . In the following, we denote W QK l := 1 √ d k W Q l ( W l K ) ⊤ , which characterizes the (trainable) interaction of d input function pair instances.

3.2 Nonlocal Attention Operator in continuum limit

As suggested by Cao [2021], we take the activation function σ as a linear operator. Then, noting that the matrix multiplication in (6) is a Riemann sum approximation of an integral (with a full derivation in Appendix A), we propose the nonlocal attention operator as the continuum limit of (6):

in which the integration is approximated by the Riemann sum in our implementation. Here,

We learn the parameters θ by minimizing the following mean squared error loss function:

The performance of the model is evaluated on test tasks with new and unseen kernels D test = { ( u test i ( x ) , f test i ( x )) } d i =1 based on the following two criteria:

Here, L 2 ( ρ ) is the empirical measure in the kernel space as defined in Lu et al. [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 to estimators significantly better than those using a single dataset. To formally understand this mechanism, we analyze a shallow 2-layer NAO in the next section.

4 Understanding the attention mechanism

To facilitate further understanding of the attention mechanism, we analyze the limit of the twolayer attention-parameterized kernel in (7) 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

which is the radial nonlocal kernel in Sec.5.1.

4.1 Limit of the two-layer attention-parameterized kernel

We show that, as the number N and the spatial mesh n x approach infinity, the limit of the two-layer attention-parameterized kernel is a double integral. Its proof is in Appendix B.1.

For simplicity, we assume that the dataset in (2) has d u = 1 and n train = 1 with a uniform mesh { x j } n x j =1 . We define the tokens by

where d = n x and { r k } N k =1 is the spatial mesh for K 's independent variable r ∈ [0 , δ ] .

Lemma 4.1. Consider the two-layer attention model in (6) -(7) with bounded parameters. For each d and N , let { x j } d j =1 and { r k } N k =1 be uniform meshes of the compact sets Ω and [0 , δ ] , and let { A j } d j =1 be the resulting uniform partition of Ω . Assume that g [ u ] in (12) is continuous on [0 , δ ] × Ω . Then,

where W QK ( x, y ) = lim d →∞ ∑ d j,j ′ =1 W QK [ j, j ′ ] 1 A j × A j ′ ( x, y ) is the scaled L 2 (Ω × Ω) limit of the parameter matrix W QK [ j ′ , j ] = ∑ d k l =1 W Q [ j, l ] · W K [ j ′ , l ] and W P,u ( r ) = lim N →∞ ∑ N k =1 W P,u [ k ] 1 [ r k -1 ,r k ) ( r ) . 4.2 Space of identifiability for the kernels

For a 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 (2) are sampled from continuous functions { u η i } d u ,n train i,η =1 with a compact support. Then, the function space the loss function in (9) has a unique minimizer K ( s ) = K [ u η 1: d , f η 1: d ; θ ]( s ) is the closure of a data-adaptive RKHS H G with a reproducing kernel ¯ G determined by the training data:

where ρ ′ is the density of the empirical measure ρ defined by

The above space is data-adaptive since the integral kernel ¯ G depends on data. It characterizes the information in the training data for estimating the nonlocal kernel K ( s ) = K [ u η 1: d , f η 1: d ; θ ]( 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 K ( s ) is not the kernel map. The minimizer is a fixed estimator for the training dataset and does not provide any information for estimating the kernel from another dataset.

and the function G is defined by G ( r, s ) := 1 n train d ∑ n train η =1 ∑ d u i =1 ∫ Ω g [ u η i ]( r, x ) g [ u η i ]( s, x ) dx.

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 classical 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 deconvolution to estimate the kernel. Thus, regularization is necessary, where two main challenges present: (i) the selection of a 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 [ u 1: d , f 1: d ; θ ]( s ) , with the parameter θ estimated from the training datasets, acts on the token ( u 1: d , f 1: d ) of ( u, f ) to provide an estimator. It passes the prior information 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. In Section B.2, we further show that a regularized estimator depends nonlinearly on the data pair ( u, f ) . In particular, similar to Lemma 4.1, there is an RKHS determined by the data pair ( u, f ) . The regularized estimator suggests that the kernel map can involve a component quadratic in the feature g [ u ] , similar to the limit form of the attention model in Lemma 4.1.

5 Experiments

We assess the performance of NAO on a wide range of physics modeling and discovery datasets. Our evaluation focuses on several key aspects: 1) we demonstrate the merits of the continuous and linear attention mechanism, compare the performance with the baseline discrete attention model (denoted as Discrete-NAO), the softmax attention mechanism (denoted as Softmax-NAO), NAO with input on u only (denoted as NAO-u), the convolution-based attention mechanism (denoted as AFNO [Guibas et al., 2021]), and an MLP-based encoder architecture that maps the datum [ u 1: d , f 1: d ] directly to a latent kernel (denoted as Autoencoder); 2) we measure the generalizability, in particular, the zero-shot prediction performance in 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 W P,u and W P,f with a 3-layer MLP with hidden dimensions (32 , 64) and LeakyReLU activation. Experiments are conducted on a 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

In this example, we consider the learning of nonlocal diffusion operators, in the form:

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 a radial kernel γ η . It finds broad physical applications in describing fracture mechanics [Silling, 2000], anomalous diffusion behaviors [Bucur et al., 2016], and the homogenization of multiscale systems [Du et al., 2020].

0

2

4

6

8

|x-y|

0

0.1

0.2

0.3

0.4

K(|x-y|)

Ground-truth in-distribution test kernel

Discrete-NAO: kernel error =55.24%

NAO, same resolution: kernel error=22.65%

NAO, cross resolution: kernel error=20.79%

0

2

4

6

8

|x-y|

0

0.2

0.4

0.6

0.8

1

1.2

K(|x-y|)

Ground-truth out-of-distribution test kernel

Discrete-NAO: kernel error =64.23%

NAO: kernel error=33.77%

In this context, our goal is to learn the operator L as well as to discover the hidden mechanism, namely the kernel K [ u 1: d , f 1: d ; θ ]( x, y ) = γ η ( | y -x | ) . In the form of the operator in (12), we have K ( r ) = γ η ( r ) and g [ u ]( r, x ) = u ( x + r ) + u ( x -r ) -2 u ( x ) for r ∈ [0 , δ ] .

To generate the training data, we consider 7 sine-type kernels

Here, η denotes task index. We generate 4530 data pairs ( g η [ u ] , f η ) with a fixed resolution ∆ x = 0 . 0125 for each task, where the loading function L γ η [ u η ] = f η is computed by the adaptive GaussKronrod 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 (17) with η = 5 (denoted as the 'in-distribution (ID) test' system), and the other following a different rule (denoted as the 'out-of-distribution (OOD) test1' system):

Both the operator error (10) and the kernel error (11) are provided in Table 1. While the former measures the error of the learned forward PDE solver (i.e., 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 an ablation study on 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 k = 10 , and data resolution ∆ x = 0 . 0125 . When comparing the operator errors, both Discrete-NAO 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 to provide a surrogate of L γ ood with < 10% error, who outperforms its discrete mixer counterpart by 68 . 62% , indicating that NAO learns a more generalizable mechanism. This argument is further affirmed when comparing the kernel errors, where NAO substantially outperforms all baselines by at least 81 . 39% in the ID test and 65 . 21% in the OOD test. This study verifies our analysis in Section 4: 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 k = 10 to d k = 5 and d k = 20 . Note that d k determines the rank bound of W QK , the matrix that characterizes the interaction between different data pairs. Discrete-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 low-dimension in the kernel space.

Alleviating 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 to d = 30 , making it more illposed 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 error also increases monotonically in the ID operator test, but there exists no consistent pattern in other test criteria. Figure 2 shows the learned test kernels in both the ID and OODtasks. It shows that Discrete-NAO learns highly oscillatory kernels, while our continuous NAO only has a discrepancy near | x -y | = 0 . Note that when | x -y | = 0 , we have u ( y ) -u ( x ) = 0 in the ground-truth operator (16), and hence the kernel value at this point does not change the operator

ur

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 ∆ x = 0 . 0125 on a dataset corresponding to ∆ x = 0 . 025 , and plot the results in Figure 2 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 et al. [2020a], the high-fidelity synthetic simulation data for this example are described by the Darcy flow. Here, the physical domain is Ω = [0 , 1] 2 , b ( x ) is the permeability field, and the Darcy's equation has the form:

In this context, we aim to learn the solution operator of Darcy's equation and compute the pressure field p ( x ) . We consider two study scenarios. 1) g → p : each task has a fixed microstructure b ( x ) , and our goal is to learn the (linear) solution operator mapping from each loading field g to the corresponding solution field p . In this case, the kernel K acts as the Green's function of (19), and can be approximated by the inverse of the stiffness matrix. 2) b → p : each task has a fixed loading field g ( x ) , and our goal is to learn the (nonlinear) solution operator mapping from the permeability field b to the corresponding solution field p .

We report the operator learning results in Table 2, 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 groundtruth inverse of stiffness matrix in Figure 3. Although both Discrete-NAO and NAO capture the major pattern, the kernel from Discrete-NAO again shows a spurious oscillation in regions where the ground-truth kernel has zero value. On the other hand, by exploring the kernel map in the integrated knowledge space, the learned kernel from NAO does not exhibit such spurious modes.

To demonstrate the physical interpretability of the learned kernel, in the first row of Figure 4 we show the ground-truth microstructure b ( x ) , a test loading field instance g ( x ) , and the corresponding solution p ( x ) . By taking the summation of the kernel strength on each row, one can discover the interaction strength of each material point x with its neighbors. As this strength is related to the permeability field b ( x ) , the underlying microstructure can be recovered accordingly. In the bottom row of Figure 4, we demonstrate the discovered microstructure of this test task. We note that the discovered microstructure is smoothed out due to the continuous setting of our learned kernel (as shown in the bottom left plot), and a thresholding step is performed to discover the two-phase microstructure. 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(x)

p(x)

1

1

1

0.8

0.8

200

0.8

0.2

0.6

0.6

100

0.6

0

0.4

0.4

0

0.4

-0.2

-100

-0.4

0.2

0.2

0.2

-200

-0.6

0

0

0

0

0.5

1

0

0.5

1

0

0.5

1

{|K[u,f](x,y)| dy

Recovered b(x)

1

1

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0

0

0

0.5

0

0.5

the applied Dirichlet-type boundary condition ( p ( x ) = 0 on ∂ Ω ) in all samples, which leads to the measurement pairs ( p ( x ) , g ( x )) containing no information near the domain boundary ∂ Ω and makes it 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 using 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 loading/response data pairs are provided. Two generalization scenarios are considered. 1) We mix the 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 to an OOD test. The corresponding results are listed in Table 3, where NAO again outperforms its discrete counterpart.

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 particular, NAO learns the function-to-function mapping based on an integral NO architecture and provides a 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 to new and unseen physical systems. As such, the kernel map explores in the function space of identifiability, resolving the enduring 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-distribution generalizability tasks.

Broader Impacts: Beyond its merits in forward/inverse PDE modeling, our work represents an initial exploration in understanding the attention mechanism in physics modeling, and paves a theoretical path towards building a foundation model in scientific ML.

Limitations: Due to limited computational resource, our experiments focus on learning from a small to medium number ( < 500 ) of similar physical systems. It would be beneficial to expand the coverage and enable learning across different types of physical systems.

Acknowledgments and Disclosure of Funding

S. Jafarzadeh would like to 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 DMS1753031 and the National Institute of Health under award 1R01GM157589-01. Portions of this research were conducted on Lehigh University's Research Computing infrastructure partially supported by NSF Award 2019035. F. Lu would like to acknowledge support by NSF DMS-2238486.

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 in 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, irrevocable, 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 https://www.energy.gov/downloads/doe-public-access-plan.

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):789803, 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. Physicsinformed neural networks (PINNs) for fluid mechanics: A review. Acta Mechanica Sinica , pages 1-12, 2022.

Shuhao Cao. Choose a transformer: Fourier or galerkin. Advances in neural information processing systems , 34:24924-24940, 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. Reviews 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.09854 , 2023.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Genevieve Coorey, Gemma A Figtree, David F Fletcher, Victoria J Snelson, Stephen Thomas Vernon, David Winlaw, Stuart M Grieve, Alistair McEwan, Jean Yee Hwa Yang, Pierre Qian, Kieran 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.01799 , 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 nonlocal effects: Mathematical and computational issues. Contemporary mathematics , 754:115-140, 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-153, 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 , 42(1):105-126, 1998.

Somdatta Goswami, Aniruddha Bora, Yue Yu, and George Em Karniadakis. Physics-informed neural operators. 2022 arXiv preprint arXiv:2207.05748 , 2022.

John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13587 , 2021.

Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differential equations. Advances in neural information processing systems , 34:24048-24062, 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. In International Conference on 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 Biomechanics , 117:110124, 2021.

Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv preprint arXiv:1912.12180 , 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 , 425:116914, 2024.

Hanyang Jiang, Yuehaw Khoo, and Haizhao Yang. Reinforced inverse scattering. arXiv preprint arXiv:2206.04186 , 2022.

George Em Karniadakis, Ioannis 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 rnns: Fast autoregressive transformers with linear attention. In International conference on machine 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-2358, 2019.

James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint arXiv:2105.03824 , 2021.

Emma Lejeune. Mechanical mnist: A benchmark dataset for mechanical metamodels. Extreme Mechanics Letters , 36:100659, 2020.

Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. arXiv preprint arXiv:2205.13671 , 2022.

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03485 , 2020a.

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhattacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. Advances in Neural Information Processing Systems , 33:NeurIPS 2020, 2020b.

Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier NeuralOperator for Parametric Partial Differential Equations. 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-informed neural operator for learning partial differential equations. 2021 arXiv preprint arXiv:2111.03794 , 2021.

Hanxiao Liu, Zihang Dai, David So, and Quoc V Le. Pay attention to mlps. Advances in neural information processing systems , 34:9204-9215, 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 Artificial 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 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:2405.09572 , 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 identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv:1910.03193 , 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):218-229, 2021.

Roberto Molinaro, Yunan Yang, Björn Engquist, and Siddhartha Mishra. Neural inverse operators for solving pde inverse problems. arXiv preprint arXiv: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.12398 , 2023.

Daniel Obmann, Johannes Schwab, and Markus Haltmeier. Deep synthesis regularization of inverse problems. arXiv preprint arXiv:2002.00155 , 2020.

Yong Zheng Ong, Zuowei Shen, and Haizhao Yang. IAE-Net: integral autoencoders for discretization-invariant learning. arXiv preprint arXiv:2203.05142 , 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):033429, 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.

Cynthia Rudin, Chaofan Chen, Zhi Chen, Haiyang Huang, Lesia Semenova, and Chudi Zhong. Interpretable 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. Journal of the Mechanics and Physics of Solids , 48(1):175-209, 2000.

Jingmin Sun, Yuxuan Liu, Zecheng Zhang, and Hayden Schaeffer. Towards a 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 Unterthiner, 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:2426124272, 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: Feedforward networks for image classification with data-efficient training. IEEE Transactions on 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: a unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv:1908.11775 , 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, Łukasz 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 of 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.06523 , 2019.

Liu Yang and Stanley J Osher. Pde generalization of in-context operator networks: A study on 1d scalar nonlinear conservation laws. arXiv preprint arXiv:2401.07364 , 2024.

Liu Yang, Xuhui Meng, and George Em Karniadakis. B-pinns: Bayesian physics-informed neural networks for forward and inverse pde problems with noisy data. Journal of Computational Physics , 425:109913, 2021.

Zhanhong Ye, Xiang Huang, Leheng Chen, Hongsheng Liu, Zidong Wang, and Bin Dong. Pdeformer: Towards a foundation model for one-dimensional partial differential equations. arXiv preprint arXiv:2402.12652 , 2024.

Özdo˘ gan Yilmaz. Seismic data analysis , volume 1. 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 , 374:113553, 2021.

Huaiqian You, Yue Yu, Marta D'Elia, Tian Gao, and Stewart Silling. Nonlocal kernel network (NKN): A stable and resolution-independent deep neural network. Journal of Computational Physics , page arXiv preprint arXiv:2201.02217, 2022.

Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and E Weinan. Deep potential molecular dynamics: a scalable model with the accuracy of quantum mechanics. Physical Review Letters , 120(14):143001, 2018.

Zecheng Zhang. Modno: Multi operator learning with distributed neural operators. arXiv preprint arXiv:2404.02892 , 2024.

A Riemann sum approximation derivations

In this section, we show that the discrete attention operator (6) 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 ∆ x . Denoting the l -th layer out-

That means, d k only controls the rank bound when lifting each point-wise feature via W QK l , while W QK l characterizes the (trainable) interaction of d input function pair instances. Moreover, when α ≤ N , X ( l ) [ α, β ] = g ( l ) β ( x α ) . When α > N , X ( l ) [ α, β ] = f ( l ) β ( x α -N ) . Then,

Similarly,

With the Riemann sum approximation: ∫ Ω p ( x ) dx ≈ ∆ x D ∑ N x k =1 p ( x k ) , one can further reformulate above derivations as:

Therefore, the attention mechanism of each layer is in fact an integral operator after a rescaling:

with the kernel K ( l ) ( x, y ) defined as:

For the L -th layer update, we denote the approximated value of f β ( x α ) as ˜ f β ( x α ) := X out [ α, β ] , then

Hence, a (rescaled) continuous limit of the kernel writes:

B Proofs and connection with regularized estimator

B.1 Proofs

Proof of Lemma 4.1. With { r k } N k =1 , we can write the attention-based kernel in (7) as

Then, as d →∞ is achieved by refining the spatial mesh, viewing the summation in j as Riemann sum,

where the integral exists since g [ u ]( r, x ) W QK ( x, y ) is bounded. Sending also the number of tokens, N , to infinity, we obtain the limit attention model in (14).

Proof of Lemma 4.2. The proof is adapted from Lu et al. [2023, 2022]. Write K ( r ) = K [ u η 1: d , f η 1: d ]( r ) . Notice that the loss function in (9) can be expanded as

where L ¯ G : L 2 ρ → L 2 ρ is the integral operator and K D is the Riesz representation of the bounded linear functional

Thus, the quadratic loss function has a unique minimizer in Null( L ¯ G ) ⊥ .

Meanwhile, since the data pairs are continuous with compact support, the function ¯ G is a squareintegrable reproducing kernel. Thus, the operator L ¯ G is compact and H G = L 1 / 2 ¯ G L 2 ρ . Then, Null( L ¯ G ) ⊥ = H G , where the closure is with respect to L 2 ρ .

B.2 Connection with regularized estimator

Consider the inverse problem of estimating the nonlocal kernel K given a data pair ( u, f ) . In the classical variational approach, one seeks the minimizer of the following loss function

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 is crucial to produce a stable solution.

To connect with the attention-based model, we consider regularizing using an RKHS H W with a square-integrable reproducing kernel W . One seeks an estimator in H W by regularizing the loss with the ∥ K ∥ 2 H W , and minimizes the regularized loss function

The next lemma shows that the regularized estimator is a nonlinear function of the data pair ( u, f ) , where the nonlinearity comes from the kernel G u and the parameter λ ∗ .

Lemma B.1. The regularized loss function in E λ,W ( K ) in (24) is

where G u is defined in (27) . Its minimizer is where L G u and L W are integral operators with integral kernels G u defined in (27) and W ,

λ ∗ is the optimal hyper-parameter controlling the strength of regularization, and K u,f ( r ) = ∫ Ω g [ u ]( r, x ) f ( x ) dx is a function determined by the data ( u, f ) .

When there is no prior information on the regularization, which happens often for the learning of the kernel, one can use the data-adaptive RKHS H G u with the reproducing kernel G u determined by data:

Lu et al. [2023] 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 as a discretization of the continuous inverse problem. Assume that the integrands are compactly supported and when the integrals are approximated by Riemann sums, we can write loss function for discrete K = ( K ( r 1 ) , . . . , K ( r N )) ⊤ ∈ R N × 1 as

where, recalling the definition of the token u 1: d in (5),

The minimizer of this discrete loss function with the optimal hyper-parameter λ ∗ is

̂ K = ( G u + λ ∗ W ) -1 K u,f . In particular, when taking W = G u and using the Neumann series ( λ -1 ∗ G 2 u + I ) -1 = ∑ ∞ k =0 ( -1) k λ -k ∗ G 2 k u , we have

In particular, λ ∗ depends on both G u and K u,f . Hence, the estimator ̂ K is nonlinear in G u and K u,f , and it is important to make the attention depend nonlinearly on the token u 1: d , as in (7).

Proof of Lemma B.1. Since K is radial and noticing that

since ξ ∈ R 1 , we can write

Meanwhile, by the Riesz representation theorem, there exists a function K u,f ∈ L 2 (0 , ∞ ) such that

Combining these two equations, we can write the loss function as

Then, we can write the regularized loss function E λ,W ( K ) as

Selecting the optimal hyper-parameter λ ∗ , which depends on both ( u, f ) and W , and setting the Fréchet derivative of E λ,W over L 2 (0 , ∞ ) to be zero, we obtain the regularized estimator in (26).

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 i } i =1 , 2 with u 1 = cos( x ) 1 [ -π,π ] ( x ) and u 2 ( x ) = sin(2 x ) 1 [ -π,π ] ( x ) . In the 'single task' setting, to create more diverse samples, the single kernel acts on a set of 14 functions: u k = cos( kx ) 1 [ -π,π ] ( x ) , k = 1 , · · · , 7 and u k ( x ) = sin( kx ) 1 [ -π,π ] ( x ) , k = 8 , · · · , 14 . In the ground-truth model, the integral L γ η [ u i ] is computed by the adaptive Gauss-Kronrod 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 ∆ x ∈ 0 . 0125 × { 1 , 2 } , we take values { u i , f i } N i =1 =

0

1

2

3

4

5

6

7

|x-y|

0

0.2

0.4

0.6

0.8

K(|x-y|)

Ground-truth out-of-distribution test1 kernel

NAO (sin only): kernel error =12.16%

NAO (sin+cos+poly): kernel error =6.92%

NAO (sin+cos+poly, less samples): kernel error =24.04%

0

1

2

3

4

5

6

7

|x-y|

0

0.5

1

1.5

K(|x-y|)

Ground-truth out-of-distribution test2 kernel

VAE (sin+cos+poly): kernel error=128.08%

NAO (sin+cos+poly, less parameter): kernel error =94.38%

NAO (sin+cos+poly): kernel error =10.48%

NAO (sin+cos+poly, less sample): kernel error =25.58%

{ u i ( x j ) , f i ( x j ) : x j ∈ [ -40 , 40] , j = 1 , . . . , J } N i =1 , where x j is the grid point on the uniform mesh of size ∆ x . We form a training sample of each task by taking d pairs from this task. When taking the token size d and k function pairs, each task contains ⌊ 2265 d ⌋ × k samples.

More diverse tasks. To further evaluate the generalization capability of NAO as a 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: γ sin η ( r ) = exp( -ηr ) sin(6 r ) 1 [0 , 11] ( r ) , η = 1 , 2 , 3 , 4 , 6 , 7 , 8 .

cosine-type kernels: γ cos η ( r ) = 10 -r 20 cos( ηr )(10 -r ) 1 [0 , 10] ( r ) , η = 0 , 1 , 2 , 3 , 4 , 5 , 6 .

polynomial-type kernels: γ p oly η ( r ) = exp( -0 . 1 r ) p η ( r -10 10 ) 1 [0 , 10] ( r ) , η = 1 , 2 , 3 , 4 , 5 , 6 , 7 , where p η is the degreeη Legendre polynomial.

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 4 and Figure 5. In addition to the original setting corresponding to all 'sine' kernels, in part II of Table 4 (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, γ ( r ) = exp( -r ) sin(6 r ) 1 [0 , 11] ( r ) , 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 a Gaussian-type kernel γ ood 2 ( r ) = exp( -0 . 5 r 2 ) , which is substantially different from all training tasks.

As shown in Table 4, 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 . 66% 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 d k 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 × 1 subjected to Dirichlet boundary conditions. The problem setting is: -∇ ( b ( x ) ∇ p ( x )) = g ( x ) 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 is represented by the location-dependent conductivity b ( x ) . p ( x ) is the source term, and the hydraulic

height g ( x ) is the solution. For each data instance, we solve the equation on a 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 obtained via a Gaussian random field generator. For each microstructure, we solve the Darcy problem considering all 100 source terms, resulting in a dataset of N = 10 , 000 function pairs in the form of { p i ( x j ) , g i ( x j ) } N i =1 , and j = 1 , 2 , · · · , 441 where x j 's 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 [ u 1: d , f 1: d ] = K [ u σ (1: d ) , f σ (1: d ) ] , where σ is the permutation operator. Hence, we augment the training dataset by permuting the function pairs in each task. Specifically, with 100 microstructures (tasks) and 100 function pairs per task, we randomly permut the function pairs and take 100 function pairs for 100 times per task. As a result, we can generate a total of 10000 samples (9000 for training and 1000 for testing) in the form of { u η 1:100 , f η 1:100 } 10000 η =1 .

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 is similar to the approach in Lejeune [2020]. The material property (stiffness) of the heterogeneous medium is constructed assuming a linear scaling between 1 and 100, according to the gray-scale bitmap of the MNIST images, which results in a set of 2D square domains with properties that vary according to MNIST digit patterns. The problem setting for this data set is the equilibrium equation: -∇ · P ( x ) = f ( x ) subjected to Dirichlet boundary condition of zero displacement and nonzero variable external forces: f ( x ) . P ( x ) = ˆ P ( I + ∇ u ( x )) is the stress tensor and is 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

[2020]. From MNIST, we took 50 samples of each digit resulting in a total of 500 microstructures. We consider 200 different external forces f ( x ) 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 × 140 uniform mesh. We then downsample from the finite element nodes to get values for the solution u ( x ) and external force f ( x ) on a the coarser 29 × 29 equi-spaced grid. The resulting dataset is of the form: { u i ( x j ) , f i ( x j ) } N i =1 , and j = 1 , 2 , ..., 841 where x j 's are the equi-spaced sampled spatial points.

Similar to Example 2, we perform the permutation trick to augment the training data. In 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,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 × d × 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 × h 1 + h 1 × h 2 ) . Thus, the total number of trainable parameters for a continuous-kernel NAO is O ( L × d × d k + d × h 1 + h 1 × h 2 ) .

The computational complexity of NAO is quadratic in the length of the input and linear in the data size, with O ([ d 2 (3 N + d k ) + 6 N 2 d ] 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 2 d k +2 Nd 2 +4 N 2 d ) flops, and the kernel map takes O ( d 2 d k + Nd 2 +2 N 2 d ) flops; thus, the total is O ( d 2 (3 N + d k ) + 6 N 2 d ) flops. In inverse PDE problems, we generally have d ≪ N , and hence the complexity of NAO is O ( N 2 d ) per layer per sample.

Compared 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 [Vaswani et al., 2017], convolution neural network (CNN), and graph neural network (GNN). As discussed in Vaswani et al. [2017], these models have a similar computational complexity, if not any higher. In particular, the complexity of the original attention model is O ( N 2 d ) , and the complexity of CNN is O ( kNd 2 ) with k being the kernel size, and a full GNN is of complexity O ( N 2 d 2 ) .

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 our claims in the abstract and introduction, including 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 framework, and offers the capability in solving ill-posed inverse PDE problems.

2. Limitations

Question: Does the paper discuss the limitations of 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 set of assumptions and a 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 conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We have elaborated the detailed architectures of the proposed model in the main body of the paper as well as data generation details in Appendix C. The code and accompanying data have been released at https://github.com/fishmoon1234/NAO .

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code and accompanying data have been released at https://github. com/fishmoon1234/NAO .

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 Experiments section, with the full details provided in Appendix C.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Error bars are not reported due to the limited computational resource. Nevertheless, we have tuned the hyperparameters of all the considered models and reported extensive results in a large variety of settings, and they all support our claim.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The computing resources are reported in the first paragraph of the Experiments section. In particular, all the experiments are conducted on a single NVIDIA GeForce RTX 3090 GPU with 24 GB memory.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conform with the NeurIPS Code of Ethics.

10. 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 on the societal impacts can be found in the Conclusion section.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

12. Licenses for existing assets

Question: Are the creators or 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?

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 documentation provided alongside the assets?

Answer: [Yes]

Justification: The code and accompanying data have been released at https://github. com/fishmoon1234/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 paper include the full text of instructions given to participants and screenshots, if applicable, as 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 or 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?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe idea is interesting and clever. The nature of sharing knowledge among multiple systems is somewhat related to in-context learning in LLMs.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nFor 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?\n\nLIMITATIONS:\nAgree with the limitation section.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

comments: 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.

STRENGTHS:
- mathematical proof is comprehensive
- ablation study has compared performance among several variants
- 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

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-08-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-08-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nCan 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-08-05"}