Paper: Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields (ICLR-2024)

Reviewer comment (Writing):
Abbreviations should not be used before they are introduced. The abstract already refers to "S2EF" and "IS2RE results" (clarified in the numerical results), refers to network architectures without citation and even uses variables ($L_\max = 2$ and $L_\max=3$) whose meaning is assumed to be known.

Author response:
Thanks for pointing these out! We have updated the abstract.

Reviewer comment (Theory):
It is not explained why forces make the denoising problem on non-equilibrium structures well-posed and also not how the forces are obtained. Is it correct to assume that the potential of any structure can be computed and that the forces are the gradient thereof? If so, it would be just two additional sentences of explanation that make the work much more clear.

Author response:
The ground-truth force labels are obtained from DFT / MD calculations depending on the dataset and are from the training set.

For predicting forces, as mentioned in Section 3.1, we either take the gradient of potential energy or directly predict them from node embeddings.

Restating from the introduction and Section 3.2.2, denoising non-equilibrium structures without encoding forces as input is ill-posed because there are many possible target structures to denoise to. This is unlike the case for equilibrium structures which are at local minima, and so there exists a unique target structure to denoise to given perturbed structures. Encoding atomic forces in DeNS helps identify the unique non-equilibrium structure to denoise to. This is also illustrated in Figure 1. More specifically, since we train GNNs with $S_{\text{non-eq}}'$ and $F(S_{\text{non-eq}})$ as inputs and $\text{Noise}(S_{\text{non-eq}}, S_{\text{non-eq}}')$ as outputs, they implicitly learn to leverage $F(S_{\text{non-eq}})$ to reconstruct $S_{\text{non-eq}}$ instead of predicting any arbitrary non-equilibrium structures. (Here we use $S_{\text{non-eq}}'$ to denote corrupted non-equilibrium structures.)

Reviewer comment (Writing):
The property "equivariant" is frequently used. While I know what it means, it was unclear to me what kind of equivariance (with respect to which transformations) is meant and why that influences the ability to encode forces. Rotations and translations of 3d coordinates/vectors?

Author response:
Yes, the equivariance in this work refers to equivariance to 3D rotation and invariance to 3D translation, which is commonly implied in the context of learning force fields. Besides, we also discuss some equivariant networks and provide pointers to detailed background in Section A.1.

As mentioned in Section 3.2.2, the node embeddings of equivariant networks consist of different type-$L$ vectors, where $L$ denotes degree, and forces belong to type-1 vectors. Therefore, to encode forces, we can simply add forces to the part of type-1 vectors in node embeddings. 

Please let us know if we can help clarify anything else.

Reviewer comment (Writing):
I do not know the term "type-L vectors" - what does it mean? I tried googling but did not have direct success. Thus, I'd recommend defining it. Also, what is "the projection of $f$ into type-L vectors with spherical harmonics"? Representing the function f with L-many coefficients corresponding to spherical harmonics basis functions?

Author response:
Type-$L$ vectors are (2$L$+1)-dimensional vectors that rotate $L$ times faster than 3D Euclidean vectors when we rotate the 3D coordinates. They can be obtained by taking tensor products of 3D relative positions (type-1 vectors). 

The projection of $f$ (type-1 vector) into type-$L$ vectors with spherical harmonics means that we use spherical harmonics $Y^{(L)}$ to transform type-1 vectors into type-$L$ vectors so that we can consider higher frequency (higher $L$).

We note that the details can be found in [1] as mentioned in Section A.1.  

Reference:  
[1] Liao et al. Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs. ICLR 2023.

Reviewer comment (Writing):
It would have helped me to cite something for "SO(3) linear layer".

Author response:
Thanks! We will add the citation of e3nn [1] after SO(3) linear layers. 

Reference:  
[1] Geiger et al. e3nn: Euclidean Neural Networks. ArXiv 2022.

Reviewer comment (Writing):
Please double-check the manuscript for typos.

Author response:
Thanks! We will fix them.

Reviewer comment (Experiments):
The ablation in Table 1 (b) indicates that $p_{DeNS}=0.5$ is the best, but also the largest tested value. Wouldn't it make sense to ablate values $>0.5$?

Author response:
We added experiments with $p_{\text{DeNS}}$ = 0.75 for 12, 20 and 30 epochs in Table 1(b). $p_{\text{DeNS}}$ = 0.5 is still the best hyperparameter choice, but this makes Table 1(b) more complete. Thanks for the suggestion!

Reviewer comment (Experiments):
Maybe could add more results from the denoising framework with other backbones to provide a comprehensive understanding of its performance.

Author response:
We evaluate EquiformerV2 [1] and eSCN [2] on OC20 S2EF-2M dataset, Equiformer [3] and SEGNN [4] on MD17 dataset, and EquiformerV2 on OC22 dataset. These architectures improve the SoTA on each dataset. Of course more can always be done, but we believe 4 architectures spanning 3 datasets provides a sufficiently general evaluation of the proposed denoising approach. As a point of comparison, prior works cover less breadth in architectures – Noisy Nodes [5] only uses GNS [6], pre-training via denoising [7] uses GNS [6], GNS-TAT and TorchMD-NET [8], and fractional denoising [9] only conducts experiments with TorchMD-NET [8].

Reference:  
[1] Liao et al. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. ArXiv 2023.  
[2] Passaro et al. Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs. ICML 2023.  
[3] Liao et al. Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs. ICLR 2023.   
[4] Brandstetter et al. Geometric and Physical Quantities Improve E(3) Equivariant Message Passing. ICLR 2022.  
[5] Godwin et al. Simple GNN Regularisation for 3D Molecular Property Prediction and Beyond. ICLR 2022.   
[6] Sanchez-Gonzalez et al. Learning to simulate complex physics with graph networks. ICML 2020.  
[7] Zaidi et al. Pre-training via Denoising for Molecular Property Prediction. ICLR 2023.   
[8] Thölke et al. TorchMD-NET: Equivariant Transformers for Neural Network based Molecular Potentials. ICLR 2022.  
[9] Feng et al. Fractional Denoising for 3D Molecular Pre-training. ICML 2023.

Reviewer comment (Novelty):
I believe that developing a general AI-based molecular dynamics (MD) method is more crucial than specifically designing and tuning for the OC dataset. A more generalized approach could be beneficial to the community by focusing on the broader application of deep learning-based MD methods across different systems, such as drug molecules, crsytal materials or polymers.

Author response:
First, we show that DeNS improves performance on MD17 (see **General Response 3** and Table 4 in the paper).

Second, we do not incorporate any constraints related to catalysis into DeNS, and therefore DeNS can be applied to all non-equilibrium structures, including those in MD simulations. From that perspective, DeNS is indeed generally applicable. Moreover, as mentioned in Section D.1 and D.2 (partial denoising, decaying the DeNS coefficient), DeNS can be specifically tuned for MD to further improve performance.

Finally, now that we have a general method that seems to work well, we agree that pushing further on MD simulations across a broader range of chemistries (drug molecules, crystal materials, polymers) is the logical next step for future work!

Reviewer comment (Experiments):
If the improvements are not solely due to data augmentation, but rather are attributed to the inverse or dual denoising setting, it would be valuable to explore the generalization ability of the S2EF system in other systems.

Author response:
The work of Noisy Nodes [1] has explored this. Comparing the first two rows in  Table 5 in [1], they find that simply adding data augmentation without denoising results in small performance gain, suggesting denoising is necessary. Besides, in Figure 3(B) and 3(C) in [1], they also show that adding noise as data augmentation results in strong regularization.

Our validation and testing S2EF evaluations on OC20 and OC22 already test for generalizability since some of the splits involve out-of-distribution catalysts or adsorbates not seen during training. 

Reference:  
[1] Godwin et al. Simple GNN Regularisation for 3D Molecular Property Prediction and Beyond. ICLR 2022. https://arxiv.org/abs/2106.07971

Reviewer comment (Presentation):
The tunable standard deviation (σ) denoising strategy appears to be crucial, it might be beneficial to incorporate visualizations of σ during the training process for better illustration.

Author response:
Thanks for the suggestion! We visualize how structures in OC20, OC22 and MD17 datasets become after adding noise of different scales in Section F in the appendix of the revision.

Reviewer comment (Novelty):
The original idea is not groundbreaking: auxiliary tasks are known, the specific case of denoising as well, here the novelty is only in feeding also the forces as input.

Author response:
We compare the contributions of previous works and this work in **General Response 2**. In summary, the main contribution is that we generalize denoising to all atomistic datasets, including both non-equilibrium and equilibrium structures, with a simple modification of force encoding. All prior works on denoising do not encode forces, which is an ablation of our proposed approach and performs worse (Table 1(e) on OC20 and **General Response 4** on MD17).

Reviewer comment (Evaluation):
It is not always clear from the results shown in tables that the proposed method improves the SOTA significantly.

Author response:
Training with DeNS improves state-of-the-art models on all the datasets we consider. The improvements are clear on MD17 (Table 5 and **General Response 3**), OC20 S2EF-2M (Table 1(d) and **General Response 5**) and OC20 S2EF-All+MD datasets (Table 2 and **General Response 5**). The performance gain is sometimes less on OC22 (Table 4 and **General Response 6**). The main point is that with this simple force encoding, we can generalize denoising to the broader set of non-equilibrium structures and improve all the state-of-the-art models.

Reviewer comment (Evaluation):
The paper has some typos, but most importantly, experiments are discussed very quickly (too quickly). More space should be devoted to discuss results.

Author response:
Thanks! We will fix them.  

We move the results of different hyper-parameters on OC22 (Table 4) to the appendix and incorporate some discussions on the results in **General Response 3, 5, 6** into the manuscript.

Reviewer comment (Novelty):
Since the results are not strikingly better when using DeNS on the SOTA, and given it involves a number of additional hyperparameters (that admittedly do not need very narrow fine-tuning, but still involve more work and potential for problems), it is not clear yet whether the contribution would be used widely.

Author response:
We disagree on results not being substantially better than SoTA. To restate, please see **General Response 3** (MD17), **General Response 5** (OC20) and **General Response 6** (OC22).

Any new method involves at least some additional hyper-parameters [1, 2, 3, 4]. The hyperparameters for DeNS are easily transferable across architectures without much tuning. Specifically, when training eSCN with DeNS onOC20 S2EF-2M (Table 1(d)), we use the same hyper-parameters as training EquiformerV2, and that achieves better performance, matching the results of EquiformerV2 without DeNS and saving 1.94$\times$ training time without any hyper-parameter tuning. Similarly, for MD17 in Table 5, we directly use the same hyper-parameters tuned for Equiformer and training with DeNS improves two architecture variants on all the molecules. Besides, some hyper-parameters such as $p_{\text{DeNS}}$ and $\sigma_{\text{high}}$ also apply across datasets like OC20 and OC22, and this makes tuning hyper-parameters easier. 

Reference:  
[1] Wang et al. Denoise pretraining on nonequilibrium molecules for accurate and transferable neural potentials. Journal of Chemical Theory and Computation, 2023.   
[2] Godwin et al. Simple GNN Regularisation for 3D Molecular Property Prediction and Beyond. ICLR 2022.  
[3] Zaidi et al. Pre-training via Denoising for Molecular Property Prediction. ICLR 2023.  
[4] Feng et al. Fractional Denoising for 3D Molecular Pre-training. ICML 2023.

Reviewer comment (Experiments):
Table 1d is probably the most convincing result, as differences between models seem more significant. It is unclear whether it is computed on the validation set(s) or train set.

Author response:
The force and energy MAEs in Table 1(d) are all on the validation set as mentioned in the caption.

Reviewer comment (Experiments):
Trainings are very heavy.

Author response:
Our work improves sample efficiency, thus **reducing the training cost**. In Table 1(d) and Table 4, we show that training with DeNS is better than scaling up computationally expensive hyperparameters like $L_{max}$ and training epochs, resulting in 2$\times$ to 3$\times$ better training efficiency. The training time of baseline models has little to do with this work.

Reviewer comment (Experiments):
There is some hyper-parameter tuning, and DeNS is used on top of the best models (equiformerV2).

Author response:
Please refer to **7. in Response to Reviewer AVUJ** for hyper-parameter tuning. 

The reason we use the Equiformer series is because they are the current state-of-the-art and to see how well DeNS can improve training efficiency and performance compared to trivially scaling up computationally expensive hyper-parameters like maximum degrees and training epochs as done in their papers.

Reviewer comment (Evaluation):
Most of the differences reported are not very significant. The suggestion is to outline in green when a metric is significantly better in one model than in others.

Author response:
Please see **General Response 3** (MD17), **General Response 5** (OC20) and **General Response 6** (OC22). We will state the takeaways more clearly in the paper.

Reviewer comment (Evaluation):
If the method is mostly able to speed up training to achieve equivalent accuracy, this should be explicitly stated (even though it will weaken the paper's claim, it will at the same time strengthen the submission by clarifying the contribution).

Author response:
The training time reported in this work is similar to that reported in previous works [1, 2]. The proposed method aims at and is capable of reducing the training time to achieve better results.

Reference:  
[1] Passaro et al. Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs. ICML 2023.  
[2] Liao et al. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. ArXiv 2023.

Reviewer comment (Evaluation):
Table 2 does not show a significant improvement from using DeNS.

Author response:
We disagree. DeNS improves performance on all datasets we considered. Please see **General Response 3** (MD17), **General Response 5** (OC20) and **General Response 6** (OC22) for discussions on the performance. The proposed method can improve both training efficiency and error instead of only aiming at shortening training time.

Most of the investment in computational cost is one-time to exhaustively evaluate design choices for reporting in this paper. We believe the identified training recipe and hyperparameters are generally applicable and will help practitioners improve accuracy and reduce training cost on their datasets.

Reviewer comment (Evaluation):
The improvement on the IS2RE task is limited, and in some cases DeNS even hurts performance. This calls into question the real-world applicability of the performance gains claimed.

Author response:
With the updated OC20 All+MD results, EquiformerV2 with DeNS is better on all validation and testing S2EF splits, and about the same on IS2RE, compared to EquiformerV2 without DeNS. Consistent with findings in prior works, where eSCN achieves lower energy and force MAE but higher IS2RE MAE than SCN, we find that improvements on S2EF don’t always translate to improvements on IS2RE. We posit that IS2RE performance might be bottlenecked by other design choices when running structural relaxations (optimizer settings, stopping criterion, etc.) and not S2EF performance. We will look into this in future work. 

For OC22, please see **General Response 6**.

Reviewer comment (Reproducibility):
In each iteration, the model uses either standard training or DeNS training. As DeNS training requires the forces to be encoded into the input atom features, clarification is needed on how the force encoding module would be used in standard training, which instead uses the forces as labels.

Author response:
During standard training, we do not feed in the forces. Specifically, the force encoding module is an SO(3) linear layer, which is fed a tensor filled with all zeros during standard training. Force encoding is only used during DeNS training. We do not use any force labels in the validation and testing sets.

Reviewer comment (Experiments):
In Table 5, the authors show that DeNS training is more efficient and results in larger performance gains than increasing max degrees of irreps. However, they changed the model from EquiformerV2 to EquiformerV1 for this investigation. Since EquiformerV2 is claimed to largely benefit from scaling the max degrees of irreps, clarification is needed on how EquiformerV2 with different max degrees of irreps would behave in the same setting as Table 5.

Author response:
We use Equiformer(V1) simply because EquiformerV2 [1] does not report results on MD17 and that makes the comparison difficult.

The benefits of using higher degrees depend on datasets. As we can see in Table 5, on MD17, increasing $L_{max}$ from 2 to 3 in Equiformer(V1) leads to small gains in energy and force MAE. Therefore, this makes increasing $L_{max}$ further and using EquiformerV2 less well-motivated. It is likely that for this small dataset, using higher $L_{max}$ just leads to margin gains. However, training with DeNS achieves better energy and force errors across all the molecules and improves training efficiency by 3.1$\times$ compared to increasing $L_{max}$ from 2 to 3.

[1] Liao et al. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. ArXiv 2023.

Reviewer comment (Experiments):
The concept of “encoding force as input” finds extensive discussion in the paper. However, this approach needs corroborative proof from experiments. Results from Table 1e indicate energy prediction outcomes remain the same without force encoding. The possibility of label leakage and its contribution to improvement in force prediction, as discussed under question 2, needs examination.

Author response:
We conduct additional ablations with and without force encoding on MD17 in **General Response 4**. In summary, denoising with force encoding achieves the best results across all molecules. Denoising without force encoding is worse than no denoising for most molecules and is worse than with force encoding.

On OC20 S2EF-2M (Table 1(e)), indeed the energy MAE is the same with and without force encoding, but the force MAE is significantly worse (~4%) without force encoding. We should look at both energy and force prediction at the same time when comparing two methods since there is always a trade-off between energy MAE and force MAE. For example, we can slightly increase energy coefficient so that we further reduce energy MAE but increase force MAE. In this case, we can have both lower energy and force MAE when we use force encoding. Moreover, using force encoding also leads to significantly better sample efficiency – training DeNS with force encoding for 12 epochs improves force MAE by 4% compared to without force encoding (Index 1 and Index 2 in Table 1(e)), as opposed to training EquiformerV2 without DeNS for 20 epochs instead of 12 (3% better force MAE for 66% more compute, Table 1(d)), suggesting that force encoding is quite important.

Reviewer comment (Novelty):
Similar problem-solving approaches have been published. A comparative discussion highlighting the distinctions and superiority of this paper's proposed methodology would prove advantageous.

Author response:
Please see **General Response 2** for a comparison to previous works. 

The main difference in our denoising formulation lies in that we additionally encode forces as input. This generalizes denoising to non-equilibrium structures since the atomic forces help identify the unique non-equilibrium structure to denoise to. Without conditioning on forces, the problem of denoising non-equilibrium structures is ill-posed since there are many possible target structures to denoise to (see Figure 1).

Previous works [1, 2, 3] are designed for equilibrium structures and are the same as our ablation of denoising without force encoding. While [4] applies the same approach to non-equilibrium structures in ANI-1 and ANI-1x datasets, we show that denoising non-equilibrium structures without force encoding is worse than with force encoding, and can sometimes even be worse than training without denoising altogether (**General Response 4**). 

Reference:  
[1] Godwin et al. Simple GNN Regularisation for 3D Molecular Property Prediction and Beyond. ICLR 2022.  
[2] Zaidi et al. Pre-training via Denoising for Molecular Property Prediction. ICLR 2023.   
[3] Feng et al. Fractional Denoising for 3D Molecular Pre-training. ICML 2023.  
[4] Wang et al. Denoise pretraining on nonequilibrium molecules for accurate and transferable neural potentials. Journal of Chemical Theory and Computation 2023.

Reviewer comment (Novelty):
The significant results were, to a large extent, achieved through Equiformer. Against Equiformer's backdrop, the improvements contributed by DeNS appear minimal.

Author response:
We respectfully disagree. Equiformer is the current state-of-the-art in equivariant Transformers for atomistic modeling, but is computationally expensive. The reason we first apply DeNS to Equiformer is to see how well DeNS can improve training efficiency and performance, compared to trivially scaling up computationally expensive hyper-parameters like maximum degrees and training epochs as was done in their paper.

We show that DeNS improves sample efficiency and performance across datasets – MD17 (**General Response 3**), OC20 (**General Response 5**), OC22 (**General Response 6**) – and architectures – Equiformer (MD17), SEGNN (MD17), EquiformerV2 (OC20) and eSCN (OC20) – demonstrating its generality.

Reviewer comment (Novelty):
While encoding input forces can mitigate the non-equilibrium denoising's ill-posed problem, [1] shows that denoising without force input is also plausible. Does this undermine your motivation and imply the redundancy of input force encoding?

Author response:
We believe it doesn’t. The formulation in [1] is the same as our “no force encoding” ablation, which we show is worse than DeNS with force encoding on OC20 (Table 1(e)) and MD17 (**General Response 4**). Conceptually, without force encoding, the target of denoising non-equilibrum structures can be arbitrary and  ill-posed (see Figure 1). 

On some datasets where the non-equilibrium structures are closely related to the ones at equilibrium, denoising without force encoding can still work, especially if the gain of mapping from noisy distributions to data distributions outweighs the issue of the targets being ill-posed. In fact, we have a similar observation in Table 1(e), where training DeNS without force encoding (Index 2) is better than training without any denoising as in Table 1(d). Nevertheless, DeNS with force encoding is a safer and more general choice, enabling further improvement (Index 1 in Table 1(e)).  

Reference:  
[1] Wang et al. Denoise pretraining on nonequilibrium molecules for accurate and transferable neural potentials. Journal of Chemical Theory and Computation, 2023.

Reviewer comment (Experiments):
Could there be label leakage when encoding forces as input for energy and force predictions of structures?

Author response:
No, there is no label leakage for forces since we do not encode any force labels during evaluation (i.e., on the validation and testing sets). We only use force labels during denoising on the training set. In fact, on the OC20 testing set, the energy and force labels are not even publicly available. Evaluation on the testing sets is performed by a centralized server on EvalAI maintained by the OC20 authors.

To clarify the training procedure, we have both the upper blue block and the lower red block in Figure 2 during training but only have the upper blue block when evaluating the results on the validation and testing sets.

We clarify this in the revision.

Reviewer comment (Theory):
Given that force prediction is one of your experiments, should you consider adding the force prediction loss to eq. (6)?

Author response:
No, Equation 6 corresponds to the objective during DeNS (lower red block in Figure 2). Force prediction loss is already covered in Equation 1 (upper blue block in Figure 2). As stated in “Auxiliary Task” in Section 3.2.3., “Specifically, given a batch of structures, for each structure, we decide whether we optimize the objective of DeNS (Equation 6) or the objective of the original task (Equation 1).” When training DeNS (Equation 6), since we already encode forces, predicting forces would be trivial and thus we do not add the force prediction loss to Equation 6.

Reviewer comment (Reproducibility):
Provision of the pseudocode for DeNS training would be beneficial. This can offer insights into the usage of Multi-Scale Noise and other hyperparameters like p_{ DeNS }.

Author response:
Thanks for the suggestion! We add this to the appendix (Section E).