DenoiseVAE: Learning Molecule-Adaptive Noise Distributions for Denoising-based 3D Molecular Pre-training

Authors: Yurou Liu, Jiahao Chen, Rui Jiao, Jiangmeng Li, Wenbing Huang, Bing Su
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: 3D Molecular pre-training via denoising, Molecular property prediction

Abstract:
Denoising learning of 3D molecules learns molecular representations by imposing noises into the equilibrium conformation and predicting the added noises to recover the equilibrium conformation, which essentially captures the information of molecular force fields. Due to the specificity of Potential Energy Surfaces, the probabilities of physically reasonable noises for each atom in different molecules are different. However, existing methods apply the shared heuristic hand-crafted noise sampling strategy to all molecules, resulting in inaccurate force field learning. In this paper, we propose a novel 3D molecular pre-training method, namely DenoiseVAE, which employs a Noise Generator to acquire atom-specific noise distributions for different molecules. It utilizes the stochastic reparameterization technique to sample noisy conformations from the generated distributions, which are inputted into a Denoising Module for denoising. The Noise Generator and the Denoising Module are jointly learned in a manner conforming with the paradigm of Variational Auto Encoder. Consequently, the sampled noisy conformations can be more diverse, adaptive, and informative, and thus DenoiseVAE can learn representations that better reveal the molecular force fields. Extensive experiments show that DenoiseVAE outperforms the current state-of-the-art methods on various molecular property prediction tasks, demonstrating the effectiveness of it.

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730091575057

thread_type:
official_review

<thread_object>
creator: reviewer_r3gr

<content>
summary:
<Summary>

 - The authors propose a denoising-based 3D molecular pre-training method, called DenoiseVAE, which employs a learnable noise generation strategy instead of existing hand-crafted strategies. This allows for adaptive, atom-specific noise distributions across different molecules.

 - The authors introduce a variational approach to jointly optimize the Denoising Module and Noise Generator of DenoiseVAE. Here, the denoising objective encourages generated noise to adhere more closely to physical constraints, enabling the recovery of equilibrium conformations. Additionally, a KL divergence-based regularization term is applied to prevent low noise intensity and increase the diversity of sampled conformations.

 - Theoretically, the authors demonstrate that optimizing their pre-training objective is equivalent to maximizing the evidence lower bound (ELBO) of the log-likelihood.

 - Extensive experiments reveal that DenoiseVAE outperforms existing denoising methods across various datasets for both molecular and complex property predictions.

strengths:
<Strengths>

 - DenoiseVAE employs a learnable noise generation strategy that allows for adaptive, atom-specific noise distributions across different molecules.

 - The authors introduce a variational approach to jointly optimize the Denoising Module and Noise Generator of DenoiseVAE. Here, the denoising objective encourages generated noise to adhere more closely to physical constraints.

 - A KL divergence-based regularization term is applied to prevent low noise intensity and increase the diversity of sampled conformations.

 - Optimizing pre-training objective is proven to be equivalent to maximizing the evidence lower bound (ELBO) of the log-likelihood.

 - DenoiseVAE outperforms existing denoising methods across various datasets for both molecular and complex property predictions.

weaknesses:
<Limitations>

  - The Boltzmann distribution is commonly used in the classical regime. However, for precise coordinate computation of interacting atoms, a quantum approach is necessary. This involves, for example, treating the energy function as a quantized operator and finding energy eigenvalues for a given basis. However, this article does not carefully consider such approaches.

 - This approach essentially adopts a force field method, which may be inadequate for precise atomic-level computations.

 - Additionally, the energy is modeled as a simple harmonic potential, lacking any exchange correlation or other accurate potential forms.

questions:
<Question>

 - The prior distribution constrains the atom coordinate distribution to follow a predefined Gaussian form. But what is the range of these coordinate distributions? Are they confined within a narrow range close to the prior distribution, or are they spread across a finite range? If so, what is that range?

 - The article addresses 3D coordinate noise by adding Gaussian noise based on a learnable distribution, but what about rotations?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730654930909

thread_type:
official_review

<thread_object>
creator: reviewer_nlii

<content>
summary:
This paper presents DenoiseVAE for 3D molecular pre-training that adapts to the anisotropic nature of molecular force fields. The authors proposed Noise Generator, which learns molecule-specific noise distributions for each atom, allowing DenoiseVAE to generate conformations that align closely with a molecule's potential energy surface. The network architecture is a variational autoencoder, where noisy molecular conformations are denoised using a denoising module. The authors claim that this method leads to more accurate representations for downstream molecular tasks. Experiments on molecular property prediction tasks including QM9 and MD17 are conducted to demonstrate the effectiveness of the proposed method.

strengths:
1. The use of adaptive noise generation is novel. Previous works use hand-crafted or uniform noise across molecules, but DenoiseVAE uses an automated, atom-specific noise generation method with stochastic parameterization, which respects the unique structural and chemical properties of each molecule.

2. The paper proposes theoretical derivations based on PES which aligns with quantum chemistry principles. The derivation for Theorem 1 based on ELBO offers theoretical insights into the effectiveness of the proposed method.

3. Experimental results are promising and convincing: the proposed method achieved state-of-the-art performance across all benchmarks.

weaknesses:
1. The use of DenoiseVAE introduces additional computational complexity. The training of the VAE with molecule-specific noise sampling could be computationally expensive or impossible for large datasets or complicated molecules. A more thorough analysis of training time and resource requirements could be more helpful.

2. It is unclear to me how the proposed DenoiseVAE can be adapted across different dataset scales. It would be great if the authors could offer a series of studies on the performance of DenoiseVAE on different sizes of datasets.

3. I can imagine the training of DenoiseVAE can be very unstable and will be sensitive to initialization and optimization settings. It would be great if the authors could show convergence of loss curves across all datasets.

questions:
1. How can different formulations of noise distributions affect the result? Currently, the noises are Gaussians with diagonal covariance, what if the covariance is non-zero at off-diagonal positions? E.g., the neighboring atoms connected by a bond will have non-zero covariance.

2. How well will the model perform if the training and testing are on molecules of different sizes? Say the model is trained on small molecules but is tested for large molecular structures such as proteins.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730669234128

thread_type:
official_review

<thread_object>
creator: reviewer_ho1z

<content>
summary:
The This work proposes a 3D molecule pretraining method with learnable atom-specific noise generation, as opposed to the common practice of hand-crafted and shared noise. The pre-trained representation shows higher performance than baselines in most of the downstream energy prediction tasks. Furthermore, the authors show the learned noise intensities correspond with the atomic reactivity and rigidity.

strengths:
1. This work addresses an important problem in molecular representation learning. The learnable and molecule-specific per-atom noise distribution is a reasonable setup and aligns better with the physical intuitions.
2. The authors provide comprehensive theoretical analysis of the rationale of their method.
3. The proposed method shows consistent improvement over the baselines in the majority of the tasks, while also being more parameter efficient.
4. The method has good interpretability and the learned noise patterns are physically relevant.

weaknesses:
The results are overall solid. However, some additional experimental and implementation details about the downstream tasks should be provided to help better understand and assess the results. See Questions.

questions:
Major:
It’s a bit unclear how the downstream predictive tasks are performed after pretraining. Specifically:

(1) Is the force prediction achieved by a task-specific prediction head? If so, what is its architecture and how is it trained (e.g. end-to-end fine tuning)?

(2) What are the features used for the prediction (e.g. pooled EGNN outputs or latent embedding)?

Minor:
(1) Following the previous questions, as the model is considered a VAE, how is the latent space obtained?

(2) Is it possible to visualize the learned energy landscape (e.g. for MD17 samples)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698807449

thread_type:
official_review

<thread_object>
creator: reviewer_d6p5

<content>
summary:
This paper introduces NoiseVAE, a novel molecular pre-training method, for learning atom-specific noise distributions in different molecules. The variational autoencoder (VAE) based model consists of a Noise Generator and a Denoising Module that are jointly trained. The pretraining objective is to learn atom-specific noise distributions that align better with quantum mechanical theories than empirically derived strategies. The results show that the pretraining DenoiseVAE can generate noises specific to atomistic environments and improve the prediction accuracy on molecular properties such as the HOMO-LUMO gap and potential energy.

strengths:
### Originality
*My expertise in related literature is limited.*
1. The NoiseVAE model leverages a VAE for a more robust atom-specific and physics-related generation of noise distribution.

### Quality
1. The experiments/results are very comprehensive.
1. The shown results on the downstream tasks suggest the robustness of the generated noise in providing quantum chemical information for accurately predicting molecular properties.

### Clarity
1. The pretraining objective and method for NoiseVAE are well introduced with details.

### Significance
1. Exploring the conformational space efficiently is of significance in quantum chemical simulations. NoiseVAE provides a way to efficiently sample conformations around a given geometry.

weaknesses:
While I appreciate the amount of information and results presented in the paper, I find the experiments and results a bit difficult to follow.

- The authors showed proof that the proposed DenoiseVAE provides a higher theoretical evidence lower bound (ELBO) guarantee for the real conformation distribution of isoenergetic molecules.

- Subsequently, to show the effectiveness of the DenoiseVAE, the authors conducted downstream tasks with the pre-trained DenoiseVAE for predicting molecular properties.

- However, the authors did not provide a clear explanation of how DenoiseVAE was used in the downstream molecular property prediction tasks. To the best of my understanding, the DenoiseVAE was used to generate noise distributions for the input molecules. However, how were these noise distributions leveraged was not clear to me.

- The tables with results are also a bit misleading. For example, tables 1 and 2 mentioned "force prediction", but the 12 properties in Table 1 were not force-related and the values in Table 2 seemed more energy-related than force-related.

- In addition, the paper lacks details on the training process of the NoiseVAE model. Specifically, the pretraining is performed with the PCQM4Mv2 dataset, but the paper only briefly mentions that the dataset has 3.4 million organic molecules.

questions:
The authors are encouraged to provide more details concerning the above-mentioned weaknesses. For example,
1. What is the significance of the PCQM4Mv2 dataset in the pretraining of the NoiseVAE model?
   - How does the dataset help in learning atom-specific noise distributions?
1. How are the noise distributions generated by DenoiseVAE used in the downstream molecular property prediction tasks?
   - What is the architecture of the property predictor(s) used in the downstream tasks? What type of model is used?
   - How are the noise distributions used in the property prediction tasks, in addition to the input 3D molecular geometry?
1. What are the correct metrics/criteria for the results in Tables 1 and 2?

Additionally, the authors may consider addressing the following questions:
1. The paper provides proof that the DenoiseVAE provides a higher theoretical evidence lower bound (ELBO) guarantee for the real conformation distribution of isoenergetic molecules. Is there a more direct way to show the correctness of the noise distributions generated by DenoiseVAE with the ground truth by first-principles calculations?
   - For example, in Figure 4, the authors showed the generated noise distributions for two molecules. How close are these distributions when benchmarked against first-principles simulations?
1. How does the DenoiseVAE model help with conformational sampling in quantum chemical simulations?
   - The DenoiseVAE model generates noise distributions that are specific to a 3D molecular geometry (to the best of my understanding). By sampling from these noise distributions, are the sampled conformations more likely to be energetically favorable/stable than randomly sampling from a preset Gaussian distribution?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces DenoiseVAE, a denoising-based 3D molecular pre-training framework that learns molecule-adaptive, atom-specific noise distributions. Unlike prior methods relying on uniform or hand-crafted noise, the proposed approach employs a Noise Generator to produce customized noise distributions for each atom, while using a standard denoising objective to approximate molecular force fields. This leads to potentially more physically informed molecular representations.

All reviewers acknowledged the novelty and reasonableness of learning atom-specific noise distributions, and they found the empirical improvements over baseline methods to be well demonstrated. However, several reviewers raised concerns about the physical fidelity of the generated conformations. They suggested more direct comparisons with ground-truth force fields or DFT-calculated distributions to better validate the physical authenticity. While the authors did provide some responses and additional experiments, these concerns were not fully resolved during the rebuttal. Additionally, some reviewers recommended integrating known rotational, vibrational, or quantum constraints into the model or loss function. Although the authors partially addressed these points, the reservations were not entirely alleviated.

Overall, all reviewers reached a consensus that the paper presents a meaningful contribution to the denoising paradigm for molecular representation learning, despite reservations regarding its physical fidelity. The AC agrees with the reviewers and recommends acceptance.