Paper: Improved Operator Learning by Orthogonal Attention (ICLR-2024)

Reviewer comment (Reproducibility):
I would like to see more details and clarifications regarding the choices made for hyperparameters in the experiments, such as the size of $k$ and $d'$.

Author response:
Q1 : Details and clarifications about the choices of hyperparameters. 
  
The size of $d'$ is always fixed to be equal to $d$ in $h_{i}^{(l)}$ to reduce the choice of hyperparamters in our experiments. $d'$ is typically 8 (for Elasticity, Plasticity, Pipe and Airfoil) or 16 (for Darcy and NS2d), which is much less than the width of the neural operator (128).

Reviewer comment (Experiments):
A per-epoch runtime comparison with other methods would also help clarify the model complexity.

Author response:
Q2: Runtime comparison. 
  
In the revision, we provide the runtime comparison on the Darcy benchmark in the **appendix B**. The results are shown as below:
  
| Model | FNO | Galerkin | Fourier | LSM |ONO (LT)| ONO (NT) |
|:---------|:---------:|:---------:|:---------:|:---------:|:---------:|:---------:|
| Run Time (s) | 3.81 |  9.88 | 58.16 | 9.05  | 7.87 | 34.46 |
  
Our implementation includes two types of NN blocks: Linear Transformer ("LT") and Nystrom Transformer ("NT"). The results demonstrate that the time cost of our neural operator is acceptable compared to the baselines, and  the complexity primarily originates from the NN block. We believe that these findings provide valuable insights into the model complexity. Thanks to your suggestion.

Reviewer comment (Experiments):
I am wondering whether multi-head attention can be naturally applied in the Orthogonal Attention blocks and if the authors have considered or tried that.

Author response:
Q3: Applying multi-head attention. 
  
We think it's an insightful suggestion to improve our proposed attention. We assume that the orthogonalization process is applied after the input is transformed into multiple attention heads. Theoretically, it will become kernel integral operator with multiple kernels, which may enrich the ability. We leave this as an interesting direction for future work.

Reviewer comment (Presentation):
Figure 2 is a bit confusing and misleading when compared with Equation 5. A more detailed caption might help clarify the process.

Author response:
Q5: The caption of Figure 2. 
  
We notice that Equation 5 can be misleading in the computation order. In practice, we will compute as $\tilde{h}^{(l+1)}_{i}= \hat{\psi}^{(l)}_i\mathrm{diag}(\hat{\boldsymbol{\mu}}^{(l)})[{\hat{\psi}^{(l)}_i}{}^{\top}( h^{(l)}_i w^{(l)}_V)]$ to maintain a linear complexity in our attention module. We have revised both the caption and the equation in the revision. Thanks to your comment!

Reviewer comment (Evaluation):
It is unclear how to justify that the method mitigates the overfitting problem, which is one of the central claims of the paper.

Author response:
Q2: The measurement of the overfitting.  
  
We mainly demonstrate the issue of overfitting by comparing the performance of our proposed neural operator with varying amounts of training data against the baseline Geo-FNO. In the latest revision of our paper, we have incorporated another well-acknowledged neural operator, DeepONet, as an additional baseline for comparison. The results highlight that our model achieves comparable performance while requiring significantly less training data. Additionally, our model exhibits lower performance degradation compared to the baselines. We also experimented with increasing the number of layers to 30 on the Elasticity and Plasticity benchmark. Our neural operator demonstrates significant improvement without exhibiting overfitting. This is in contrast to the performance of Geo-FNO, which shows a decline in performance when the number of layers exceeded 4 in [1]. This observation suggests that our proposed neural operator offers an effective regularization mechanism that contributes to its enhanced performance compared to the baseline.  
We believe that these results effectively showcase the problem of overfitting and highlight the regularization benefits of our proposed orthogonal attention.

Reviewer comment (Experiments):
The impact of the method versus the models' size should be studied.

Author response:
Q1: Performance versus model sizes. 
  
We have conducted experiments about the width and the number of layers of our model in our paper. Please see **section 4.4** for more details. The results show that our neural operator improves as we increase the width or layers. We further increased the layers to 30 and parameters to 10M, leading to notable enhancements on both the Elasticity and Plasticity benchmarks.

Reviewer comment (Experiments):
what are the performance versus model sizes?

Author response:
Q1: Performance versus model sizes. 
  
We have conducted experiments about the width and the number of layers of our model in our paper. Please see **section 4.4** for more details. The results show that our neural operator improves as we increase the width or layers. We further increased the layers to 30 and parameters to 10M, leading to notable enhancements on both the Elasticity and Plasticity benchmarks.

Reviewer comment (Evaluation):
how do you measure the overfitting?

Author response:
Q2: The measurement of the overfitting.  
  
We mainly demonstrate the issue of overfitting by comparing the performance of our proposed neural operator with varying amounts of training data against the baseline Geo-FNO. In the latest revision of our paper, we have incorporated another well-acknowledged neural operator, DeepONet, as an additional baseline for comparison. The results highlight that our model achieves comparable performance while requiring significantly less training data. Additionally, our model exhibits lower performance degradation compared to the baselines. We also experimented with increasing the number of layers to 30 on the Elasticity and Plasticity benchmark. Our neural operator demonstrates significant improvement without exhibiting overfitting. This is in contrast to the performance of Geo-FNO, which shows a decline in performance when the number of layers exceeded 4 in [1]. This observation suggests that our proposed neural operator offers an effective regularization mechanism that contributes to its enhanced performance compared to the baseline.  
We believe that these results effectively showcase the problem of overfitting and highlight the regularization benefits of our proposed orthogonal attention.

Reviewer comment (Experiments):
The motivation of avoiding overfitting with regularization is reasonable, but the paper lacks experiments that directly demonstrate overfitting issues in baseline models to substantiate the need for orthogonal regularization. Adding such empirical analysis could strengthen the motivation.

Author response:
Thanks to your comments about this. We have performed an experiment on the Elasticity benchmark, varying the number of training data, to illustrate the issue of overfitting when data is limited. Our results indicate that our neural operator exhibits a smaller decrease in performance with fewer training data compared to the baseline model Geo-FNO. We have also included another well-acknowledged neural operator, DeepONet, as an additional baseline in our revised version. This addition serves to further highlight the issue and strengthen the motivation behind our work. What's more, we have also experimented with increasing the number of layers to 30 on the Elasticity benchmark. In contrast, Geo-FNO showed a decline in performance when exceeding 4 layers in [1]. This suggests that our proposed neural operator incorporates effective regularization mechanisms, contributing to its superior performance compared to the baseline model.

Reviewer comment (Experiments):
The ablation study verifies the usefulness of the orthogonalization component. However, it does not isolate the impact of the two-pathway architecture itself. Additional experiments are needed to demonstrate the benefits of the disentangled design.

Author response:
Thanks for your interest in the disentangled design. We think that the orthogonal attention naturally requires two pathways and their inputs: the eigenfunctions and the input function, making it challenging to replace the disentangled design.

Reviewer comment (Experiments):
The long-term impact could be boosted by testing on real-world datasets and problems beyond standard benchmarks to showcase effectiveness in complex practical settings. The paper could also provide a more comprehensive evaluation of the model's performance across a broader spectrum of PDEs, including those with more complex boundary conditions and non-linearities. This would not only demonstrate the robustness of the model but also identify potential limitations in its current form.

Author response:
We agree that more complex practical settings can evaluate the neural operator more comprehensively. We are also interested in the performance in practical scenarios. However, we emphasize that our current empirical studies align with the related works in this field and can prove the effectiveness of our method. We will explore more complex tasks in the next version.

Reviewer comment (Experiments):
I understand the motivation that the top line in figure 1 updates the PDE solution so that strong regularization is needed, and that is why the proposed orthogonal attention is incorporated, along with linear attention. However, technically, both the top line and the bottom line are simply nonlinear functions, so have the authors tried to incorporate the proposed attention into both lines to see if it further improves the generalization?

Author response:
Thanks to this comment. We think that the bottom line requires expressive structures to extract features from the input data and approximate the eigenfunctions rather than structures with regularization that may limit the overall expressiveness of the neural operator. In fact, we conducted such an experiment on the Darcy benchmark and observed that the neural operator exhibits unsatisfactory performance (0.036 versus 0.0072 on Darcy). Hence, we recommend employing a well-studied linear attention instead of orthogonal attention in the bottom line.

Reviewer comment (Reproducibility):
The current parametrization requires solving a Cholesky decomposition for each batch of the data during training, and I wonder whether there could be ways to simplify the process. For example, could we fix the orthonormal matrix before training, and learn mu as a function of the input data? In that way, we could avoid taking the inverse of a whole matrix each iteration but rather taking the inverse of individual values in mu.

Author response:
We appreciate your insightful advice regarding simplifying the orthogonalization process. We understand your concern about the complexity involved in the process. The complexity of our proposed process mainly relies on the Cholesky decomposition of the covariance matrix and matrix inversion. We think that both processes maintain an acceptable complexity of O($8^3$) (for Elasticity, Plasticity, Pipe, and Airfoil) and O($16^3$) (for Darcy and NS2d). In our latest revision, we have provided the time cost compared with baselines in **appendix B**, which demonstrates that the complexity is reasonable. We acknowledge that the process can be further simplified and we are grateful for your suggestion. We will explore this in our future work.