Nonlocal Attention Operator: Materializing Hidden Knowledge Towards Interpretable Physics Discovery

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 $\theta$ estimated from the training dataset (2). As such, it maps from the token $(\mathbf{u}_{1:d},\mathbf{f}_{1:d})$ of the dataset $\{(u_{i},f_{i})\}_{i=1}^{d_{u}}$ to a kernel estimator, acting as an inverse PDE solver.

A major 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})\}_{i=1}^{d_{u}}$ to tokens $(\mathbf{u}_{1:d},\mathbf{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}\}_{k=1}^{N}$.

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

Here, $\theta_{l}=(\bm{W}_{l}^{Q}\in^{d\times d_{k}},\bm{W}_{l}^{K}\in^{d\times d_{k}})$ and the attention function is

The kernel map is defined as

where $\theta:=\{W^{P,u}\in^{N\times N},W^{P,f}\in^{N\times N},\{\bm{W}_{l}^{Q},\bm{W}_{l}^{K}\}_{l=1}^{L}\}$ are learnable parameters. Here, we note that $d_{k}$ only controls the rank bound when lifting each point-wise feature via $\bm{W}_{l}^{Q}(\bm{W}_{l}^{K})^{\top}$. In the following, we denote $\bm{W}^{QK}_{l}:=\frac{1}{\sqrt{d_{k}}}\bm{W}_{l}^{Q}(\bm{W}_{l}^{K})^{\top}$, which characterizes the (trainable) interaction of $d$ input function pair instances.

As suggested by Cao (2021), we take the activation function $\sigma$ 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 $\theta$ by minimizing the following mean squared error loss function:

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

Here, $L^{2}(\rho)$ is the empirical measure in the kernel space as defined in Lu et al. (2023). Note that $\theta$ 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.

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}\}_{j=1}^{n_{x}}$. We define the tokens by

where $d=n_{x}$ and $\{r_{k}\}_{k=1}^{N}$ is the spatial mesh for $K$’s independent variable $r\in[0,\delta]$.

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.

The above space is data-adaptive since the integral kernel $\bar{G}$ depends on data. It characterizes the information in the training data for estimating the nonlocal kernel $K(s)=K[\mathbf{u}^{\eta}_{1:d},\mathbf{f}^{\eta}_{1:d};\theta](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.

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[\mathbf{u}_{1:d},\mathbf{f}_{1:d};\theta](s)$, with the parameter $\theta$ estimated from the training datasets, acts on the token $(\mathbf{u}_{1:d},\mathbf{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.

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 $\gamma_{\eta}$. 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).

In this context, our goal is to learn the operator $\mathcal{L}$ as well as to discover the hidden mechanism, namely the kernel $K[\bm{u}_{1:d},\bm{f}_{1:d};\theta](x,y)=\gamma_{\eta}({\left|y-x\right|})$. In the form of the operator in (19), we have $K(r)=\gamma_{\eta}(r)$ and $g[u](r,x)=u(x+r)+u(x-r)-2u(x)$for $r\in[0,\delta]$.

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

Here, $\eta$ denotes task index. We generate $4530$ data pairs $(g^{\eta}[u],f^{\eta})$ with a fixed resolution $\Delta x=0.0125$ for each task, where the loading function $\mathcal{L}_{\gamma_{\eta}}[u^{\eta}]=f^{\eta}$ 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 $\frac{4530}{302}=15$ samples. We consider two test kernels: one following the same rule of (24) with $\eta=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 (17) and the kernel error (18) 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 $\Delta 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 $\mathcal{L}_{\gamma_{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 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 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 OOD tasks. 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 (23), and hence the kernel value at this point does not change the operator 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 $\Delta x=0.0125$ on a dataset corresponding to $\Delta 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\%$).

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: $\gamma^{\sin}_{\eta}(r)=\exp(-\eta r)\sin(6r)\mathbf{1}_{[0,11]}(r)$, $\eta=1,2,3,4,6,7,8$.

• •

  cosine-type kernels: $\gamma^{\cos}_{\eta}(r)=\frac{10-r}{20}\cos(\eta r)(10-r)\mathbf{1}_{[0,10]}(r)$, $\eta=0,1,2,3,4,5,6$.

• •

  polynomial-type kernels: $\gamma^{\text{p}oly}_{\eta}(r)=\exp(-0.1r)p_{\eta}\left(\frac{r-10}{10}\right)\mathbf{1}_{[0,10]}(r)$, $\eta=1,2,3,4,5,6,7$, where $p_{\eta}$ is the degree-$\eta$ 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 2 and Figure 3. In addition to the original setting corresponding to all “sine” kernels, in part II of Table 2 (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, $\gamma(r)=\exp(-r)\sin(6r)\mathbf{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 $\gamma_{ood2}(r)=\frac{\exp(-0.5r^{2})}{\sqrt{2\pi}}$, which is substantially different from all training tasks.

As shown in Table 2, 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.

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 ${\Omega}=[0,1]^{2}$, $b(\bm{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(\bm{x})$. We consider two study scenarios. 1) $g\rightarrow p$: each task has a fixed microstructure $b(\bm{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 (26), and can be approximated by the inverse of the stiffness matrix. 2) $b\rightarrow p$: each task has a fixed loading field $g(\bm{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 3, 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 4. 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 5 we show the ground-truth microstructure $b(\bm{x})$, a test loading field instance $g(\bm{x})$, and the corresponding solution $p(\bm{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 5, 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 the applied Dirichlet-type boundary condition ($p(x)=0$ on $\partial\Omega$) in all samples, which leads to the measurement pairs $(p(x),g(x))$ containing no information near the domain boundary $\partial\Omega$ and makes it impossible to identify the kernel on boundaries.

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 4, where NAO again outperforms its discrete counterpart, especially in the small-data regime and OOD setting.