Paper: Energy-conserving equivariant GNN for elasticity of lattice architected metamaterials (ICLR-2024)

Reviewer comment (Novelty):
The novelty is somewhat limited in the sense that the proposed model is a direct utilization of MACE, equipped with the proposed layer that satisfies PSD.

Author response:
We recognize that the reviewer is worried about the lack of ML-specific novelty of this work.
We would like to address these concerns on two levels:
(i) it is not trivial to enforce positive semi-definiteness of a 4-th order tensor while maintaining equivariance (Reviewer azcZ pointed out a need for formal mathematical proofs, which we now include in the manuscript.) Therefore, we believe that this constitutes significant contribution to the community.
(ii) The paper is submitted in the area of *applications to physical sciences*, and another key contribution is that we are making the dataset publicly available. There has been a lack of datasets and corresponding method development for higher-order tensors, which we are hoping to address by this work.

Reviewer comment (Experiments):
The dataset is a bit limited and the experiments lack of comparison with other methods. If there are no significant deep learning based (or even GNN-based) methods to compare, it would also help a lot if any FE-based method could be involved. The readers may be curious about how the proposed model can achieve better accuracy/efficiency tradeoff compared with traditional solvers.

Author response:
It is true that the dataset lacks comparisons with other methods -- that is because there is presently a lack of methods which focus on lattices and 4-th order tensors.
The two previously presented GNN methods for lattices are based on CGC and NNConv, which are the models against which we benchmark our MACE-based model.
Moreover, we have included runtime comparisons for the three model families, and for FE baseline.
We show that the ML models achieve a speedup of 3 orders of magnitude.

Reviewer comment (Experiments):
The implication of the practical usage of the method is limited. Since this paper is developed purely based on practical considerations, it would be better if the paper could show how this model can help in other related/similar tasks or downstream tasks that could potentially benefit from this method, other than just predict the stiffness tensor.

Author response:
We thank the reviewer for pointing out other tasks that could benefit from this methods other than just the prediction of the stiffness tensor.
The presented methods is applicable to any 4-th order tensor. Examples include diffusion tensor in MRI,  piezo-optical tensor and the elasto-optical tensor.
We have incorporated this statement in the Abstract and Conclusion.

Reviewer comment (Experiments):
How does the method perform compared with other solvers, deep learning-based or even not?

Author response:
It is true that the dataset lacks comparisons with other methods -- that is because there is presently a lack of methods which focus on lattices and 4-th order tensors.
The two previously presented GNN methods for lattices are based on CGC and NNConv, which are the models against which we benchmark our MACE-based model.
Moreover, we have included runtime comparisons for the three model families, and for FE baseline.
We show that the ML models achieve a speedup of 3 orders of magnitude.

Reviewer comment (Experiments):
Are there other tasks/datasets that can benefit from the development of such GNN-based method that is specifically designed to regress on the stiffness?

Author response:
We thank the reviewer for pointing out other tasks that could benefit from this methods other than just the prediction of the stiffness tensor.
The presented methods is applicable to any 4-th order tensor. Examples include diffusion tensor in MRI,  piezo-optical tensor and the elasto-optical tensor.
We have incorporated this statement in the Abstract and Conclusion.

Reviewer comment (Experiments):
Even if the paper discusses the limitations of previous Cholesky/eigenvalue decomposition-based methods that permit PSD constraint, it would help a lot if this point is also verified by experimental results/ablation studies.

Author response:
We do not provide benchmarks with Cholesky/eigenvalue based methods because the Choleksy method fails on the basic requirement of equivariance, and the eigenvalue-based method is unstable during training due to the underlying instability of gradients of eigenvector computation.

Reviewer comment (Presentation):
It seems that most of the important information which is novel to the deep learning community is put into the appendix. For example, definitions of stress and strain tensors, or their relation to the stiffness tensor. Even more importantly, only in the appendix one can read why the stiffness tensor can be represented as a matrix, and how the positive definiteness of a matrix can be ensured.

Author response:
We expand on the solid mechanics section in the background, noting that for energy conservarion, the stiffness tensor must be positive semi-definite. We moved the details of MACE model which might feel repetitive into the appendix.

Reviewer comment (Presentation):
On the other hand, background and related work are a bit repetitive.

Author response:
We expand on the solid mechanics section in the background, noting that for energy conservarion, the stiffness tensor must be positive semi-definite. We moved the details of MACE model which might feel repetitive into the appendix.

Reviewer comment (Reproducibility):
Code / or pseudocode would be pretty helpful to understand output layers and how the stiffness tensor is ensembled.

Author response:
In the Methods section, we included a schematic of the model, so it is clear where the PSD layer is applied.  
We expanded the Methods section to include more details about the dataset and the positive semi-definite (PSD) layer.  
We explain how the fourth-order tensor is represented by $6\times6$ matrix and how it can be made positive definite.  
In addition, we have written a mathematical proof that our PSD layer in Mandel notation is equivariant. This is now in the Appendix.

Reviewer comment (Experiments):
There are no runtime comparisons in the paper, especially since e3nn based models are known to be slow this would be interesting to know.

Author response:
We provide inference time comparisons for the three families of ML models and FE baseline.  
While the e3nn-based model is slower than the ML counterparts, all ML models are approximately 3 orders of magnitude faster than FE.

Reviewer comment (Experiments):
Related to runtime comparisons, comparisons to FE methods are needed to get a better understanding of the presented performances. I understand that FE models are used as ground truth and that this might be tricky to obtain, but for example one could estimate runtime comparisons (especially since it is stated the FE methods are very slow) and report performance differences for nodal perturbations to get a perspective for the reported loss values? The latter should give a feeling to what nodal perturbations the presented losses are comparable.

Author response:
We provide inference time comparisons for the three families of ML models and FE baseline.  
While the e3nn-based model is slower than the ML counterparts, all ML models are approximately 3 orders of magnitude faster than FE.

Reviewer comment (Experiments):
In Figure 2 validation loss of CGC+tr seems to be still dropping sharply, is it possible that training was stopped too early?

Author response:
While it appears in Figure 2 (now Fig 3) that validation loss for CGC+tr continues to drop sharply, this is exaggerated by the log-scale nature of the x-axis.  
All models were trained to convergence.

Reviewer comment (Novelty):
This work chooses to predict a 4-th ordered tensor – “stiffness tensor”, whose symmetry has been maintained well with existing models like TFN, PAINN. This paper does not propose theoretical innovations in maintaining equivariance, so this methodology and the emphasis on “fourth-order” in abstract are not directly related (the Positive Semi-Definite Layer is based on physical meaning, regardless of whether it is a fourth-order tensor).

Author response:
The reviewer states that 

>    the existing models (like TFN and PAINN) have been able to fit a 4-th order tensor and that this paper does not propose theoretical innovations in maintaining equivariance.


While it is true that models like TFN and PAINN can fit 4-th-order tensors, there is a lack of such examples in the machine learning community. 
For instance, the PAINN paper shows an example of fitting rank-2 polarizability tensor.[1,p5] 
Very few, if any, examples of rank-4 tensors have been reported in the machine learning community.
We believe that it is one of the contributions of this paper (in the area of *applications to physical sciences*) to make available a real-world dataset of rank-4 tensors, as this could fuel further developments in the community. 

Secondly, Positive Semi-Definite (PSD) layer brings about new challenges related to equivariance.
Maintaining equivariance of the 4-th order tensor is more challenging than for the rank-2 tensor.
A new theoretical contribution in maintaining equivariance, consistent with positive-definiteness of a 4-th order tensor, is needed and is presented here.

Reviewer comment (Theory):
Since the whole backbone is built with existing MACE layer, the main novelty is to propose the “Positive Semi-Definite Layer” (let us call it “PSD-layer”). To maintain the equivariance of the total model, we must ensure each layer in such a model is equivariant. But for the PSD-layer, the PSD matrix \(A\) is based on \(M\), which is created by arranging \(n(n+1)/2\) entries of the output of an equivariant model. Is it an equivariant operator to arrange \(n(n+1)/2\) entries to get \(M\)? This requires a formal mathematical proof.

Author response:
The reviewer raises valid questions. Based on these, we have made the following revisions to the manuscript.


1.  We included in the Appendix section **Proof of equivariance of PSD layer in Mandel notation**
in which we mathematically prove that our PSD layer is equivariant.
1.  The reviewer voices concerns that function $\rho$ might lead to a reduced representational capacity of the model.
While we are unable to run further experiments at this short timescale, we will make effort to address reviewer's points in due course.
Empirically and from the point of view of using this model in practice, the performance hit due to the possible reduction in representational capacity is minor as seen in Table 1.
1.  We have included section **Speedup using machine learning methods** in the results section
where we compare inference speed of the three families of models and FE.
We show that the ML models bring a speedup of 3 orders of magnitude.
Since FE results are the ground truth data for ML training, it is not meaningful to compare the accuracy of FE with the accuracy of ML models.
1.  The reasoning behind our decision to discuss models based on CGC, NNConv, and MACE is as follows.
The only GNN models, to our best knowledge, that have been used for lattices are CGC and NNConv. [2,3]
Therefore we decided to benchmark our model against the previous baseline used in the literature.
The reason why we chose MACE as the basis for our model (as opposed to TFN, NequIP etc) is that
it has been reported by the computational chemistry community as the best equivariant model for molecules.
As a side note, in paragraph **Correlation order** on p.9 we note that MACE with $\nu=1$ is equivalent to TFN.

Reviewer comment (Theory):
Regarding the same PSD-layer formulation, is it an equivariant operator of the function \(\rho\) to turn \(M\) into \(A\)? In the end of Section A.2.2, “If the overall model had been equivariant with respect to vectors in \(U\), it will remain equivariant after eigenvalues are made positive” is not trivial, and it may also require a formal mathematical proof.

Author response:
The reviewer raises valid questions. Based on these, we have made the following revisions to the manuscript.


1.  We included in the Appendix section **Proof of equivariance of PSD layer in Mandel notation**
in which we mathematically prove that our PSD layer is equivariant.
1.  The reviewer voices concerns that function $\rho$ might lead to a reduced representational capacity of the model.
While we are unable to run further experiments at this short timescale, we will make effort to address reviewer's points in due course.
Empirically and from the point of view of using this model in practice, the performance hit due to the possible reduction in representational capacity is minor as seen in Table 1.
1.  We have included section **Speedup using machine learning methods** in the results section
where we compare inference speed of the three families of models and FE.
We show that the ML models bring a speedup of 3 orders of magnitude.
Since FE results are the ground truth data for ML training, it is not meaningful to compare the accuracy of FE with the accuracy of ML models.
1.  The reasoning behind our decision to discuss models based on CGC, NNConv, and MACE is as follows.
The only GNN models, to our best knowledge, that have been used for lattices are CGC and NNConv. [2,3]
Therefore we decided to benchmark our model against the previous baseline used in the literature.
The reason why we chose MACE as the basis for our model (as opposed to TFN, NequIP etc) is that
it has been reported by the computational chemistry community as the best equivariant model for molecules.
As a side note, in paragraph **Correlation order** on p.9 we note that MACE with $\nu=1$ is equivalent to TFN.

Reviewer comment (Theory):
Assuming that the previous question (function \(\rho\) is equivariant) has a good proof, it is necessary to discuss whether this design will reduce the representational capacity of the entire neural network. Assume that this neural network is a bijection (we consider an equivalence class divided based on different perspectives as an input). But the function \(\rho\) may not be a bijection, e.g., \((\pm \Lambda)^2\) will get the same result. The Wigner-D based network builds a faithful representation of the group, but the induction of function \(\rho\) will lead to an unfaithful representation, which requires more analysis. For example, an analysis for the increased stiffness-based errors of CGC+ve, NNConv+ve, and MACE+ve in Table 1 may be a manifestation of the decline in the ability of network representation.

Author response:
The reviewer raises valid questions. Based on these, we have made the following revisions to the manuscript.


1.  We included in the Appendix section **Proof of equivariance of PSD layer in Mandel notation**
in which we mathematically prove that our PSD layer is equivariant.
1.  The reviewer voices concerns that function $\rho$ might lead to a reduced representational capacity of the model.
While we are unable to run further experiments at this short timescale, we will make effort to address reviewer's points in due course.
Empirically and from the point of view of using this model in practice, the performance hit due to the possible reduction in representational capacity is minor as seen in Table 1.
1.  We have included section **Speedup using machine learning methods** in the results section
where we compare inference speed of the three families of models and FE.
We show that the ML models bring a speedup of 3 orders of magnitude.
Since FE results are the ground truth data for ML training, it is not meaningful to compare the accuracy of FE with the accuracy of ML models.
1.  The reasoning behind our decision to discuss models based on CGC, NNConv, and MACE is as follows.
The only GNN models, to our best knowledge, that have been used for lattices are CGC and NNConv. [2,3]
Therefore we decided to benchmark our model against the previous baseline used in the literature.
The reason why we chose MACE as the basis for our model (as opposed to TFN, NequIP etc) is that
it has been reported by the computational chemistry community as the best equivariant model for molecules.
As a side note, in paragraph **Correlation order** on p.9 we note that MACE with $\nu=1$ is equivalent to TFN.

Reviewer comment (Experiments):
Assuming the PSD-layer is equivariant, more powerful baselines should be discussed, especially the geometric GNNs like TFN [1], NequIP [2], SCN [3], GMN [4], SE(3)-Transformer [5].

Author response:
The reviewer raises valid questions. Based on these, we have made the following revisions to the manuscript.


1.  We included in the Appendix section **Proof of equivariance of PSD layer in Mandel notation**
in which we mathematically prove that our PSD layer is equivariant.
1.  The reviewer voices concerns that function $\rho$ might lead to a reduced representational capacity of the model.
While we are unable to run further experiments at this short timescale, we will make effort to address reviewer's points in due course.
Empirically and from the point of view of using this model in practice, the performance hit due to the possible reduction in representational capacity is minor as seen in Table 1.
1.  We have included section **Speedup using machine learning methods** in the results section
where we compare inference speed of the three families of models and FE.
We show that the ML models bring a speedup of 3 orders of magnitude.
Since FE results are the ground truth data for ML training, it is not meaningful to compare the accuracy of FE with the accuracy of ML models.
1.  The reasoning behind our decision to discuss models based on CGC, NNConv, and MACE is as follows.
The only GNN models, to our best knowledge, that have been used for lattices are CGC and NNConv. [2,3]
Therefore we decided to benchmark our model against the previous baseline used in the literature.
The reason why we chose MACE as the basis for our model (as opposed to TFN, NequIP etc) is that
it has been reported by the computational chemistry community as the best equivariant model for molecules.
As a side note, in paragraph **Correlation order** on p.9 we note that MACE with $\nu=1$ is equivalent to TFN.