Pretraining Strategy for Neural Potentials (2402.15921v2)

Abstract: We propose a mask pretraining method for Graph Neural Networks (GNNs) to improve their performance on fitting potential energy surfaces, particularly in water systems. GNNs are pretrained by recovering spatial information related to masked-out atoms from molecules, then transferred and finetuned on atomic forcefields. Through such pretraining, GNNs learn meaningful prior about structural and underlying physical information of molecule systems that are useful for downstream tasks. From comprehensive experiments and ablation studies, we show that the proposed method improves the accuracy and convergence speed compared to GNNs trained from scratch or using other pretraining techniques such as denoising. On the other hand, our pretraining method is suitable for both energy-centric and force-centric GNNs. This approach showcases its potential to enhance the performance and data efficiency of GNNs in fitting molecular force fields.

• The paper presents a novel mask pretraining strategy that enhances GNN performance by learning robust molecular representations.
• The method achieves significant improvements in the accuracy of force and energy predictions while accelerating model convergence.
• Results demonstrate the approach's broad applicability across GNN architectures, offering promising advances for computational chemistry and materials science.

Improving Graph Neural Networks for Molecular Dynamics through Mask Pretraining

Introduction to Molecular Dynamics and the Role of GNNs

In the field of computational chemistry and materials science, accurately modeling molecular dynamics is essential for a variety of applications, from drug discovery to the design of new materials. Central to these studies is the potential energy surface (PES), which describes the energy of a system with respect to the positions of all atoms involved. Traditionally, methods such as Density Functional Theory (DFT) have been employed to paper PES with high accuracy, albeit at a significant computational cost, especially for large systems. Alternatively, empirical interatomic potentials offer faster computations but lack the accuracy and adaptability seen in ab initio methods.

With advancements in ML, new methods have emerged to approximate PES with near ab initio accuracy at a fraction of the computational cost. Notably, the incorporation of Graph Neural Networks (GNNs), which directly learn from atomic coordinates, represents a significant advance in the modeling of molecular dynamics. Particularly, the development of neural network-based potentials using GNNs has shown promise in accurately representing PES while maintaining computational efficiency.

The Challenge with GNN-based Neural Potentials

Despite their potential, the performance of GNN-based neural potentials heavily depends on the availability of large, accurately labeled datasets, which are expensive and challenging to generate for complex molecular systems. To address this bottleneck, various pretraining strategies have been investigated, aiming to improve the data efficiency and performance of GNNs in molecular modeling tasks. These strategies involve pretraining GNNs on related tasks with ample available data before finetuning them on the target molecular dynamics prediction tasks.

Our Contribution: Mask Pretraining for Enhancing GNN Performance

We propose a novel pretraining strategy specifically designed for improving the performance of GNN-based neural potentials. Our approach involves pretraining GNNs by masking out spatial information of certain atoms—in our case, a hydrogen atom in water molecules—and then predicting the masked information based on the context provided by the unmasked parts of the molecule. This pretraining task forces GNNs to learn meaningful representations of molecular structures and the underlying physical interactions without the need for explicit labeling, thus enhancing the model's ability to predict potential energy and forces with improved accuracy and speed upon subsequent finetuning.

Experimental Validation

Our comprehensive experiments showcase the effectiveness of the proposed pretraining method. When compared to GNNs trained from scratch or those pretrained using other methods such as denoising, our approach consistently delivers superior performance across different metrics. Notably, we observe significant improvements in both the accuracy of force and energy predictions and the speed of model convergence. These results are robust across different GNN architectures, including those focused on force-centric and energy-centric predictions, highlighting the general applicability of our pretraining strategy.

Future Directions and Implications

The findings from our work open up new avenues for further exploration and application of GNNs in the paper of molecular dynamics. By demonstrating the utility of mask pretraining, we not only provide a novel method for enhancing the performance of neural potentials but also underscore the potential of pretraining strategies in overcoming data limitations inherent in modeling complex molecular systems. Future work may extend our approach to larger and more varied datasets or explore the integration of our pretraining strategy with other GNN architectures and learning paradigms.

In conclusion, our paper introduces a promising direction for the use of GNNs in computational chemistry and materials science, potentially accelerating advancements in drug discovery, materials design, and beyond. The ability to accurately model molecular dynamics with improved data efficiency has far-reaching implications, opening the door to new discoveries and innovations across sciences.