Paper: Stabilized E(n)-Equivariant Graph Neural Networks-assisted Generative Models (ICLR-2024)

Reviewer comment (Evaluation):
This work analyzes the gradient update of a single-layer EGNN but this reviewer does not see an analysis of special reasons leading to unstable ENF and EDM training. Only a single sentence states that “In both ENF and EDM, the graph node coordinates keep changing during the generative process, and abnormal coordinate updates may occur.” However, the generative process is the inference process while this paper is dealing with the training instability. Therefore, the ENF and EDM training instability is not accurately summarized. Is EGNN on the given training data stable? Is ENF and EDM unstable on the given training data? If the answers to both questions are yes, then this might indicate that there are special reasons resulting in the instability of ENF and EDM training. If ENF and EDM training instability is caused by EGNN training only (no special stuff), then the discussions over the ENF and EDM are redundant and not necessary. In general, this reviewer thinks the preliminary discussions over the training instability are not sufficient.

Author response:
**Q2. What are the special reasons for numerical instabilities in ENF and EDM training? Or the instability of ENF and EDM is all caused by EGNN?**

**Reply:** The instability issue of ENF and EDM we study originates from the EGNN model in the training phase, and we have pointed this out in the first paragraph of Section 1.2 of our paper. We focus on ENF and EDM for the following two main reasons:

- Though EGNN is a general E(n)-equivariant graph neural network, one of the most remarkable applications of EGNN is for equivariant generative modeling. If we look at the citations of the EGNN paper, a very large portion are related to equivariant generative modeling.

- Besides stabilizing training ENF and EDM, the proposed regularization scheme brings additional important benefits to ENF and EDM. 
For ENF, because solving a high-dimensional ODE is required in both training and testing, the proposed regularization can also regularize the Lipschitz constant of the ODE to improve the computational efficiency of ENF. For EDM, the proposed regularization allows EGNN to use an unnormalized coordinate update, and such an unnormalized coordinate update can improve the performance of EDM. 

In the revision, we have further improved the discussion of the instability issues of training ENF and EDM in Section 1.2 and Section 3.

Finally, in the concluding remarks section, we have added the discussion that the training instability in ENF and EDM we study comes from the EGNN itself. Studying the efficacy of the proposed regularization for other EGNN applications can be an interesting future work.

Reviewer comment (Experiments):
The proposed regularizer is a sum of gradient calculations, which is extremely hard to compute. The summations are computed over all atom pairs (i, j) and EGNN layers L with at least O(N**2*L) time complexity. Hence, although the proposed regularizer is a very effective approach to restricting the gradient norm, the calculation of the proposed regularizer could not be scalable to either larger particle systems or deep EGNN architectures.

Author response:
**Q1. What is the time complexity of the calculation of the proposed regularizer? How could the proposed approach scale to larger particle systems and deep EGNN architectures?**

**Reply:** The graphs considered in the molecular generation tasks are sparse, and therefore the time complexity of the regularization term is $O(NL)$ with $N$ and $L$ being the number of particles and the depth of EGNN, respectively. 

All summations are taken over the neighborhood of a given graph node, we have fixed all inappropriate notations in the revision.

In our paper, we have also compared the computational time of the EDM with and without regularization, showing the extra computational time is not very significant. The computational overhead due to the regularization term is much smaller than the baseline computational time.

Reviewer comment (Theory):
What is the time complexity of the calculation of the proposed regularizer? How could the proposed approach scale to larger particle systems and deep EGNN architectures?

Author response:
**Q1. What is the time complexity of the calculation of the proposed regularizer? How could the proposed approach scale to larger particle systems and deep EGNN architectures?**

**Reply:** The graphs considered in the molecular generation tasks are sparse, and therefore the time complexity of the regularization term is $O(NL)$ with $N$ and $L$ being the number of particles and the depth of EGNN, respectively. 

All summations are taken over the neighborhood of a given graph node, we have fixed all inappropriate notations in the revision.

In our paper, we have also compared the computational time of the EDM with and without regularization, showing the extra computational time is not very significant. The computational overhead due to the regularization term is much smaller than the baseline computational time.

Reviewer comment (Theory):
What are the special reasons for numerical instabilities in ENF and EDM training? Or the instability of ENF and EDM is all caused by EGNN?

Author response:
**Q2. What are the special reasons for numerical instabilities in ENF and EDM training? Or the instability of ENF and EDM is all caused by EGNN?**

**Reply:** The instability issue of ENF and EDM we study originates from the EGNN model in the training phase, and we have pointed this out in the first paragraph of Section 1.2 of our paper. We focus on ENF and EDM for the following two main reasons:

- Though EGNN is a general E(n)-equivariant graph neural network, one of the most remarkable applications of EGNN is for equivariant generative modeling. If we look at the citations of the EGNN paper, a very large portion are related to equivariant generative modeling.

- Besides stabilizing training ENF and EDM, the proposed regularization scheme brings additional important benefits to ENF and EDM. 
For ENF, because solving a high-dimensional ODE is required in both training and testing, the proposed regularization can also regularize the Lipschitz constant of the ODE to improve the computational efficiency of ENF. For EDM, the proposed regularization allows EGNN to use an unnormalized coordinate update, and such an unnormalized coordinate update can improve the performance of EDM. 

In the revision, we have further improved the discussion of the instability issues of training ENF and EDM in Section 1.2 and Section 3.

Finally, in the concluding remarks section, we have added the discussion that the training instability in ENF and EDM we study comes from the EGNN itself. Studying the efficacy of the proposed regularization for other EGNN applications can be an interesting future work.

Reviewer comment (Novelty):
The contribution is incremental. The contribution is just dividing by a normalization factor the well known EGCL, and adding a regularization factor to training. These are natural modifications of EGCL, and any other network, and it is suspected that many practitioners apply such tricks often.

Author response:
**Reply:** We respectfully disagree that our contribution is incremental. We fear there are some misunderstandings of our contribution. Please allow us to clarify our contribution below:

- First, dividing by a normalization factor, or using the normalized coordinate update is not our contribution, which is a default use in ENF and EDM. Moreover, ***the regularization we proposed is not by adding a factor to the denominator in the normalized coordinate update***. Instead, our proposed regularization scheme is given in Equation 4 in our paper.

- Second, the ***analysis of the instability issues in training ENF and EDM*** Section 3 is our first contribution. In particular, we show that the training instability remains even when using the normalized coordinate update. By performing a sensitivity analysis in Section 3.2, we identify a source of the training instability issue. In particular, the term $\frac{\partial {\bf x}\_i^{l+1}}{\partial {\bf x}\_j^l}$ is the most sensitive term to the coordinate update in EGNN during the training process.

- Third, in Section 4, based on the expression of $\frac{\partial {\bf x}\_i^{l+1}}{\partial {\bf x}\_j^l}$, we propose a new regularization term to improve the training of ENF and EDM. ***The new regularization is given by $\sqrt{\sum_{i,j\in\mathcal{N}(i),l}( \frac{\partial \phi_x({\bf m}_{ij}^l) }{\partial  ||{\bf x}_i^l-{\bf x}_j^l ||^2  }  )^2 }$, which is novel and highly nontrivial.*** 
Besides stabilizing training ENF and EDM, we further show that ***the proposed regularization can accelerate the training and testing of ENF and improve the performance of EDM***.
In Section 5, we numerically verify the efficacy of the proposed regularization and confirm our theoretical results.

We believe the above contributions are significant.

Reviewer comment (Theory):
Page 4: “In contrast, normalizing coordinates can avoid abnormal coordinate updates from large differences among coordinates of neighboring nodes” - but on the flip side, without the $+1$ regularization in the denominator, it is unstable to small coordinates. But with the $+1$ normalization, close-by nodes contribute a very small difference. How do you then choose the scale of the coordinates for the $+1$ to work well? Why do you use $+1$ and not $+b$ for some $b$ that depends on the characteristic target distance between modes?

Author response:
**Reply:** +1 in the denominator is not the regularization scheme we proposed, it was used in the original ENF and EDM models. Again, our proposed regularization is given by Equation 4.

Reviewer comment (Reproducibility):
Please explain how you construct a graph from the node locations and features.

Author response:
**Reply:** We follow the original ENF paper to construct the graphs for DW4 and LJ13. For QM9, the graph is given in the benchmark datasets; see reference [1] listed below for details. Specifically, the graph is padded with nodes containing zero-valued features until all molecules share the same number of nodes. Then the graph edges are constructed by connecting all nodes with non-zero features. We follow this same graph construction procedure from the baseline ENF and EDM models. The exact code snippet is provided in reference [2] listed below and is executed when the graphs are collated into mini-batches for training.

[1] Anderson et al. Cormorant: Covariant Molecular Neural Networks. NeurIPS 2019.

[2] https://github.com/vgsatorras/en_flows/blob/main/qm9/data/collate.py#L78C4-L83C64

Reviewer comment (Theory):
Section 3.2 Sensitivity Analysis: The normalized EGCL is different from the unnormalized one. Why don’t you compute the derivatives of the normalized version if this is the method you propose? Also, it is strange to directly write the derivative of $\phi_x$ with respect to $\|x_i-x_j\|$. You need to use the chain rule, and first differentiate with respect to $m_{i,j}$.

Author response:
**Reply:** We first stress again that the normalized version of EGCL is not what we proposed. 

The notation $\frac{\partial \phi_x }{\partial ||{\bf x}_i-{\bf x}_j||^2}$ is clear while its detailed expression is obtained by using the chain rule. In what follows, we would like to provide basic examples to illustrate this.

Let $w=f(z)$, $z=g(y)$, and $y=h(x)$. Though the exact expression $\frac{\partial w}{\partial x}$ is $\frac{\partial w}{ \partial z}\cdot \frac{\partial z}{ \partial y}\cdot \frac{\partial y}{ \partial x}$, in practice, people will simply write $\frac{\partial z}{\partial x}$ rather than always write down the exact expression, which is rather redundant.

Reviewer comment (Writing):
Notations: you did not define $I_3$.

Author response:
**Reply:** Thank you for pointing this out, $I_3$ is the $3\times 3$ identity matrix. We have clarified this in the revision.

Reviewer comment (Theory):
Section 4: The derivative $\partial f^{l+1}/\partial f^l$ is not the only main contributor to the size of $\partial L (f^L)/\partial \theta$. When partitioning the parameters to the last layer parameters $\theta^L$ and the previous layers parameters $\theta^{L-1}$, $\partial L (f^L)/\partial \theta$ has two components: first, $\partial L/\partial f^L \cdot  \partial f^L\partial \theta^L$, and then $\partial L/\partial f^L \cdot  \partial f^L\partial f^{L-1}\cdot \partial f^{L-1}\partial \theta^{L-1}$. Hence, $ \partial f^{l}/\partial \theta^{l}$ are also important, and these are roughly going to depend on $\|x_i-x_j\|^l$ by induction, as each later multiplies by a factor of order $\|x_i-x_j\|$. This is another main reason to normalize $x_i-x_j$.

Author response:
**Reply:** We respectfully disagree that our analysis is partial. Notice that in Equation 6 of our paper, we have clearly stated that ***$\theta$ is the weight parameter appearing at the $l$-th layer of EGNN.***

We would like to point out again that our sensitivity analysis aims to identify the most sensitive term with respect to $||{\bf x}_i-{\bf x}_j||$ in the gradient $\frac{\partial L}{\partial \theta}$. Through this sensitivity analysis, we propose the regularization scheme in Equation 4.

Reviewer comment (Theory):
The analysis in this paper is partial. There is no systematic analysis of all components that contribute to the gradient.

Author response:
**Reply:** We respectfully disagree that our analysis is partial. Notice that in Equation 6 of our paper, we have clearly stated that ***$\theta$ is the weight parameter appearing at the $l$-th layer of EGNN.***

We would like to point out again that our sensitivity analysis aims to identify the most sensitive term with respect to $||{\bf x}_i-{\bf x}_j||$ in the gradient $\frac{\partial L}{\partial \theta}$. Through this sensitivity analysis, we propose the regularization scheme in Equation 4.

Reviewer comment (Theory):
The analysis only examines the unnormalized layer. If the goal is to show that the normalized layer is more stable, a gradient analysis of the normalized layer should also be provided, along with the full gradient with respect to the model parameters, as these are the gradients used in training.

Author response:
**Reply:** We believe there is a misunderstanding, we are not aiming to show the normalized layer is more stable than the unnormalized layer. ***What we are showing is that using no matter normalized or unnormalized coordinate updates, training ENF and EDM suffers from the instabilities issue.***

In our paper, due to the page limit, we only write down the results for the unnormalized coordinate update. For the normalized coordinate update, a similar result can be derived by merging the normalization factor into $\phi_x({\bf m}_{ij}^l)$ in Equation 5.

Reviewer comment (Experiments):
Experiments: why are EDMs and ENF not compared against each other and other models in molecule generation in one table? Only a comparison between EDM and other methods is shown.

Author response:
**Reply:** In our original submission, we separated the results of ENF and EDM into two tables since ENF (normalizing flows) and EDM (diffusion models) are quite different models. In the revision, we have merged the two tables.