Nonlocal Attention Operator: Materializing Hidden Knowledge Towards Interpretable Physics Discovery

We thank the reviewer for their valuable time and constructive suggestions. We have generated more families of kernels, and performed additional tests on three different training settings (diverse tasks, single task with diverse function pairs, and diverse tasks with a reduced number of samples). Additionally, we have created an OOD test task with a substantially different kernel from the trained ones (i.e., Gaussian-type kernel versus sine, cosine and polynomial-type kernels). All additional test results can be found in the table of the one-page PDF file. We will also modify the manuscript to make the narrative cleaner and more compact. Our other responses:

Broader range of tasks for different types of kernels: Per the reviewer's suggestion, we have added another two types of kernels into the training dataset. The training dataset is now constructed based on 21 kernels of three groups, with $15$ samples on each kernel:

sine-type kernels: $\gamma_\eta(r)=\exp(-\eta r)\sin(6r)\mathbf{1}_{[0,11]}(r)$, $\eta=1,2,3,4,6,7,8$.

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

polynomial-type kernels: $\gamma_\eta(r)=\exp(-0.1 r)p_\eta$($\frac{r-10}{10}$), $\eta=1,2,3,4,5,6,7$, where $p_\eta$ is the degree-$\eta$ Legendre polynomial.

Based on this enriched dataset, beyond the original ''sin only'' setting, three additional settings are considered in Table 1 of the attached PDF. In part II of Table 1 (denoted as ''sin+cos+poly''), we consider a ''diverse task'' setting, where all $315$ samples are employed in training. Then, in part III we consider a ''single task'' setting, where only the first sine-type kernel

$\gamma(r)=\exp(-r)\sin(6r)\mathbf{1}_{[0,11]}(r)$

is considered as the training task, with all $105$ samples on this task. Lastly, in part IV we demonstrate a ''fewer samples'' setting, where the training dataset still consists of all $21$ tasks but with only $5$ samples on each task. For the purpose of testing, besides the in-distribution (ID) and out-of-distribution (OOD) tasks in the original manuscript, we have added an additional (OOD) task with Gaussian-type kernel $\gamma_{ood2}(r)=\frac{\exp(-0.5r^2)}{\sqrt{2\pi}}$, which is substantially different from all training tasks.

As shown in Table 1, considering diverse datasets helps in both OOD tests (kernel error reduced from 9.14% to 6.92% in OOD test 1, and from 329.66% to 10.48% in OOD test 2), but not for the ID test (kernel error slightly increased from 4.03% to 5.04%). We also note that since the ''sin only'' setting only sees systems with the same kernel frequency in training, it does not generalize to the OOD test 2. In this test, including more diverse training tasks becomes necessary. Then, as the training tasks become more diverse, the intrinsic dimension of kernel space increases, requiring a larger weight matrices rank size $d_k$. When comparing the ''fewer samples'' setting with the ''sin only'' setting, the former has better task diversity but fewer samples per task. One can see that the performance on ID test deteriorates but the performance on OOD test2 improves. This observation highlights the balance between the diversity of tasks and the number of training samples per task.

Difference from learning more data pairs from a single system: Although all 7 sine-type kernels are corresponding to the same frequency, they are very different from each other. If we train the model from a single kernel, the model does not generalize at all. Then, the test error on validation task does not improve during training, and the model always predicts an all-zero kernel. As shown in the ''Single task'' setting of attached Table 1, the kernel and operator error are always around 100% for all test tasks.

Reported number of data pairs: When we talk about the data pairs for each task, as in Example 1, the reported number of data pairs refers to the total number of data pairs of each task. Then, we form a training sample of each task by taking $d$ pairs from this task. When taking the token size $d=302$, each task contains $\lfloor\frac{4530}{302}\rfloor=15$ samples. When having 7 training tasks corresponding to 7 different sine-type kernels, there are $105$ samples in total. Then, when discussing about the total number of data pairs, as in Examples 2-3, we are reporting it across all tasks. In operator learning settings, we note that the permutation of function pairs in each sample should not change the learned kernel, i.e., one should have K[$u_{1:d}$,$f_{1:d}$]=K[$u_{\sigma(1:d)}$,$f_{\sigma(1:d)}$], where $\sigma$ is the permutation operator. Hence, we have augmented the training dataset by permuting the function pairs in each task. Taking example 3 for instance, with $500$ microstructures (tasks) and $200$ function pairs per task, we have randomly permuted the function pairs and taken $100$ function pairs for $100$ times per task. As such, we have created the total 50000 samples (45000 for training and 5000 for test) in Example 3. We will clarify this in the revised manuscript.

We thank the reviewer for their valuable time and constructive suggestions. Our source code, trained models, and datasets will be made publicly available on Github upon acceptance of the paper, so as to guarantee reproducibility of all results. Our other responses:

Limited baselines: We did not compare much with other work because very few existing models are capable of discovering hidden physics or the governing physical parameters directly from data. Beyond the direct application of the attention model as in the discrete NAO setting, the most relevant work to our best knowledge is the adaptive Fourier neural operator that adopted a convolution-based attention mechanism, which we compared to in the ablation study. However, it is unable to learn a good solution across tasks (as indicated by the large errors in the in-distribution and out-of-distribution tests), nor can it discover physics interpretation for the inverse problem.

To further address the reviewer's concern, we have added an autoencoder model as an additional baseline in Example 1, see Table 1 of the one-page PDF. In this setting, an encoder is constructed to map the concatenated data pairs [$u_{1:d}$,$f_{1:d}$] to the hidden kernel K[$u_{1:d}$,$f_{1:d}$] as the latent representation, and then in the decoder the kernel is multiplied with $u$ following the last step of Eq.(6): K[$u_{1:d}$,$f_{1:d}$]$u_{1:d}$-> $f_{1:d}$. We note that the concatenated data pairs [$u_{1:d}$,$f_{1:d}$] is of size $O(Nd)$ as the input of the encoder, and the kernel $K$ is of size $N^2$ as the output. Hence, even a linear encoder contains around $N^3d$ trainable parameters. When the autoencoder model is of a similar size to NAO, its expressivity is very limited.

We thank the reviewer for the valuable comments and suggestions. Our response:

Number of trainable parameters and computational complexity: Denoting $(N,d,d_k,L,n_{train})$ as the size of the spatial mesh, the number of data pairs in each training model, the column width of the query/key weight matrices $W^Q$ and $W^K$, the number of layers, and the number of training models, the number of trainable parameters in a discrete NAO is of size $O(L\times d\times d_k+N^2)$. For continuous-kernel NAO, taking a three-layer MLP as a dense net of size $(d, h_1, h_2, 1)$ for the trainable kernels $W^{P,u}$ and $W^{P,f}$ for example, its the corresponding number of trainable parameters is $O(d\times h_1+h_1\times h_2)$. Thus, the total number of trainable parameters for a continuous-kernel NAO is $O(L\times d\times d_k+d\times h_1+h_1\times h_2)$.

The computational complexity of NAO is quadratic in the length of the input and linear in the data size, with $O([d^2(3N+d_k) + 6N^2d] Ln_{train})$ flops in each epoch in the optimization. It is computed as follows for each layer and the data of each training model: the attention function takes $O(d^2d_k+ 2Nd^2+4N^2d )$ flops, and the kernel map takes $O(d^2d_k+ Nd^2+2N^2d)$ flops; thus, the total is $O(d^2(3N+d_k) + 6N^2d)$ flops. In inverse PDE problems, we generally have $d\ll N$, and hence the complexity of NAO is $O(N^2d)$ per layer per sample.

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

[1] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017).

Data efficiency: The minimal data for NAO to perform well requires a balance between multiple factors, including the intrinsic dimension (smoothness) of the kernel, the token size $d$ for each sample, the rank $d_k$ of the weight matrices, the number of samples in each task, and the diversity of training tasks. Thus, it is difficult to have a theoretical guarantee. In the original manuscript, with all other factors fixed, we have tested the dependence on the token dimension and the rank of the weight matrices in Table 1. To further elaborate the importance of sample number and task diversity, in Table 1 of the one-page PDF we have further perform tests on three settings: 1) diverse tasks; 2) single task with diverse function pairs; and 3) diverse tasks with a reduced number of samples. In particular, in the diverse task setting, the training dataset is constructed based on 21 kernels of three groups, with $15$ samples on each kernel:

sine-type kernels: $\gamma_\eta(r)=\exp(-\eta r)\sin(6r)\mathbf{1}_{[0,11]}(r)$, $\eta=1,2,3,4,6,7,8$.

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

polynomial-type kernels: $\gamma_\eta(r)=\exp(-0.1 r)p_\eta$( $\frac{r-10}{10}$ ), $\eta=1,2,3,4,5,6,7$, where $p_\eta$ is the degree-$\eta$ Legendre polynomial.

In the ''diverse task'' setting, all $315$ samples are employed in training. In the ''single task'' setting, only the first sine-type kernel:

$\gamma_1(r)=\exp(-r)\sin(6r)\mathbf{1}_{[0,11]}(r)$

is considered as the training task, with all $105$ samples on this task. Then, in the ''fewer samples'' setting, the training dataset still consists of all $21$ tasks but with only $5$ samples on each task. For the purpose of testing, besides the in-distribution (ID) and out-of-distribution (OOD) tasks in the original manuscript, we have added an additional (OOD) task with Gaussian-type kernel $\gamma_{ood2}(r)=\frac{\exp(-0.5r^2)}{\sqrt{2\pi}}$, which is substantially different from all training tasks.

As shown in the results in Table 1 of the one-page PDF, for the ID test task, a small number of training tasks are sufficient, and the model performs well in all except the ''single task'' setting. When the test task has a larger discrepancy from the training tasks, as in the OOD test2 setting, it becomes necessary to increase the training task diversity, and the required rank size $d_k$ should also increase to obtain better model expressivity. When comparing the ''fewer samples'' setting with the ''sin only'' setting, the former has better task diversity but fewer samples per task. One can see that the performance on ID test deteriorates but the performance on OOD test2 improves. This observation highlights the balance between the diversity of tasks and the number of training samples.

Combination with physics-based modeling approaches: When the physics structure is known, NAO can directly encodes such information in its architecture, such as the kernel structure in the nonlocal diffusion model in Example 1. By doing so, we anticipate to explore the kernel in a lower-dimensional manifold, which reduces the requirement on the size of training data.

On the other hand, NAO could also help physics-based modeling approaches by providing a surrogate foundation model when there is a large class of physical models to be estimated and simulated. In this setting, NAO can be first trained using data generated by these models, and then used as an efficient surrogate model. It will be useful for sampling and uncertainty quantification tasks.

We thank the reviewer again for reading our rebuttal and for the valuable further suggestions. We sincerely appreciate the reviewer's helpful discussions, which helped us to improve the quality of this work.

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

Interpretability for physical understanding: The physical meaning of the learned kernel is application-dependent: in the context of material modeling, the learned kernel characterizes the interaction between material points; in turbulence modeling problems, the interaction range of the kernel suggests the characteristic length of the eddies; in fracture mechanics, the vanishing of kernel values between two neighboring regions may indicate a crack between these two regions; just to name a few. Generally speaking, identifying the kernel is instrumental in identifying influential factors, characterizing spatial dependence, and analyzing the forward operator.

(i) The learned kernels themselves are fundamental for understanding the physical model. They indicate the range and strength of interactions between components. For example, a learned interaction kernel can reveal interactions between particles or agents. In the Darcy's problem $g\rightarrow p$ setting, when taking a summation of the kernel strength on each row, one can discover the interaction strength of each material point $x$ with its neighbors. Then, since such a strength is related to the permeability field $b(x)$, the underlying microstructure can be recovered. In Figure 2 of the one-page PDF, we demonstrate the discovered microstructure on a test specimen. We note that the discovered microstructure was smoothed-out due to the continuous setting of our learned kernel (as shown in the middle plot), a thresholding step was performed to discover the two-phase microstructure. The discovered microstructure (right plot) matches well with the hidden ground-truth microstructure (left plot), except on the regions near the domain boundary. This is because of the Dirichlet-type boundary condition setting ($p(x)=0$ on $\partial\Omega$) in all samples. As a result, the measurement pairs $(p(x),g(x))$ contain no information near the domain boundary $\partial\Omega$, making it impossible to identify the kernel on boundaries.

(ii) The learned kernel characterizes the output's spatial dependence on the input. The learned kernel in the nonlocal diffusion model of our Example 1 demonstrates this spatial dependence. Such a range of spatial dependence characterizes the characteristic lengths, which is especially of interests in multiscale problems (see, e.g., [1]).

(iii) The learned kernels facilitate the analysis of the forward operator, including its spectral properties, stiffness for computation, and error bounds for output prediction under perturbations in the input.

We will add these discussions and the additional figure to the revised manuscript.

[1] You, Huaiqian, et al. "A data-driven peridynamic continuum model for upscaling molecular dynamics." Computer Methods in Applied Mechanics and Engineering 389 (2022): 114400.

We thank the reviewer again for reading our rebuttal and for the kind response. We sincerely appreciate the reviewer's valuable time and suggestions, which helped us to improve the quality of this work.