Equivariant Masked Position Prediction for Efficient Molecular Representation

Thank you very much for your review. Your concerns and questions are very meaningful. In the revised version, we have added experiments on time cost (Appendix C.2, Table 12). You can also directly look at the response below to save time.

Responses to Questions 1:

We have established three tasks: 1. Only predicting the alpha property of QM9 using Equiformer, 2. Only predicting positions using Equiformer (Ours), 3. Only predicting noise using Equiformer (Denoising). They correspond to three models with different output heads. Each experiment was conducted on a single Nvidia A100, and the table below shows that the time consumptions of three models are similar (EMPP is slightly bigger). We argue there are two main reasons why the time consumption introduced by EMPP's output head is small: First, the output head of EMPP only considers the neighboring nodes of the masked atoms, and the rest of the nodes do not need computation. Second, before entering the grid module, the number of embedding channels has already been compressed to a small number 32 (as we mentioned in Table 5 and 6).

Additionally, you mentioned about the small improvement on MD17. We consider that is because: the training set of MD17 is a small (950 molecules) and easy-to-learn dataset. Thus, the MAE of Equiformer has reached saturation (less than 5 meV in most tasks), and the augmentation of EMPP for 950 data may be limited. The more complex QM9 dataset can better demonstrate the advantages of EMPP.

I apologize for any confusion. In your mention of weaknesses, you state, 'for non-pretraining comparisons, there are models that substantially outperform the proposed method, based also on higher-order tensors.' If you're referring to the pre-training comparison in Table 3, we used TorchMD-Net for the comparison with denoising methods (since they used it as the backbone). In fact, EMPP can be applied to any equivariant GNNs.

Responses to Questions 2:

We have published some of the sampling results in the Table 4. In our task, we want to represent fine-grained spherical distribution with dense sampling, and the grid operation does not actually take up a significant amount of time, so we use a larger sampling rate to ensure the results. In fact, when the sampling rate exceeds a certain value, the performance will stabilize: we tested the sampling rate of 80/90 during the rebuttal period, the MAE of $\\alpha$ and $\varepsilon_{HOMO}$ are 0.041/0.041 and 14.4/14.3 (the results of default setting 100 are 0.041 / 14.2).

Responses to Questions 3:

Good question! We conducted experiments to verify whether the improvement comes from our TorchMD variant. As shown in the table below, the modified TorchMD achieved very similar results to the original. This is mainly because that we only projected the node embeddings onto higher orders using spherical harmonics without changing the core operations of ET. You can also see from our Table 1 and Table 2 that EMPP introduced no changes on the Equiformer, yet it still achieved improvements. Additionally, we modified TorchMD because higher-order equivariance helps to describe a finer-grained spherical function (after grid module), which thus aids in the learning of EMPP.

(* denotes the results we that we have reproduced in the same environment.)

Finally, we greatly appreciate the pointing out of our citation error, which has been corrected. All of your questions are of significant importance in refining our paper.

Thank you very much for your review. We will address all your concerns, and to save your time, we will bold some key points.

Response to main concerns:

(1) The Gaussian distribution in Equation 14 is used to transform scalar distances $|\vec{r}\_{ij}|$ into a set of continuous and smooth features. This transformation enhances the model's expressive power and makes it easier to learn complex geometric relationships. This approach is known as Gaussian RBF, which is commonly used in equivariant GNNs, such as TorchMD-Net [1], TFN [2], SCN [3]. The directional distribution in Equation 16 serves a similar purpose. It is important to emphasize that in EMPP, the true labels are the well-defined relative positions $\vec{r}\_{ij}$. As long as the defined distribution can uniquely represent $\vec{r}\_{ij}$, the model can learn the correct information. Moreover, we have included additional experiments in Appendix C.2 (Table 10) to show that EMPP with a Dirac delta distribution also work well. However, the smooth distributions we employed can enhance training stability and yield superior results. Further details are provided in the subsequent responses. Additionally, we have revised the loss section to clarify potential misunderstandings. Thank you for your suggestion.

[1] TorchMD: A deep learning framework for molecular simulations. Stefan Doerr, et al.

[2] Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. Nathaniel Thomas, et al.

[3] Spherical Channels for Modeling Atomic Interactions. C. Lawrence Zitnick, et al.

(2).Our method is fundamentally different from denoising approaches. Denoising methods approximate an unknown potential energy surface (PES) by using Gaussian mixtures and compute its derivative (around noisy points) to create a pseudo-label for the force prediction. In contrast, EMPP does not approximate unknown physical variables. It bypasses the PES and its derivatives, instead directly modeling the relationship between atomic interactions and positions. Since the force field is a function of position that reflects overall interactions, position prediction in EMPP can indirectly capture the key features of the force field.

Responses to Weakness 1:

We have submitted a revised paper to address these concerns, which you can check. Additionally, we would like to make some explains:

(1) The role of the language model is to highlight the importance of self-supervision and data augmentation. We have modified this part to enable the readers to understand clearly.

(2) The ablation study in Section 5.4 includes two experiments (Table 4 and Figure 3). Table 4 is for the proposed method and Figure 3 reflects the fundamental shortcomings of denoising which is very important for understand our motivation. Due to space limitations, we move all other less critical ablation to the appendix.

(3) We clarify that EMPP is not to predict forces. The purpose of EMPP is to enable models to predict a physically plausible position using a set of neighbors, which essentially requires the model to implicitly model the correct atomic interactions (such as force field). Besides, there exists a difference: denoising methods approximate the PES and computing derivatives to learn forces, while EMPP models the functions of atomic interactions and atomic absolute position.

(4) During fine-tuning, EMPP serves as an auxiliary task, meaning each training molecule passes through the network twice: first to calculate the regular energy or forces loss, and second to calculate the EMPP loss, with both contributing to gradient descent. The denoising methods are also as a similar auxiliary task during fine-tuning [4,5]. The difference is that EMPP can model the relationship between labels and molecular structures (positions), thus promoting generalization of backbone to certain labels, whereas denoising as an auxiliary task does not consider labels. Our experiments (Table 1, Table 2) also demonstrate that even without pre-training (introducing additional data) EMPP can enhance model generalization, whereas denoising methods lose effectiveness without pre-training, as shown in the table below.

(* denotes the results we that we have reproduced in the same environment.)

[4] Pre-training via denoising for molecular property prediction. Zaidi S, et al.

[5] Fractional denoising for 3d molecular pre-training. Feng S, et al.

Responses to Weakness 2:

The Gaussian distribution in our mehtod is different from it in denoising. Denoising methods use it to approximate the local minima of unknown PES. We use it to represent the well-defined relative positions $|\vec{r}\_{ij}|$, like the Gaussian RBF in previous molecular models (See Responses to main concerns). The spherical representation in Equation 16 is similar, They brings a significant benefit: facilitating learning of neural networks. For example, The computer represents distributions using sampling on the radius or sphere, and the peak of the strict Dirac distribution might be missed by finite sampling. When we use smooth representation like Gaussian, computer sampling can better fit mathematical distributions, making the loss calculation accurate. The experiments in Table 10 show that training with smooth functions is more stable and easier to achieve better results.

Responses to Weakness 3:

The explanation of "force prediction" can be found in responses (3) to Weakness 1. Besides, EMPP is fundamentally different from denoising methods (See Responses to main concerns).

The implementations of EMPP and denoising are also significantly different: denoising methods use noise sampling to predict noise, whereas EMPP employs position masking and predicts positions based on unmasked atoms. In denoising, equivariant GNNs are primarily used to enforce the equivariance of the predicted noise. In contrast, we use equivariant models for an additional purpose: equivariant models are capable of describing complex interactions (Theorem 2 in [6]), particularly when the degree of interaction is high, which enables EMPP to accurately model the relationship between interactions and positions.

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

As for your point on diffusion, it is similar to denoising methods, but not the same. The similarity arises because predicting the forces from the Boltzmann distribution is precisely a score matching process [4]. However, the purpose of EDM (which you referenced as [4] in reviews) is to add artificially defined noise to the input and then recover the molecular structure from the noise, with all distributions in diffusion being artificially defined and known. The goal of denoising is not to generate molecules, but to produce plausible data to enhance generalization. However, the real physical spaces (such as PES) are unknown and denoising methods use Gaussian to approximate. Therefore, the challenge lies not in solving score matching, but in how to approximate the real physical distribution. Similarly, the citation [2] in your reviews predicts the $\sigma$ of Gaussian; however, the label of $\sigma$ used during training is still an approximation, not the curvity of true physical distribution.

Reponses to Weakness 4:

Please refer to Responses to main concerns and Weakness 2 to understand the role of Gaussian distribution we use. Additionally, you provided a great example. Please note that the purpose of the predicted distribution is to describe the ground truth $\vec{r}\_{ij}$, not generate atoms or fit the unknown distribution. Thus, we do not need to align the distribution of each neighbor; as long as it can represent $\vec{r}\_{ij}$, that is enough. Equation 14 and 16 are both unique mappings of $\vec{r}\_{ij}$.

Responses to Question 1:

Please refer to our responses (4) to Weakness 1.

Thank you. We have provided a clearer explanation of fine-tuning in the revised version, indicated in red, please review.

Responses to Question 2:

Please refer to our responses (4) to Weakness 2.

Thank you. We have provided a clearer explanation of the Gaussian assumption in the revised version, indicated in red, please review.

Responses to Question 3:

Please refer to our responses to Weakness 3.

Responses to Question 4:

Please refer to our responses to Weakness 4. As we mentioned, we do not need to align the distributions of neighbors. These distributions are not intended to generate or fit real distributions. They are projections of $\vec{r}\_{ij}$, which is real in molecule and consistent for all neighbors.

Responses to Question 5:

Thank you for your question. We have added experiments regarding $\tau$ in the Appendix C.2 (Table 13). Briefly, EMPP is not sensitive to $\tau / \sigma$. we generally use the common setting of $0.1/0.5$, and our additional experiments also show that the common setting is suitable for different tasks, and for different values, results are similar. As previously mentioned, the role of these hyper-parameters are not same to $\sigma$ in denoising (which is used to approximate unknown curvity). Thus, They can be flexibly determined.

Thank you for your review. Our revision will not affect any core ideas and implementation, but help make the paper clearer. If you have any other concerns, please do not hesitate to let us know.

Thank you for your review and recognition of our paper. We will address your concerns.

Responses to Weekness:

In Table 1 and 2, we believe that auxiliary tasks without pre-training provide more accurate reflection of the method's performance. However, while previous denoising methods have been used as auxiliary tasks during fine-tuning, their reported results all rely on pre-training and torchMD-Net (we use Equiformer in Table 1 and 2). To provide you a more fair comparison, We conducted experiments without pre-training based on TorchMD-Net (using the open-source implementation of the denoising method [1]). The results are shown in below. These experiments revealed that the performance improvement of denoising methods is limited in the absence of pre-training. In contrast, EMPP demonstrated a stronger improvement compared to denoising methods.

(* denotes the results we that we have reproduced in the same environment.)

[1] Pre-training via denoising for molecular property prediction. Zaidi S, et al.

Responses to Question:

Yes, masking multiple atoms leads to increased training time. In our response to Questions 1 of reviewer DiCi, we evaluate the time cost of EMPP. In short, when using n-mask EMPP as auxiliary task, the training time will be expand by n times (as shown in the below tabel). We are also glad to discuss the trade-off: in most cases, the 1-Mask method is sufficient, as it doubles the training time but also significantly improves performance. More masks can further enhance performance, but currently, their improvements are not significant, and the training time is expanded several times. We consider that when resources and time are abundant, multi-mask methods can be considered. Besides, EMPP does not affect the efficiency of model inference.

Our ablation study and discussion on trade-offs between performance and computational costs has been added in Appendix C.2 (Table 12).

Thank you for your review and recognition of our paper. We will address your concerns.

Responses to Weakness "Unclear Evaluation in Pre-training":

We acknowledge your guess that the poor ddiversity of QM9 limit methods. Additionally, there is an another reason: Table 1 and Table 3 utilize different backbones (Equiformer and TorchMD-Net). These models inherently introduce performance differences. The Equiformer in Table 1 is a more advanced model and performs better. In Table 3, we chose TorchMD because previous Denoising methods were based on TorchMD. To facilitate a fair comparison, we used TorchMD-Net.

Responses to Suggestion and Weakness "Too simple benchmark":

Great suggestion! we will conduct the experiments on "Geom-Drugs". If it is completed by the rebuttal deadline, we will add it in the paper. Otherwise, we will make it publicly available on GitHub (Sorry, we need time and resource to train). Now, we can first discuss other concerns.

Note that EMPP is a self-supervised method, and its purpose is to produce reasonable data within a given molecular system. From Table 1, we observe: the QM9 molecular system is relatively simple and lacks of diversity, and the performance (MAE) of Equiformer is approaching saturation, yet EMPP can still produce a significant performance improvement. We believe this is sufficient to demonstrate the effort of EMPP. If we increase the complexity and diversity of the molecular system as you mentioned, EMPP should intuitively provide even greater improvements, since it can produce more kind of reasonable data.

Responses to Questions "Physical Assumptions in Modeled Distributions":

We view our approach as a distinct modeling assumption, different from Denoising.

Denoising methods use Gaussian mixtures to approximate an unknown (Boltzmann) distribution. They then perform noise sampling on this distribution to learn the force field (the derivative of the PES). EMPP does not involve any approximation to unknown physics distribution. It models directly at the level of the force field without computing derivatives. The purpose of EMPP is to enable models to predict a physically plausible position using a set of neighbors, which essentially requires the model to implicitly model the correct atomic interactions (such as force field).

Responses to Questions "Comparison with Denoising Methods":

Note that the distribution we used is not to fit an unknown distribution, but to transform scalar distances $|\vec{r}\_{ij}|$ into a set of continuous and smooth features. This transformation known as Gaussian RBF can enhance the model's expressive power and makes it easier to learn complex geometric relationships. We have conducted corresponding ablation studies and added the results to Appendix C.2 (Table 10). In short, we used various functions to represent relative positions, all of which achieved similar improvements. However, when we directly predict the relative positions, the improvement of EMPP was small. We believe this is because the transformation $\vec{r}\_{ij}$ to a smooth function facilitates learning the information of relative positions. Moreover, this is also a common practice, in many molecular GNNs [1,2], models do not directly map vectors to length and direction, but use a smooth Gaussian RBF to map their length before feedding to the model. More details on why we use such distributions can be referred to our responses to weekness 2 of reviewer RrYB.

[1] TorchMD: A deep learning framework for molecular simulations. Stefan Doerr, et al.

[2] Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. Nathaniel Thomas, et al.

We conducted experiments on GEOM-Drug during the rebuttal period, and the introduction and analysis of this part have been incorporated into the revised version (appendix C.2). We use Equiformer as the backbone model. The training configuration follows the configuration on QM9 experiments, with the number of training epochs set to 150. We use the absolute energy of each conformation as the label. Additionally, we randomly sample 200,000 molecules from GEOM-Drug as the training set and 10,000 as the validation set (Apologies for the limited resources and time; we are unable to use the full dataset. However, 200,000 data points should still be sufficient to demonstrate the advantages of EMPP.). To ensure the validity of the validation results, SMILES that appear in the validation set do not appear in the training set (a single SMILES includes multiple conformational data). The below table show that EMPP can achieve significant performance improvements on GEOM-Drug. The energy MAE can be decreased by $32$%. We draw two conclusions: 1. EMPP is applicable to more complex organic molecular systems; 2. Most molecules in GEOM-Drug are non-equilibrium. EMPP can still achieve an significant improvement when used as an auxiliary task. This experiment shows that EMPP overcomes the limitation of denoising methods that can only approximate local minima (equilibrium structure)."

[Updated in 11.25, We trained the Drug dataset for 300 epochs to achieve a stable performance. The results show that EMPP can produce a more significant improvement (48%).]

To evaluate the impact of EMPP on molecules with different sizes, we conducted experiments by categorizing the QM9 training data into three groups based on the number of atoms: (0-16), (17-19), and (20+). These categories contain roughly equal amounts of data. In each experiment, we computed the EMPP loss using only molecules from one of these categories. As shown in the below table, the results demonstrate that EMPP consistently improves performance across different molecular sizes, with larger molecules experiencing more significant performance gains.

Thank you for your detailed and thorough responses to the questions and suggestions in my review. I greatly appreciate the effort you have put into addressing the raised concerns and providing additional experimental results. Your explanations clarify key aspects of EMPP's design and its theoretical underpinnings, as well as its practical strengths. The new experiments on GEOM-Drug and the analysis of performance scaling with molecular size are especially helpful in demonstrating EMPP's applicability to more complex molecular systems. Thank you once again for your dedication and comprehensive responses. I look forward to seeing the final version of your paper and its contributions to the field. Best regards.

Hi Ameya,

Thank you for your suggestions. We've included a comparison in the camera-ready version. A general comparison can be found at the end of Section 3.2.2 (POSITION PREDICTION) and in Section 4 (RELATED WORK). Due to space limitations, more detailed comparisons are provided in Appendix B.5. In this version, only the comparison with Symphony has been added, while the rest remains unchanged.

Best regards,

Authors