Paper: Hybrid Directional Graph Neural Network for Molecules (ICLR-2024)

Reviewer comment (Theory):
Since you are claiming that your architecture: “can enhance the expressive power suffering from the limitations of finite low-degree group representation” should this be also shown in a proposition or theorem, to complement your ablation study?

Author response:
We discuss relevant theorem and experiments in Appendix A.3. Our discussion is based on two theoretical works [1] [2]. **The theorem 2 in [1] asserts that any equivariant function can be approximated by a CG tensor product with infinite degree, though practical constraints limit it to a finite low degree.** Therefore, some complex equivariant pattern can not be fitted by CG produt with low degree (please refer to Table 2 in [2]). Concurrently, **based on the approximation properties of MLP, we recognize MLP's capability to approximate any continuous function, which includes equivariant functions (please refer to figure 1 in [1]).** Our intention is to leverage MLP to overcome the limitations of CG tensor products with low degrees. **Table 6 further illustrates MLP's proficiency in learning higher-degree patterns when embedding is in a low-degree representation.** Due to space constraints and the non-original nature of the aforementioned theorem, we provide a more detailed discussion in the Appendix.

[1] On the Universality of Rotation Equivariant Point Cloud Networks. Nadav Dym, Haggai Maron. ICLR 2021

[2] On the Expressive Power of Geometric Graph Neural Networks. Chaitanya K. Joshi, Cristian Bodnar, Simon V. Mathis, Taco Cohen, Pietro Li.

Reviewer comment (Experiments):
Could you expand experiments on other recent GNN benchmarks, follow Equiformer or the other works that you are comparing to.

Author response:
Firstly, we did not assess the MD17 dataset. Given its small size, it poses challenges for HDGNN to learn the correct equivariance (other methods use inherent equivariance, expect for SCN). **Our work is designed to enhance expressive power by relaxing equivariance, enabling networks to comprehend intricate physical interactions between atoms. Consequently, HDGNN is expected to have good performance across a large swath of chemical space. OC20 stands out as representative due to its substantial size (O(100M)) and diversity (55 elements).** Additionally, Equiformer presented some results involving data augmentation on OC20 experiments, such as noisy nodes and IS2RS. This may be attributed to Equiformer's competitive nature (OC20 challenge), where the authors tend to achieve optimal MAE. In contrast, our study focuses on only comparing the performance of pure models. The basic training environment (without augmentation) can provide an accurate evaluation.

Since part of the equivariance of HDGNN comes from learning, **SO(3)-transformation augmentation of MD17 may allow HDGNN to have a good performance. But this will lead to unfair comparisons (the results of other methods use the original dataset).** In addition, we are sorry that we cannot complete the data augmentation experiments on OD20 during the rebuttal time due to limited resources (OC20 experiments take up a lot of resources). We will publish these results in the public community in the future. **Nevertheless, outcomes derived from data augmentation may not precisely represent the true performance of the model, as it remains unclear which specific component contributes to the observed improvement. We believe that the pure model comparisons in our Table 1 and Table 2 can provide a significant evaluation for HDGNN.**

Reviewer comment (Theory):
A lot of focus is on breaking of equivariance, or relaxing this constraint. There is one element that I find very important which is not discussed however: When converting the features to a global reference frame via the Wigner-D matrix (Equation 9), the rotation is computed from the direction: "$R_{ij}$ denotes the 3D rotation matrix that transforms  $\vec{r}_{ij}$ to $[0,0,1]$". However, this is ill-defined as there is one degree of freedom left in the rotation matrix as any $R_{ij} R_{\alpha}$ with $R_{\alpha}$ any rotation around the z-axis results in the same mapping from $\vec{r}_{ij}$ to $[0,0,1]$. Thus, the equality in equation 9 is not true, and this is an error in the paper. Is this a deliberate error, because you want to break equivariance? Either way, it should be fixed or commented on. It seems that this type of trick is adopted from [Zitnick et al. 2022], although I haven't checked if the same type of error is present in that paper. However, working with local reference frames was also presented in [Kofinas et al. 2021]. See section 3.2 for their discussion on the ill-defined rotation. The claim "... was first presented in Zitnick et al" might therefore have to be also updated/nuanced.

Author response:
While acknowledging the existence of multiple rotation transformations $\mathbf{R} _{ij}$ capable of rotating $\vec{\mathbf{r}} _{ij} $ to [0,0,1], **the equality in Equation 9 remains valid**. In detail, we use two rotation parameters, namely $\theta$ and $\varphi$ (changes in the polar and azimuth angles), to denote $\mathbf{R} _{ij}$. $\varphi$ is random when $\mathbf{R} _{ij}$ rotate $\vec{\mathbf{r}} _{ij}$ to [0,0,1]. We randomly select one parameter configuration and consistently apply it throughout the subsequent processes. With the choosen $\mathbf{R} _{ij}$, we get $\mathbf{D}^{-1}(\mathbf{R} _{ij})\mathbf{D}(\mathbf{R} _{ij})=\mathbf{I}$. Consequently, we initially convert the left side of Equation 9 to $(\mathbf{D}^{-1}(\mathbf{R} _{ij})\mathbf{D}(\mathbf{R} _{ij})\mathbf{x} _{i}) \otimes (\mathbf{D}^{-1}(\mathbf{R} _{ij})\mathbf{D}(\mathbf{R} _{ij}) \mathbf{S}(\vec{\mathbf{r}} _{ij}))$. Based on the definition of $\mathbf{x}' _{i}$ and the Wigner-D matrix, this formula transforms into $(\mathbf{D}^{-1}(\mathbf{R} _{ij})\mathbf{x}' _{i}) \otimes \mathbf{D}^{-1}(\mathbf{R} _{ij})\mathbf{S}(\mathbf{C})$. Additionally. The CG product is equivariant for all rotations (refer to Equation 26 in the Appendix). Consequently, the above formula can be transformed into the right side of Equation 9.

In summary, equation 9 can not be affected by randomness because of the equivariance of CG product. The $\mathbf{D}^{-1}$ is the transpose of $\mathbf{D}$ in the implementation (They are based on the same $\mathbf{R}_{ij}$). In this case, **the outputs of the right side of equation 9 is the same when we choose different $\mathbf{R}_{ij}$. We deliberately break equivariance in Equation 11**, incorporating the non-equivariant MLP (You can refer to Responses to 7-th detailed comment). 

Reading reponse to the 5-th detailed comment together may better address your concerns.

Thank you for your valuable suggestions. We updated two parts in the new version: **1.We provide a more rigorous reference and description to the local reference frames. 2.We will add the disscusion that Equation 9 will not influence by randomness.**

Reviewer comment (Experiments):
A second weakness of the paper is the complexity of the model (figure 1) with all sorts of interactions and blocks. It is great that excellent results are obtained with the model, but is it really worth the engineering effort? In my opinion, the contribution of great results is greatly diminished if it comes at the cost of such an intricate (over-engineered?) network design. Then, what remains is the ablation that shows that relaxing equivariance could be beneficial. This is a message that contains scientific value, in my opinion, however, this point I think is not thoroughly discussed/analyzed.

Author response:
Figure 1 may appear intricate for two main reasons: 

1. We show most of the details to improve its reproducibility. Analogues can be found in other MPNNs, such as the Figure 2 in GemNet [1] and Figure 1 in Equiformer [2]. **While the design of many MPNNs may seem complex, it's important to note that most of modules are fixed. Therefore, we can bypass the complex hyper-parameter design in engineering and easily apply models to various molecular tasks.** Our approach is similar. In our ablation (Appendix E.4), we show that HDGNN mainly replies on four hyper-parameters ($K$, $C$, $L$, $L'$). 

2. **Some operations depicted in our Figure 1 share the concise high-level scientific ideas.** For instance, the neural network paths in Multi-TP share the similar non-equivariant design. **The modules in Figure 1 can be summarized to two scientific ideas: 1.the design of approximately equivariant modules; 2.the combination between equivariant modules and non-equivariant modules.** We discuss the motivation of ideas in introduction and section 2.4, the evolutions process of ideas in section 3.2 and 3.3. We have emphasized the core scientific ideas in new version and indicated their corresponding building blocks.

[1] Gasteiger J, Becker F, Günnemann S. Gemnet: Universal directional graph neural networks for molecules. Neurips 2021.

[2] Liao Y L, Smidt T. Equiformer: Equivariant graph attention transformer for 3d atomistic graph. ICLR 2023

Reviewer comment (Theory):
The main innovation of this paper, which is the partial equivariance, is not clearly explained, in my opinion. Apart from the error with respect to ill-defined rotation matrix estimation, the explanation of why one type of processes breaks equivariance and the other doesn't (Eq 11) could be better explained.

Author response:
Thank you for the suggestion. We will address this in the new version. We clarify that **two components in Equation 11 break equivariance.** This is because the non-equivariance of MLP (You can refer to Responses 5 to detailed comments). **The equivariant component in message block is METP operation.** Note that the final message is calculate by Equation 12. We use equations 8 and 11 to introduce the equivariant and non-equivariant parts of equation 12, respectively.

Reviewer comment (Experiments):
In general, I think the paper could still benefit from being more precise on the purpose of all included experiments (what precisely do they test for?).

Author response:
**The results in Table 1 and Table 2 represent the generalization ability on validation, which reveal whether models can accurately decipher the underlying physical or chemical principles inherent in molecular models. In our ablation experiments, we systematically examine how specific sub-modules in HDGNN influence the generalization ability.** In essence, generalization ability is influenced by both expressive power and equivariance. The modules in HDGNN impact these two aspects. When equivariance is severely damaged, generalization ability will also be severely reduced ("Unrestricted" in Table 3). Additionally, insufficient expressive power also constrains generalization ability. In Table 3, we employ strictly equivariant, approximately equivariant, and unconstrained operations to compute messages within our model framework. The results confirm that approximate equivariance is able to enhance generalization. Table 4 assesses our attention method, which is designed to dynamically adjust the ratio of non-equivariant to equivariant operations ($a _{ij}$ in Equation 12). This outcome illustrates our method's endeavor to strike a balance between expressive power and equivariance for improved overall generalization. Table 5 is dedicated to validating the detail structure or design idea of the core module. A more extensive evaluation is provided in the Appendix. These experiments serve to facilitate the reproducibility and customization of HDGNN.

Reviewer comment (Theory):
In the introduction it is mentioned "While equivariant neural networks exhibit appropriate inductive biases ... leading to inefficient learning of interactions between atoms." I do not agree with (or understand) this claim at all. Precisely these domains require equivariance: it is demanded by the physics. Any operation that is non-equivariant is bound to fail, unless it *learns to be equivariant*. The property of equivariance is either way essential and thus imposing equivariance should improve efficiency in learning. Perhaps the authors mean something like that it might be more efficient to start off unconstrained and let the network *learn equivariance* (through the task or augmentation). The paper is not clear on this; it is often suggested that one wants to break equivariance because it might not be the right inductive bias. Either way, I really do not understand what is meant with the statement, and I think it is wrong. Please give a counterexample where the predicted properties do not require equivariance/invariance, or update the sentence.

Author response:
Thank you for your review, we have updated this sentence to make a rigorous expression. We are sorry that our expression caused you misunderstanding and hope our rebuttal can clarify it. What we want to express is that the constraints in equivariant operation (sband-limit degree) can impact the model's expressive capacity, leading to inefficient learning. In this context, inefficient learning refers to the model's inability to capture accurate mathematical relationships between atoms.

Reviewer comment (Presentation):
In Section 2.4, the analysis of expressive power is, I think, more important than the benchmark results but is hidden in the appendix. Could the results be somehow briefly discussed in the main body, instead of an appendix that is not included in the main PDF?

Author response:
We have included the analysis of expressive power in the Appendix for two main reasons: 1. Given the constraints on text length, our primary aim is to maximize the introduction of our main contribution—HDGNN. While the analysis of expressive power serves as a crucial motivation, the intricate details of this aspect are complex. Consequently, we can only provide a concise overview of the core ideas. 2. The theorems related to expressive power in equivariant models are not our contributions. Excessive description in the main text may lead to misunderstanding of our contribution. Therefore, we only put the most important conclusions in the introduction and section 2.4.

Reviewer comment (Presentation):
Above equation 6 in the neighborhood definition please use $\mid$ (\mid) to leave some space between the left and right of |. Also equation 6 reads better if you include some white space after the comma.

Author response:
Thank you for your suggestion. We have updated additional details and discussions concerning $\mathbf{R} _{ij}$. In practice, we first produce normalized vector of $\vec{\mathbf{r}} _{ij}$, denoted as $\vec{\mathbf{a}}$. We then randomly sample an orthogonal vector $\vec{\mathbf{b}}$ of $\vec{\mathbf{r}} _{ij}$ on the sphere, and calculate another orthogonal vector $\vec{\mathbf{c}}$ through cross product. $\mathbf{R} _{ij} = [\vec{\mathbf{b}}^{T}; \vec{\mathbf{c}}^{T}; \vec{\mathbf{a}}^{T}]$. This approach ensures orthogonality of $\mathbf{R} _{ij}$ and is similar to produce $\mathbf{R} _{ij}$ by random $\varphi$.

Reviewer comment (Theory):
Please include a discussion or more details on the estimation of $R_{ij}$ from $\vec{r}_{ij}$, as mentioned in my main comments.

Author response:
**MLP breaks equivariance because its functional form is non-equivariant to SO(3)-transformation.** MLP does not conform to the definition in Equation 1, when $T(g)$ denotes SO(3)-transformation. In simpler terms, **when a 3-D vector $\mathbf{x}$ (or an irreducible representation) is input into MLP, the outcome violates the equality $NN(\mathbf{R} \mathbf{x}) \not= \mathbf{D} NN(\mathbf{x})$.** Here, $\mathbf{R}$ represents any rotation and $\mathbf{D}$ represents corresponding Wigner-D matrix. 

For example, we assume that the input $\mathbf{x}$ is based on the spherical harmonic basis. **We can use $\mathbf{D}^{-1}(\mathbf{R} _{ij})$ to rotate it back only when $NN(\mathbf{D}(\mathbf{R} _{ij}) \mathbf{x}) = \mathbf{D}(\mathbf{R} _{ij}) NN(\mathbf{x})$.** However, an untrained MLP cannot satisfy this condition (its output is not based on the spherical harmonic basis).

Randomness will amplify the impact of non-equivariance of MLP. When we replace $\mathbf{x}_i \otimes \mathbf{S}(C)$ with a untrained MLP in equation 9, different $\mathbf{R} _{ij}$ lead to different results. In this case, the equality of equation is not true, and the equivariance on the right side breaks down.

Reviewer comment (Theory):
Above equation 10: "However, the MLP structure introduces non-equivariance" — please explain this. I do not see why it would break equivariance. In other words, if equation 9 were true (if the rotation estimate were to be well-defined), then applying an MLP on $x_i'$ and then rotating it back afterwards would be fine, right? It would only break equivariance because of the ill-defined rotation estimation, but that is separate from the use of an MLP. So I wonder what is the issue here really? Could you explain this?

Author response:
Yes, the randomness is caused by free rotation angle.

Reviewer comment (Writing):
The same sentence continues with "... rotation matrix introduces randomness". Is this because of the free rotation angle, or what is meant here?

Author response:
We first clarify the equivariance of each term and purpose of each equation. Equation 10 is strictly equivariant, and $NN(\mathbf{x})$ and $NN(\mathbf{x}'')$ in equation 11 are both non-equivariant. Equation 9-11 are expansion of Equation 7 which is based on equivariant CG tensor product. We mentioned that Equation 7 is widely used in other equivariant models (the sentence after Equation 7). However, finite degree limits expressiveness. In equation 11, we aim to replace CG product by MLP to improve expressiveness.

Second, we introduce the idea of decomposition. If there is no decomposition, then Equation 9 will be approximated by $\mathbf{D}_{-1}(\mathbf{R} _{ij})NN(\mathbf{x}')$. This will cause two problems: 1. The randomness in $\mathbf{x}'$ will affect the learning of MLP. 2. We hope that MLP can capture directional features of edge. In equation 9, these features are stored in the wigner-D matrix ($\phi$ is random, but $\theta$ is well-defined). However, the input of MLP ($\mathbf{x}'=\mathbf{D}\mathbf{x}$) is a whole, and we cannot expect MLP to capture the information in $\mathbf{D}$ separately from this whole. Through the decomposition of equation 10 and equation 11, one of the branches $NN(\mathbf{x} _{i})$ is not affected by randomness. Moreover, we feed $(\mathbf{D}-\mathbf{I})\mathbf{x}$ and $\mathbf{x}$ to the MLP, making it possible to extract the directional features in $\mathbf{D}$. We discussed this decomposition in the sentences following Equation 11.

We add some content here to facilitate understanding the role of MLP. We use MLP to learn the CG product process $f(\mathbf{x}) = (\mathbf{x} \otimes \mathbf{S}(C))$. Due to MLP's ability to approximate any function, it can even extend the learned function to higher-degree CG tensor products. Table 6 verified that MLP can learn features in higher-degree representations.

Reviewer comment (Theory):
Equations 10 and 11 could benefit from more explanation. Why this decomposition? What problem does it solve? Is it to split the processing into an equivariant and non-equivariant part, through NN(x_i) or NN(x_i'')? Or what is the idea behind it?

Author response:
$\mathbf{a} _{ij}$ is the attention coefficient, they are generated by invariant branches. The formula is $MLP(\mathbf{f} _i, \mathbf{f} _j, \mathbf{c} _{ij})$, which corresponds to the green module in figure 1. Thank you for pointing this out, we have added a description of it in the new version. The attention coefficients are used to automatically adjust the weights of non-equivariant and equivariant operations. Due to limited space, the details of invariant branch $\mathbf{f}$ are put in Appendix B.1.

Reviewer comment (Writing):
In equation 12, what are the $a_{ij}$, and how are they computed?

Author response:
Equations 13 and 14 represent the Fourier transform on the sphere. Spherical harmonic can be regarded as frequency representation of Fourier transform on a sphere (Appendix A.1.2). The CG tensor product can be regarded as a special convolution process. Based on the convolution theorem, **we can use point-wise operations in the time domain to approximate the convolution process in the frequency domain.** Therefore, in equation 13, we first convert the signal in the frequency domain to the time domain, and then perform a point-wise MLP in the time domain. Finally, we convert it to the frequency domain through equation 14. We do this to reduce the amount of computation. **The purpose of the update block is to calculate the interaction between different channels of each node embedding. If we reference the CG product to each channel pair, the computational cost will become very high. This approach also cannot consider the interaction between three or more channels.** Although the point-wise MLP in the time domain is only an approximation, its complexity is much smaller and interactions between mutiple channels are considered.

Reviewer comment (Presentation):
Equations 13 and 14 describe a Fourier-based activation function, right? Whereas as "activation function" an MLP is used. I do like this approach, but I think the presentation could be simplified a bit.

Author response:
When we perform point-wise operations in time domain, the point refers to each direction vectors $\mathbf{r}$. Taking the convolution theorem as an example, assume that node embedding in the frequency domain has only two channels, denoted by $m _1$ and $m _2$. We use $p _1$, $p _2$ to denote their signals in time domain. We can use point-wise multiplication $p _1 \circ p _2$ to approximate convolution ($\circ$ is the Hadamard product). For each point $\mathbf{r}$, the output is $p _1(\mathbf{r}) * p _2(\mathbf{r})$. Now, we have C channels and use MLP as point-wise operation. The calculation at point $\mathbf{r}$ can be expressed as $MLP([p _1(r), p _2(r), ..., p _C(r)])$. Moverover, we extend to correlation so also include $\mathbf{-r}$.

Reviewer comment (Theory):
After equation 13: "we concatenate all the elements..." — why? It seems very ad-hoc, but perhaps there is a good explanation for doing this. Could you explain why this is done?

Author response:
We regard CG tensor product as a special convolution because some paths in CG tensor product are similar to the form of convolution. However, the paths outside of the convolution will be missing. We supplement the missing path through correlation. For example, if we calculate CG product between two irreducible representation $m_1$, $m_2$. We assume $L=4$, and $m_1$, $m_2$ and $m'=m_1 \otimes m_2$ contain type-0, type-1, $\cdots$, type-4 vectors. (type-$ l_1$, type-$l_2$) denotes a path based on $m_1$ and $m_2$. Taking type-2 vector of $m'$ as an example, it is the aggregation of paths (1, 1), (0 , 2), (2, 0), (1, 3), (3,1), (4,2), and (2,4) (refer to equation 5). However, convolution only contains (1, 1), (0, 2), and (2, 0) paths (refer to equation 32). This mehtod will weaken the equivariant interactions they learn. We introduce correlation whose corresponding paths are (0, 2), (2, 0), (1,3) , (3,1), (4,2), and (2,4), including paths that are not belong to convolution. In equation 5, convolution/correlation corresponds to $(l_{1}+l_{2}= l)/|(l_{1}-l_{2}|=l)$ paths. The conversion of correlation to the time domain is an operation between ($\mathbf{r}$) and ($\mathbf{-r}$). Therefore, we concatenate the ($\mathbf{r}$) and ($\mathbf{-r}$) of all channels in the time domain.

**Although the difference between convolution and correlation is only one negative sign, they correspond to different paths of CG tensor product in the frequency domain.** We discussed this in appendix B.3. Table 11 also proves that introducing correlation can improve performance.

Reviewer comment (Theory):
"Here, we extend the point-wise operation to two points by incorporating both convolution and correlation" — what does that mean? Convolution and correlation are one and the same thing, up to a flip of the kernel. This sentence makes no sense to me; please clarify what is meant by it.

Author response:
Thanks for your reviews, we have updated this sentence.

Reviewer comment (Evaluation):
Below table 1: "This highlights the ability of HDGNN to capture more general physical properties in real world scenarios." I do not see this claim substantiated. The only thing I see from the table is that the method outperforms several others (not always), but this does not imply a better capability to "capture more general physical properties". What are these "more general properties"? There are so many differences between your architecture and those in the table. In my interpretation, HDGNN might just be better engineered to solve the task but does not give any more insights than that.

Author response:
These scores ​​are the average of OOD tasks (OOD ads, OOD Cat, and OOD Both). The average scores in the Table 1 is the average of all tasks (ID, OOD ads, OOD Cat, OOD Both).

Reviewer comment (Presentation):
What are the ODD average scores between parenthesis in text (GDGNN: 564 vs SCN/Equiformer: 592/613)? I cannot find these in the table (sorry if I overlook them somehow).

Author response:
Please refer to our response to 4-th main comments.

Reviewer comment (Experiments):
"The QM9 dataset ... susceptible to overfitting" — I do not see why this is an issue. Did you indeed observe overfitting in your learning curves? If you want to test these kind of overfitting issues, perhaps the revised rMD17 would be more appropriate. In fact, I think rMD17 might be the best dataset to test your method, because the task in rMD17 (interatomic potential energy prediction, for a single molecule given different conformations) is purely driven by geometry and should be purely invariant.

Author response:
We have observed overfitting in all the tasks of QM9. For example, The final training MAE is much lower than the validation MAE (in the strictly equivariant model, they are similar, you can refer to the open training log of Equiformer). The reason for overfitting is that the dataset of QM9 is relatively small, and HDGNN cannot learn the equivariance under SO(3) transformation. If we perform SO(3) transformation augmentation, this overfitting may be alleviated. However, this method will bring an unfair comparison (other methods use the original dataset). Thanks for your suggestion, we will evaluate rMD17 in the open source community in the future (sorry that the available resources are currently limited). QM9 and IS2RE are also purely invariant, and lead a comprehensive comparison to SOTA molecular models. Many advanced methods have published code and results on them.

Reviewer comment (Theory):
"In contrast, our approach merges inherent equivariance with learned equivariance." What mechanism induces learned equivariance? This is not discussed in the paper.

Author response:
The MLP in equation 12 introduces learned equivariance. The METP introduces inherent equivariance.

Reviewer comment (Evaluation):
"While HDGNN has shown superior performance compared to strictly equivariant models" — did it? It seems sometimes it does, sometimes it doesn't. The claim should be nuanced.

Author response:
Thanks for your suggestion. We have updated.

Reviewer comment (Evaluation):
On several occasions the paper motivates the equivariance relaxation from the efficiency perspective. But what about actual speed efficiency? I see no comparison in the paper and it would be nice to see if something is gained on that level.

Author response:
Thank you for your suggestion. The experiments assessing speed efficiency are presented in Table 13 in the Appendix. Furthermore, **we have added the complexity analysis of our operations in Appendix B.2 and B.3.** This update includes a comparative assessment of the efficiency between our proposed structure and the CG tensor product.