# Pretraining Strategy for Neural Potentials

Zehua Zhang  [  Zijie Li  [  Amir Barati Farimani  [<EMAIL_ADDRESS>
###### 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.

INI]Information Network Institute, Carnegie Mellon University, Pittsburgh PA,
USA MechE]Department of Mechanical Engineering, Carnegie Mellon University,
Pittsburgh PA, USA MechE]Department of Mechanical Engineering, Carnegie Mellon
University, Pittsburgh PA, USA [ML]Machine Learning Department, Carnegie
Mellon University, Pittsburgh PA, USA [Chem]Department of Chemical
Engineering, Carnegie Mellon University, Pittsburgh PA, USA

##  1 Introduction

Molecular dynamics plays a pivotal role in elucidating the dynamic behavior of
molecules, providing essential insights into the temporal evolution of complex
systems, such as understanding the conformational changes, interactions, and
thermodynamic properties of molecules. It is thus important for fields ranging
from drug discovery to materials science. The potential energy surface serves
as the underlying landscape dictating the dynamics of a molecular system, and
making an accurate representation of it is crucial for studying the dynamics
of molecules through numerical simulations. Ab initio methods like Density
Functional Theory (DFT) provide high accuracy by accounting for the electronic
structure of atoms. Yet their high computational demands pose a significant
challenge for efficient utilization, particularly in the context of large
many-body systems. On the other hand, forces can be computed using empirical
interatomic potentials tailored to specific environments, bypassing the
electronic structures of the system. This approach significantly reduces the
computational cost compared to ab initio methods. As there is a wide variety
of interactions in molecular systems (bonded or non-bonded interactions),
finding the appropriate functional forms can be challenging 1, 2. The
difficulty in describing complex interatomic interaction presents a potential
limitation for classical MD. Machine learning (ML) has emerged as a promising
alternative for learning and fitting the potential energy surface with
accuracy close to ab initio approaches 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
14. The kernel-based regression model is one of the examples. Gradient-domain
ML (GDML) employs a kernel-based method to fit the forcefield with the
guarantee of energy conservation 15. Following GDML, Symmetric GDML (sGDML)
extends GDML to incorporate space group symmetries and dynamic nonrigid
symmetries 16. Gaussian approximation potentials (GAPs) is another notable ML
method, which decomposes the energy as the sum of atomic-centered Gaussian
basis function 17, 18 with hand-designed local environment descriptors 19, 20.
The strong fitting capacity of the neural network has made it a suitable
choice for fitting the potential energy surface. The initial efforts of
building neural-network-based potential (neural potential) have leveraged
tailored descriptors that capture the characteristics of local environment of
atoms in the system 21, 7, 22, 23, 24, 25, where Behler-Parrinello Neural
Networks (BPNN) 26, 27, 28 is one of the first examples. The recent
advancements of deep neural networks, especially graph neural networks (GNNs),
open up the possibility of learning the atomic representations directly from
the raw atomic coordinates, which provides an alternative to the manually
tailored atomic fingerprints 29, 30, 31, 32, 33, 34, 35. In GNNs, the
interaction between atoms is modeled by message passing 36 between atoms
within the local neighborhood. SchNet 33 introduces continuous convolution for
aggregating messages across nodes (atoms) to produce a continuous energy
surface. Various other works have developed different message-passing
protocols based on physics and chemical insights 31, 32, 34, 35, 37. One
notable variant is the class of equivariant GNNs. With tailored message-
passing functions and update strategies, these neural networks are equivariant
to group actions such as rotation and translation in 3D Euclidean groups. The
exploitation of property of the symmetry has made equivariant GNNs not only
more data-efficient and also more accurate 38, 39, 40, 41, 42, 43, 44.

While GNN-based neural potentials have provided a flexible and accurate end-
to-end framework for learning and predicting the potential energy, their
performance heavily relies on the amount of data. Collecting DFT-based data
can be expensive or even infeasible for large and complex molecular systems.
To improve the data efficiency of GNNs on molecular modeling tasks such as
property prediction, various pretraining strategies have been investigated.
One direction is pretraining GNNs to predict the masked attributes or context
of the graph 45, 46, 47. Another direction is pretraining GNNs by contrasting
the embeddings from the network given different views of the data 48, 49, 50,
51, 52, 53, 54, 47. Different views can be generated via randomly masking out
some of the nodes and edges in the molecular graphs 48 or directly contrasting
3D representation with 2D representation of the molecular system 55, 56. It is
worth noting that many of these methods have focused on improving the GNNs’
performance on 2D graphs, while for neural potential the input is usually 3D
structures (coordinates of atoms) of molecular systems. Several recent works
have proposed to pretrain the GNNs on 3D representation by denoising 57, 58,
59, 60. More specifically, the pretraining task is to predict the atom-wise
Gaussian noise added to the atom coordinates under equilibrium state or non-
equilibrium state. The denoising pretraining has been demonstrated to help
improve models’ performance in predicting the properties of small molecules or
single molecule systems. Previous works have viewed denoising as learning the
pseudo-gradient (or force field) of the energy surface and an appropriately
small perturbation (e.g. Gaussian noise with variance less than 0.1) can yield
meaningful learning tasks for GNNs. However, identifying the appropriate noise
scale can be difficult, such that it should not be trivially small nor
excessively large. For the water system, the overall potential energy includes
the contribution from intra-molecule interaction and inter-molecule
interaction. This makes selecting an appropriate variance much trickier than
single-molecule systems as small noise cannot yield significant enough
perturbation to the intra-molecule part while large perturbation is too
difficult for GNNs to recover the original geometric structures.

Figure 1: Framework of pretraining GNNs over masked water molecules. (a)
Selecting water molecules with a certain ratio, spatial information of
selected Hydrogen atoms (Blue) is masked to create a pretext task. (b) GNNs
are pretrained to recover displacements between masked-out atoms and other
atoms. Information on any atoms but masked-out ones is provided for GNNs to
predict such displacements. (c) The pretrained GNNs are transferred and
finetuned to predict the potential energy surface.

##  2 Methodology

###  2.1 Graph Neural Networks

Molecules are naturally suitable to be modeled as graphs, where atoms are
depicted as nodes and their interactions as edges. Message Passing Neural
Network (MPNN)36, has been a popular Graph Neural Network paradigm for its
competitive performance on complexly structured graphs, which makes it apt for
molecular property prediction. It often includes two phases during forward
propagation: message passing and readout. We define the constructed graph $G$
that contains node features $x_{v}$ and edge features $e_{vw}$ encoded from
atoms and bonds (interaction) between atoms. During the message passing stage,
hidden states $h_{v}^{t}$ at each node are updated with an update function
$f_{u}^{t}$ based on messages $m_{v}^{t+1}$ that are aggregated by a message
function $f_{m}^{t}$.

$m_{v}^{(t+1)}=\sum_{j\in\mathcal{N}(i)}f_{m}^{(t)}(h_{v}^{(t)},h_{w}^{(t)},e_{vw})$
(1)

$h_{v}^{(t+1)}=f_{u}^{(t)}\left(h_{v}^{(t)},m_{v}^{(t+1)}\right)$ (2)

Then a readout function $R$ aggregates all node features to learn a graph-
level representation.

$\hat{y}=R(\\{h_{v}^{T}|v\in G\\})$ (3)

Several research 61, 62, 45, focused on extending MPNNs to the field of
molecular attribute predictions. Among them, a subset of equivariant and
covariant MPNNs, EGNN63, GNS64, and ForceNet65 are used in our investigation
as baseline models. EGNN extends equivariance to higher orders and shows
competitive performance on molecular datasets such as QM9 66. When used to
predict molecular dynamics, EGNN first predicts energy while maintaining
rotational equivariant and energy conserving. The forces are then derived as
the negative first derivatives of energy with respect to atomic positions. On
the other hand, both GNS and ForceNet are based on encode-process-decode
architecture. Generally, a GNN is used to encode and update node embeddings,
which are then decoded and used to predict atomic forces. A Multilayer
Perceptron (MLP) finally approximates the energy by aggregating decoded
embeddings.

###  2.2 Pretraining by masking

Our proposed approach Hydrogen atom masking, similar to token masking in
Masked Language Modeling 67, refers to deliberately masking one Hydrogen atom
out of each selected water molecule with its spatial information removed from
the input. The selection ratio is a hyperparameter that directly affects our
framework’s performance. Analysis of masking ratios is further discussed in
the supplementary information section 6.1.

Specifically, before masking, we calculate displacements $\mathbf{d}_{i}$
between atoms within the same molecule and append them to inputs along with
indices of masked atoms, masked displacements and positions, atom types, and
box sizes. Models will then predict the spatial information of the masked
atoms and return their predicted displacements $\hat{\mathbf{d}}_{i}$ relative
to other atoms in the same water molecules. The to-be minimized objective
$\mathcal{L_{\text{masking}}}$ is the negative Cosine Similarity between the
prediction $\hat{\mathbf{d}}_{i}$ and ground true label, $\mathbf{d}_{i}$:

$\hat{d_{i}}=\phi_{\theta}(V,\hat{X})$ (4)
$\mathcal{L_{\text{masking}}}=1-\frac{\sum_{i=1}^{n}\hat{d_{i}}\cdot
d_{i}}{\sqrt{\sum_{i=1}^{n}\hat{d_{i}}^{2}}\cdot\sqrt{\sum_{i=1}^{n}d_{i}^{2}}},$
(5)

where $\phi_{\theta}$ denotes an GNN parametrized by $\theta$. $V$ and $X$ are
the respective encoded atom information and Cartesian coordinates of input
data samples.

The atom type we choose to mask out is based on the chemical properties of
atoms in the water molecules. We empirically discovered that masking Oxygen
atoms out causes instability during training and provides degraded results,
which could partially be attributed to that an Oxygen atom weighs 18 times as
a Hydrogen atom and exerts greater force on its neighbors. Consequently, we
decided to mask out a hydrogen atom from selected water molecules and use it
as the appropriate pretraining target.

On the other hand, the remaining Oxygen and Hydrogen atoms would have
different impacts on the predictions of masked Hydrogen atoms’ coordinates due
to their different atomic types. Thus, by reconstructing spatial information
of missing atoms, GNNs are expected to learn from restoring the structural and
physical information and the inherent structure of water systems. Moreover,
the choice of using Cosine Similarity instead of $L1/L2$ distance to calculate
the loss of displacement predictions is due to its scale invariance. It is
more reasonable to have the GNNs learn the relative positioning of the atoms
instead of the absolute values of the displacements, which may be of small
variance under different MD simulation conditions such as changes in
Temperature in the data generation process.

###  2.3 Pretraining by denoising

Denoising refers to predicting the artificial noise added to input atom
coordinates in a self-supervised fashion. The noise $\xi\in\mathbb{R}^{N\times
3}$ is usually sampled from Gaussian $\mathcal{N}(0,\sigma I)$. Such noise is
added to atom coordinates as perturbations. Taking noised atom coordinates,
atom types, and box sizes as input, GNNs return the predicted noise
$\hat{\mathbf{\xi}}$ added to the coordinates. The denoising objective is MSE
Loss, aiming to capture and minimize the difference between predicted noise
and true noise.

The denoising training objective is shown below.

$\hat{\xi_{i}}=\phi_{\theta}(V,\hat{X})$ (6)
$\mathcal{L_{\text{denoising}}}=\|\hat{\mathbf{\xi}}-\xi\|^{2}$ (7)

where $\phi_{\theta}$ denotes an GNN parametrized by $\theta$. $V$ and $X$ are
the encoded atom information and noised Cartesian coordinates of input data
samples.

The capacity to predict perturbation exerted over atom coordinates helps GNNs
develop an understanding of reasonable positions and distances relative to
each other. As analyzed in several prior works, denoising atomic coordinates
is related to learning a pseudo-force field at the perturbed states 57, 68,
60, 69. For a given molecule with $N$ atoms and its corresponding coordinates
$X$, the probability of this conformation ${X}$ is $p(X)\sim\exp{(-E(X))}$
following the Boltzmann distribution, with $E(X)$ being the potential energy
at $X$. As shown in Vincent 70, the network $s_{\theta}$ the minimizes the
following objective:

$\frac{1}{2}\mathbb{E}_{q_{\sigma}(\tilde{X}|X)p(X)}[||s_{\theta}(\tilde{X})-\nabla_{\tilde{X}}(\log
q_{\sigma}(\tilde{X}|X)||_{2}^{2}]$ (8)

also satisfies $s_{\theta}(X)=\nabla_{X}\log q_{\sigma}(X)$, with
$\tilde{X}=X+\xi$ and $\xi\sim\mathcal{N}(0,\sigma I)$. When $\sigma$ is small
enough, $s_{\theta}$ serves as a good estimate of the score function of the
true data distribution $\nabla_{X}\log p(X)$. Hence, the denoising pretraining
using Gaussian noise with small variance is equivalent to estimating the score
function or the negative gradient of the energy. However, this requires the
variance to be small enough and we will show in the experiment that it is not
a necessary condition to achieve good performance in practice.

Given denoising is a popular method to pretrain GNNs when approximating
Molecular Dynamics, we also pretrained baseline models using denoising to
offer a comparison, as indicated in section 3.4.

##  3 Experiment

###  3.1 Experiment Set up

The whole training pipeline is of two stages: pretraining and finetuning. In
the pretraining stage, all models are firstly trained on a pretext task in a
self-supervised way. Pretrained model weights are then transferred and
finetuned with supervision during the finetuning stage. This transfer aims to
leverage the knowledge learned about molecular systems during pretraining.
Such knowledge embedded in the GNN parameters is expected to improve the
performance of the GNN predicting force and energy labels with a strong prior.
Specifically, we pretrain the baseline models with the aforementioned mask
pretraining strategy in section 2.2. The overall architecture is illustrated
in Figure1. This paradigm is model-agnostic and can easily be applied to
different models.

The datasets used in this work are generated from classical molecular dynamics
(MD) and density functional theory (DFT). The classical MD simulations are
performed using OpenMM 71 with Tip3p forcefield 72. The systems considered are
uniform single-atom or single-molecule systems with different simulation box
sizes. Periodic Boundary Conditions (PBCs) are set in all directions, and a
cut-off distance of 10 angstroms was used. NVT ensemble with a Nosé–Hoover
chain thermostat 73, 74 is used for the simulation. We simulate each
configuration of the molecular system to $50000$ steps with a time step size
of $2.0$ femtoseconds and save the state of the system every 50 steps. Each
configuration is generated by initializing particles in the system with random
positions. A total of $10$ configurations are for generating the data. The DFT
simulation data of water molecules are obtained from Cheng et al. 75, which
uses the revised Perdew–Burke–Ernzerhof (RPBE) functional 76 with Van der
Waals correction 77.

Two datasets, RPBE and Tip3p, are both split into a train-validation set and a
test set with a 90% and 10% ratio, and the train-validation set is further
split into a train set and a validation set following the same proportion. In
our works, we reproduced three previously mentioned Graph Neural Networks,
EGNN63, GNS64, and ForceNet65 as baseline models to evaluate our method’s
performance.

During Pretraining, baseline models are trained on RPBE and Tip3p datasets for
100 and 25 epochs, respectively. The pretrained models are later finetuned
over the same dataset. The finetuning training epochs are 300 epochs and 50
epochs respectively. The difference between pretraining and finetuning epochs
mainly comes from the size of the dataset. RPBE has 7241 data samples.
Meanwhile, in the Tip3p dataset, for a total of 10000 data points, each of
them comes with 258 water molecules. Thus, GNNs should be trained for more
epochs to make up for fewer data samples in the RPBE dataset. To better
investigate the effectiveness of our masking pretraining strategy, we also
trained GNNs from scratch without pretraining, and we kept the rest of the
configurations the same, to compare with their pretrained counterparts.
Training from scratch epochs are 300 epochs and 50 epochs respectively as a
fair comparison.

We employ the AdamW optimizer 78 for both pretraining and finetuning. Learning
rates during all stages of training are set to 1e-4 and eventually reduced to
1e-7, with a weight decay rate of 1e-4. To create an equal comparison, we
apply the same batch sizes, 16 and 10, for RPBE and Tip3p datasets across
training stages and models. Notably, we observed that EGNN is highly sensitive
to larger batch sizes, as training would become fairly unstable. The cutoff
radius for water systems is set to 3.4 angstroms.

###  3.2 Experiment Results

Datasets | Models | Pretrained | Force RMSE | Energy RMSE | Force Improv | Energy Improv | Masking Ratio
---|---|---|---|---|---|---|---
RPBE | EGNN |  | 220.412 | 1597.309 |  |  |
EGNN | ✓ | 89.798 | 906.139 | -58.87% | -42.19% | 0.25
GNS |  | 214.970 | 1842.212 |  |  |
GNS | ✓ | 103.403 | 2217.729 | -51.80% | 20.31% | 0.5
ForceNet |  | - | - |  |  |
ForceNet | ✓ | 231.297 | - | - | - | 0.5
Tip3p | EGNN |  | 61.481 | 685.071 |  |  |
EGNN | ✓ | 39.291 | 241.853 | -36.09% | -64.70% | 0.9
GNS |  | 52.021 | 638.216 |  |  |
GNS | ✓ | 49.104 | 540.406 | -5.61% | -15.33% | 0.5
ForceNet |  | - | - |  |  |
ForceNet | ✓ | - | - | - | - | 0.5

Table 1: Quantitative force and energy performance analysis on RPBE and Tip3p
datasets. Performance improvements are calculated between pretrained models
and their trained-from-scratch counterparts. The units for force and energy
RMSE are ($meV/$\mathrm{\SIUnitSymbolAngstrom}$$) and ($meV$).

Following the described training setup, we conducted experiments to
investigate if masking pretraining would help the finetuning objective over
energy and force prediction accuracy on molecule potential datasets. All the
configurations are trained on four different seeds, and we report the averaged
results. As shown in Table 1, we compared the RMSE of predictions of force and
energy for our baseline models. The results show that our proposed masking
pretraining strategy provides a significant improvement over both objectives
across datasets. Most of the RMSE objectives are reduced by pretraining with
our proposed masking pretext task. The average improvement of force and energy
RMSE for EGNN and GNS is $38.09\%$ and $25.48\%$. Particularly, the pretrained
EGNN model 63 gained more than $35\%$ performance improvement on all
objectives. Masking pretraining also helps the GNS model 64 reduce $28.71\%$
force RMSE. This implies that our approach is model-agnostic (equivariant and
covariant models) and can be generalized to both force-centric models and
energy-centric models.

It is worth noting that the ForceNet model65 is hardly converging on both
datasets without pretraining. Since ForceNet’s results are several orders of
magnitude higher than other models’ results, we did not report out-of-scale
results for clarity. However, we observe that finetuning ForceNet becomes more
stable after pretraining in comparison to ForceNet trained from scratch. We
also observe a similar trend that masked pretraining can effectively stabilize
the training of other models. We argue that masked pretraining could
facilitate more stable training and faster convergence. We elaborate on this
point with training logs in the supplemental information section 6.2.

###  3.3 Comparison over models trained from scratch with additional epochs

Figure 2: Force and Energy performance comparison between models trained with
regular and extended epochs: 300 epochs and 400 epochs on the RPBE dataset,
and 50 epochs and 75 epochs on the Tip3p dataset.
$\text{Improvement}=1-\frac{{RMSE}_{\text{(pretrained + regular
epochs)}}}{{RMSE}_{\text{(train from scratch + regular/extended epochs)}}}$. A
negative percentage indicates better performance after pretraining.

Extended experiments are also conducted to help demonstrate the efficacy and
necessity of our two-stage training approach in comparison to simply extending
the training duration without pretraining. In contrast to trained from
scratch, pretrained models have yielded significant improvements on the RPBE
dataset, as shown in figure 2. In the first row of figure 2, We compared the
energy and force RMSE of our pretrained models to the models trained from
scratch for either regular 300 or an extended 400 epochs. The results show
that our training strategy of 100 epochs of pretraining followed by 300 epochs
of finetuning still outperforms the elongated train-from-scratch models.

The performance gap in force RMSE between EGNN and GNS models trained on the
RPBE dataset with and without pretraining sustains when compared to an
extended training duration. It underscores the effectiveness of pretraining
over merely increasing the training iterations. On Tip3p, the difference
between the pretrained models and the model trained from scratch with 75
epochs is reduced on force RMSE. Energy-wise, a similar trend is observed. It
is worth noting that the pretrained EGNN on the RPBE dataset achieves a higher
energy RMSE improvement when compared with EGNN trained from scratch with 400
epochs instead of 300 epochs. We attribute it to the overfitting of the model
to this dataset. This observation indicates that for models converging on
smaller datasets given certain epochs, starting from a point where the model
has learned meaningful molecule featurization (with pretraining), the network
is less likely to converge to worse local minima during the final training
phase.

Pretraining gives the models a better initialization and facilitates models to
extract meaningful representation from the unlabelled data, resulting in more
adept and competent models. Therefore, our findings suggests the integration
of pretraining as an essential component in building machine learning
potential training pipelines.

###  3.4 Comparison with Denoising

Figure 3:  Force and Energy performance comparison between models pretrained
with masking and denoising tasks

We provide the comparative performance metrics of our proposed masking and
denoising methods, which is critical in evaluating the efficacy of
computational models in predicting molecular properties. Figure 3 provided a
direct comparison between these two methods across two datasets, for force and
energy predictions. To create fair comparisons, we held masking pretrained
models’ results from table 1 constant. We implemented denoising as described
in Sec 2.3. The mean and standard deviation of the Gaussian Noise added are
set to be 0 and 0.5 throughout the experiments.

We observe that in the RPBE dataset, the performance of denoising is similar
and competitive to that of our method. We attribute this observation to the
effectiveness of the denoising approach in single molecule and attribute label
pairs datasets such as ISO17, and ANI-1x as demonstrated in recent research
57, 60. Containing on average 27 water molecules, data points in the RPBE
dataset are comparatively smaller than those in the Tip3p dataset.
Importantly, denoising has shown difficulties scaling to the Tip3p dataset.
EGNN model under the masking condition has a energy RMSE of about 241.85
$meV$. The performance deteriorates significantly when denoising is applied,
with poor convergence and RMSE rising to approximately 2442.75
$meV/$\mathrm{\SIUnitSymbolAngstrom}$$ by the end of the finetuning stage.
Denoising also exhibits degraded performance on the Energy RMSE metric when
pretraining using the GNS model on the Tip3p dataset. Both EGNN and GNS
pretrained using masked pretraining have shown stable convergence during
training without suffering from scalability and achieved superior performance.

##  4 Conclusion

We present a new method of pretraining for improving the performance of GNNs
on forcefield learning. We demonstrate that our method could achieve better
accuracy on both force and energy prediction. Moreover, it also facilitates
more stable training and exhibits better convergence in comparison to models
without pretraining or pretrained with other pretext tasks such as denoising.
These results indicate the effectiveness of our method for further utilization
of GNNs to facilitate the development of neural potentials. Scaling our method
with larger and more advanced GNNs, training on larger-scale datasets and
molecular systems, and extending our methods to auto-regressive models are
promising directions for future work.

##  5 Data Availability Statements

The data and code that support the findings of this project can be found at:
https://github.com/StevenZhang904/Masked_Pretrain.git.

## References

  * Harrison et al. 2018 Harrison, J. A.; Schall, J. D.; Maskey, S.; Mikulski, P. T.; Knippenberg, M. T.; Morrow, B. H. Review of force fields and intermolecular potentials used in atomistic computational materials research. _Applied Physics Reviews_ 2018, _5_ , 031104
  * Unke et al. 2021 Unke, O. T.; Chmiela, S.; Gastegger, M.; Schütt, K. T.; Sauceda, H. E.; Müller, K.-R. SpookyNet: Learning force fields with electronic degrees of freedom and nonlocal effects. _Nature Communications_ 2021, _12_ , 7273
  * Noé et al. 2020 Noé, F.; Tkatchenko, A.; Müller, K.-R.; Clementi, C. Machine Learning for Molecular Simulation. _Annual Review of Physical Chemistry_ 2020, _71_ , 361–390, PMID: 32092281
  * Gkeka et al. 2020 Gkeka, P.; Stoltz, G.; Barati Farimani, A.; Belkacemi, Z.; Ceriotti, M.; Chodera, J. D.; Dinner, A. R.; Ferguson, A. L.; Maillet, J.-B.; Minoux, H.; Peter, C.; Pietrucci, F.; Silveira, A.; Tkatchenko, A.; Trstanova, Z.; Wiewiora, R.; Lelièvre, T. Machine Learning Force Fields and Coarse-Grained Variables in Molecular Dynamics: Application to Materials and Biological Systems. _Journal of Chemical Theory and Computation_ 2020, _16_ , 4757–4775, PMID: 32559068
  * Unke et al. 2021 Unke, O. T.; Chmiela, S.; Sauceda, H. E.; Gastegger, M.; Poltavsky, I.; Schütt, K. T.; Tkatchenko, A.; Müller, K.-R. Machine Learning Force Fields. 2021
  * Botu et al. 2017 Botu, V.; Batra, R.; Chapman, J.; Ramprasad, R. Machine Learning Force Fields: Construction, Validation, and Outlook. _The Journal of Physical Chemistry C_ 2017, _121_ , 511–522
  * Li et al. 2017 Li, Y.; Li, H.; Pickard, F. C.; Narayanan, B.; Sen, F. G.; Chan, M. K. Y.; Sankaranarayanan, S. K. R. S.; Brooks, B. R.; Roux, B. Machine Learning Force Field Parameters from Ab Initio Data. _Journal of Chemical Theory and Computation_ 2017, _13_ , 4492–4503, PMID: 28800233
  * Chmiela et al. 2018 Chmiela, S.; Sauceda, H. E.; Müller, K.-R.; Tkatchenko, A. Towards exact molecular dynamics simulations with machine-learned force fields. _Nature Communications_ 2018, _9_ , 3887
  * Deringer et al. 2019 Deringer, V. L.; Caro, M. A.; Csányi, G. Machine Learning Interatomic Potentials as Emerging Tools for Materials Science. _Advanced Materials_ 2019, _31_ , 1902765
  * Behler 2016 Behler, J. Perspective: Machine learning potentials for atomistic simulations. _The Journal of Chemical Physics_ 2016, _145_ , 170901
  * Eshet et al. 2010 Eshet, H.; Khaliullin, R. Z.; Kühne, T. D.; Behler, J.; Parrinello, M. Ab initio quality neural-network potential for sodium. _Phys. Rev. B_ 2010, _81_ , 184107
  * Artrith et al. 2011 Artrith, N.; Morawietz, T.; Behler, J. High-dimensional neural-network potentials for multicomponent systems: Applications to zinc oxide. _Phys. Rev. B_ 2011, _83_ , 153101
  * Wang et al. 2018 Wang, H.; Zhang, L.; Han, J.; E, W. DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics. _Computer Physics Communications_ 2018, _228_ , 178–184
  * Artrith et al. 2017 Artrith, N.; Urban, A.; Ceder, G. Efficient and accurate machine-learning interpolation of atomic energies in compositions with many species. _Phys. Rev. B_ 2017, _96_ , 014112
  * Chmiela et al. 2017 Chmiela, S.; Tkatchenko, A.; Sauceda, H. E.; Poltavsky, I.; Schütt, K. T.; Müller, K.-R. Machine learning of accurate energy-conserving molecular force fields. _Science Advances_ 2017, _3_
  * Chmiela et al. 2018 Chmiela, S.; Sauceda, H. E.; Müller, K.-R.; Tkatchenko, A. Towards exact molecular dynamics simulations with machine-learned force fields. _Nature communications_ 2018, _9_ , 1–10
  * Bartók et al. 2010 Bartók, A. P.; Payne, M. C.; Kondor, R.; Csányi, G. Gaussian approximation potentials: The accuracy of quantum mechanics, without the electrons. _Physical review letters_ 2010, _104_ , 136403
  * Bartók et al. 2017 Bartók, A. P.; De, S.; Poelking, C.; Bernstein, N.; Kermode, J. R.; Csányi, G.; Ceriotti, M. Machine learning unifies the modeling of materials and molecules. _Science advances_ 2017, _3_ , e1701816
  * Bartók et al. 2013 Bartók, A. P.; Kondor, R.; Csányi, G. On representing chemical environments. _Physical Review B_ 2013, _87_ , 184115
  * Grisafi et al. 2018 Grisafi, A.; Wilkins, D. M.; Csányi, G.; Ceriotti, M. Symmetry-adapted machine learning for tensorial properties of atomistic systems. _Physical review letters_ 2018, _120_ , 036002
  * Bartók et al. 2017 Bartók, A. P.; De, S.; Poelking, C.; Bernstein, N.; Kermode, J. R.; Csányi, G.; Ceriotti, M. Machine learning unifies the modeling of materials and molecules. _Science Advances_ 2017, _3_
  * Smith et al. 2017 Smith, J. S.; Isayev, O.; Roitberg, A. E. ANI-1: an extensible neural network potential with DFT accuracy at force field computational cost. _Chem. Sci._ 2017, _8_ , 3192–3203
  * Christensen et al. 2020 Christensen, A. S.; Bratholm, L. A.; Faber, F. A.; Anatole von Lilienfeld, O. FCHL revisited: Faster and more accurate quantum machine learning. _The Journal of Chemical Physics_ 2020, _152_ , 044107
  * Zhang et al. 2018 Zhang, L.; Han, J.; Wang, H.; Car, R.; E, W. Deep Potential Molecular Dynamics: A Scalable Model with the Accuracy of Quantum Mechanics. _Phys. Rev. Lett._ 2018, _120_ , 143001
  * Carbogno et al. 2008 Carbogno, C.; Behler, J.; Groß, A.; Reuter, K. Fingerprints for Spin-Selection Rules in the Interaction Dynamics of ${\mathrm{O}}_{2}$ at Al(111). _Phys. Rev. Lett._ 2008, _101_ , 096104
  * Behler 2011 Behler, J. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. _The Journal of Chemical Physics_ 2011, _134_ , 074106
  * Behler and Parrinello 2007 Behler, J.; Parrinello, M. Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces. _Phys. Rev. Lett._ 2007, _98_ , 146401
  * Behler 2017 Behler, J. First Principles Neural Network Potentials for Reactive Simulations of Large Molecular and Condensed Systems. _Angewandte Chemie International Edition_ 2017, _56_ , 12828–12840
  * Duvenaud et al. 2015 Duvenaud, D. K.; Maclaurin, D.; Iparraguirre, J.; Bombarell, R.; Hirzel, T.; Aspuru-Guzik, A.; Adams, R. P. Convolutional Networks on Graphs for Learning Molecular Fingerprints. Advances in Neural Information Processing Systems. 2015
  * Kearnes et al. 2016 Kearnes, S.; McCloskey, K.; Berndl, M.; Pande, V.; Riley, P. Molecular graph convolutions: moving beyond fingerprints. _Journal of Computer-Aided Molecular Design_ 2016, _30_ , 595–608
  * Schütt et al. 2017 Schütt, K. T.; Arbabzadah, F.; Chmiela, S.; Müller, K. R.; Tkatchenko, A. Quantum-chemical insights from deep tensor neural networks. _Nature Communications_ 2017, _8_ , 13890
  * Unke and Meuwly 2019 Unke, O. T.; Meuwly, M. PhysNet: A Neural Network for Predicting Energies, Forces, Dipole Moments, and Partial Charges. _Journal of Chemical Theory and Computation_ 2019, _15_ , 3678–3693, PMID: 31042390
  * Schütt et al. 2017 Schütt, K. T.; Kindermans, P.-J.; Sauceda, H. E.; Chmiela, S.; Tkatchenko, A.; Müller, K.-R. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. 2017
  * Lubbers et al. 2018 Lubbers, N.; Smith, J. S.; Barros, K. Hierarchical modeling of molecular energies using a deep neural network. _The Journal of Chemical Physics_ 2018, _148_ , 241715
  * Imbalzano et al. 2018 Imbalzano, G.; Anelli, A.; Giofré, D.; Klees, S.; Behler, J.; Ceriotti, M. Automatic selection of atomic fingerprints and reference configurations for machine-learning potentials. _The Journal of Chemical Physics_ 2018, _148_ , 241730
  * Gilmer et al. 2017 Gilmer, J.; Schoenholz, S. S.; Riley, P. F.; Vinyals, O.; Dahl, G. E. Neural Message Passing for Quantum Chemistry. Proceedings of the 34th International Conference on Machine Learning. 2017; pp 1263–1272
  * Klicpera et al. 2020 Klicpera, J.; Groß, J.; Günnemann, S. Directional Message Passing for Molecular Graphs. International Conference on Learning Representations. 2020
  * Thomas et al. 2018 Thomas, N.; Smidt, T.; Kearnes, S.; Yang, L.; Li, L.; Kohlhoff, K.; Riley, P. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. _arXiv preprint arXiv:1802.08219_ 2018,
  * Fuchs et al. 2020 Fuchs, F.; Worrall, D.; Fischer, V.; Welling, M. Se (3)-transformers: 3d roto-translation equivariant attention networks. _Advances in Neural Information Processing Systems_ 2020, _33_ , 1970–1981
  * Anderson et al. 2019 Anderson, B.; Hy, T. S.; Kondor, R. Cormorant: Covariant molecular neural networks. _Advances in neural information processing systems_ 2019, _32_
  * Jing et al. 2021 Jing, B.; Eismann, S.; Suriana, P.; Townshend, R. J. L.; Dror, R. Learning from Protein Structure with Geometric Vector Perceptrons. International Conference on Learning Representations. 2021
  * Villar et al. 2021 Villar, S.; Hogg, D. W.; Storey-Fisher, K.; Yao, W.; Blum-Smith, B. Scalars are universal: Equivariant machine learning, structured like classical physics. Advances in Neural Information Processing Systems. 2021
  * Gasteiger et al. 2021 Gasteiger, J.; Becker, F.; Günnemann, S. Gemnet: Universal directional graph neural networks for molecules. _Advances in Neural Information Processing Systems_ 2021, _34_ , 6790–6802
  * Batzner et al. 2022 Batzner, S.; Musaelian, A.; Sun, L.; Geiger, M.; Mailoa, J. P.; Kornbluth, M.; Molinari, N.; Smidt, T. E.; Kozinsky, B. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. _Nature communications_ 2022, _13_ , 1–11
  * Hu* et al. 2020 Hu*, W.; Liu*, B.; Gomes, J.; Zitnik, M.; Liang, P.; Pande, V.; Leskovec, J. Strategies for Pre-training Graph Neural Networks. International Conference on Learning Representations. 2020
  * Rong et al. 2020 Rong, Y.; Bian, Y.; Xu, T.; Xie, W.; Wei, Y.; Huang, W.; Huang, J. Self-supervised graph transformer on large-scale molecular data. _Advances in Neural Information Processing Systems_ 2020, _33_ , 12559–12571
  * Zhang et al. 2021 Zhang, Z.; Liu, Q.; Wang, H.; Lu, C.; Lee, C.-K. Motif-based graph self-supervised learning for molecular property prediction. _Advances in Neural Information Processing Systems_ 2021, _34_ , 15870–15882
  * Wang et al. 2021 Wang, Y.; Wang, J.; Cao, Z.; Farimani, A. B. MolCLR: Molecular Contrastive Learning of Representations via Graph Neural Networks. 2021
  * Liu et al. 2022 Liu, Y.; Jin, M.; Pan, S.; Zhou, C.; Zheng, Y.; Xia, F.; Yu, P. Graph self-supervised learning: A survey. _IEEE Transactions on Knowledge and Data Engineering_ 2022,
  * Krishnan et al. 2022 Krishnan, R.; Rajpurkar, P.; Topol, E. J. Self-supervised learning in medicine and healthcare. _Nature Biomedical Engineering_ 2022, 1–7
  * Magar et al. 2022 Magar, R.; Wang, Y.; Barati Farimani, A. Crystal twins: self-supervised learning for crystalline material property prediction. _npj Computational Materials_ 2022, _8_ , 1–8
  * Cao et al. 2022 Cao, Z.; Magar, R.; Wang, Y.; Farimani, A. B. MOFormer: Self-Supervised Transformer model for Metal-Organic Framework Property Prediction. _arXiv preprint arXiv:2210.14188_ 2022,
  * Zhang et al. 2020 Zhang, S.; Hu, Z.; Subramonian, A.; Sun, Y. Motif-driven contrastive learning of graph representations. _arXiv preprint arXiv:2012.12533_ 2020,
  * Wang et al. 2022 Wang, Y.; Magar, R.; Liang, C.; Barati Farimani, A. Improving Molecular Contrastive Learning via Faulty Negative Mitigation and Decomposed Fragment Contrast. _Journal of Chemical Information and Modeling_ 2022,
  * Stärk et al. 2022 Stärk, H.; Beaini, D.; Corso, G.; Tossou, P.; Dallago, C.; Günnemann, S.; Liò, P. 3d infomax improves gnns for molecular property prediction. International Conference on Machine Learning. 2022; pp 20479–20502
  * Liu et al. 2022 Liu, S.; Wang, H.; Liu, W.; Lasenby, J.; Guo, H.; Tang, J. Pre-training Molecular Graph Representation with 3D Geometry. International Conference on Learning Representations. 2022
  * Zaidi et al. 2022 Zaidi, S.; Schaarschmidt, M.; Martens, J.; Kim, H.; Teh, Y. W.; Sanchez-Gonzalez, A.; Battaglia, P.; Pascanu, R.; Godwin, J. Pre-training via Denoising for Molecular Property Prediction. _arXiv preprint arXiv:2206.00133_ 2022,
  * Liu et al. 2022 Liu, S.; Guo, H.; Tang, J. Molecular geometry pretraining with se (3)-invariant denoising distance matching. _arXiv preprint arXiv:2206.13602_ 2022,
  * Zhou et al. 2023 Zhou, G.; Gao, Z.; Ding, Q.; Zheng, H.; Xu, H.; Wei, Z.; Zhang, L.; Ke, G. Uni-Mol: A Universal 3D Molecular Representation Learning Framework. The Eleventh International Conference on Learning Representations. 2023
  * Wang et al. 2023 Wang, Y.; Xu, C.; Li, Z.; Barati Farimani, A. Denoise Pretraining on Nonequilibrium Molecules for Accurate and Transferable Neural Potentials. _Journal of Chemical Theory and Computation_ 2023, _19_ , 5077–5087
  * Jiang et al. 2021 Jiang, D.; Zhenxing Wu, C.-Y. H.; Chen, G.; Liao, B.; Wang, Z.; Shen, C.; Cao, D.; Wu, J.; Hou, T. Could graph neural networks learn better molecular representation for drug discovery? A comparison study of descriptor-based and graph-based models. _Journal of Cheminformatics_ 2021,
  * Rong et al. 2020 Rong, Y.; Bian, Y.; Xu, T.; Xie, W.; Wei, Y.; Huang, W.; Huang, J. Self-Supervised Graph Transformer on Large-Scale Molecular Data. 2020,
  * Satorras et al. 2021 Satorras, V. G.; Hoogeboom, E.; Welling, M. E (n) equivariant graph neural networks. International conference on machine learning. 2021; pp 9323–9332
  * Kumar and Vantassel 2022 Kumar, K.; Vantassel, J. GNS: A generalizable Graph Neural Network-based simulator for particulate and fluid modeling. 2022,
  * Hu et al. 2021 Hu, W.; Shuaibi, M.; Das, A.; Goyal, S.; Sriram, A.; Leskovec, J.; Parikh, D.; Zitnick, C. L. ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations. 2021
  * Raghunathan Ramakrishnan and von Lilienfeld 2014 Raghunathan Ramakrishnan, M. R., Pavlo O. Dral; von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. _Scientific Data_ 2014,
  * Devlin et al. 2019 Devlin, J.; Chang, M.-W.; Lee, K.; Toutanova, K. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. 2019
  * Xie et al. 2022 Xie, T.; Fu, X.; Ganea, O.-E.; Barzilay, R.; Jaakkola, T. S. Crystal Diffusion Variational Autoencoder for Periodic Material Generation. International Conference on Learning Representations. 2022
  * Arts et al. 2023 Arts, M.; Satorras, V. G.; Huang, C.-W.; Zuegner, D.; Federici, M.; Clementi, C.; Noé, F.; Pinsler, R.; Berg, R. v. d. Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamics. _arXiv preprint arXiv:2302.00600_ 2023,
  * Vincent 2011 Vincent, P. A Connection Between Score Matching and Denoising Autoencoders. _Neural Computation_ 2011, _23_ , 1661–1674
  * Eastman et al. 2017 Eastman, P.; Swails, J.; Chodera, J. D.; McGibbon, R. T.; Zhao, Y.; Beauchamp, K. A.; Wang, L.-P.; Simmonett, A. C.; Harrigan, M. P.; Stern, C. D.; Wiewiora, R. P.; Brooks, B. R.; Pande, V. S. OpenMM 7: Rapid development of high performance algorithms for molecular dynamics. _PLOS Computational Biology_ 2017, _13_ , 1–17
  * Jorgensen et al. 1983 Jorgensen, W. L.; Chandrasekhar, J.; Madura, J. D.; Impey, R. W.; Klein, M. L. Comparison of simple potential functions for simulating liquid water. _The Journal of Chemical Physics_ 1983, _79_ , 926–935
  * Nosé 1984 Nosé, S. A unified formulation of the constant temperature molecular dynamics methods. _The Journal of Chemical Physics_ 1984, _81_ , 511–519
  * Hoover 1985 Hoover, W. G. Canonical dynamics: Equilibrium phase-space distributions. _Phys. Rev. A_ 1985, _31_ , 1695–1697
  * Cheng et al. 2019 Cheng, B.; Engel, E. A.; Behler, J.; Dellago, C.; Ceriotti, M. Ab initio thermodynamics of liquid and solid water. _Proceedings of the National Academy of Sciences_ 2019, _116_ , 1110–1115
  * Hammer et al. 1999 Hammer, B.; Hansen, L. B.; Nørskov, J. K. Improved adsorption energetics within density-functional theory using revised Perdew-Burke-Ernzerhof functionals. _Phys. Rev. B_ 1999, _59_ , 7413–7421
  * Grimme et al. 2010 Grimme, S.; Antony, J.; Ehrlich, S.; Krieg, H. A consistent and accurate ab initio parametrization of density functional dispersion correction (DFT-D) for the 94 elements H-Pu. _The Journal of Chemical Physics_ 2010, _132_ , 154104
  * Loshchilov and Hutter 2019 Loshchilov, I.; Hutter, F. Decoupled Weight Decay Regularization. 2019,

##  6 Supplemental Information

###  6.1 Choice of Masking Ratios

Throughout the experiments, we found that many factors directly influence
final performance, and among all, the choice of masking ratio is rather
important to analyze. Masking ratio controls the difficulty of pretraining
tasks during pretraining, thus determining how much pretrained models learn
about the water molecule systems and underlying dynamics. A higher masking
ratio for pretraining leaves the model with the harder task of recovering more
missing atoms. We conducted additional experiments on training EGNN with three
different masking ratios and compared the validation RMSE curve with the
trained-from-scratch ones on both datasets.

In general, a higher masking ratio on a more diverse and challenging dataset
like the RPBE dataset would make the model harder to learn good prior from it
with limited epochs of pretraining and lead to degraded results, as shown in
figure 4. A relatively lower masking ratio of 0.25 achieves better convergence
throughout the training duration. Meanwhile, when trained on the Tip3p dataset
as shown in figure 5 where the configuration is less diverse, the model can
restore the information with 90% of water molecules having masked hydrogen
atoms. In addition, when trained on the Tip3p dataset, all three choices of
masking ratio could achieve better performance, indicating the effectiveness
of our pretraining strategy. In practice, we suggest selecting a masking ratio
based on the size of the dataset to avoid negative transfer and justify the
trade-off between extra training time required by the pretraining stage and
performance gain.

Figure 4: Validation RMSE curve of EGNN model trained on RPBE dataset with
different masking ratios. The results of the masking ratio of 0.25 are
selected and reported in table 1, shown in the above figures as the bold
yellow line. Figure 5: Validation RMSE curve of EGNN model trained on Tip3p
dataset with different masking ratios. The results of the masking ratio of 0.9
are selected and reported in table 1, shown in the above figures as the bold
yellow line.

###  6.2 ForceNet Convergence Analysis

(a) Validation RMSE curve of ForceNet finetuned on RPBE dataset

(b) Validation RMSE curve of ForceNet finetuned on Tip3p dataset

Figure 6: Finetuning validation RMSE curve of ForceNet model trained with and
without pretraining

As shown in figure 6, a few epochs of masking pretraining could effectively
stabilize the training. Pretrained ForceNet models have smoother validation
RMSE curves in both datasets, while curves of trained-from-scratch
counterparts are more fluctuated and barely show signs of convergence.
Specifically, ForceNet trained from scratch on the RPBE dataset suffers from
gradient explosion on both energy and force objectives. It could probably be
attributed to RPBE being a smaller dataset. Small datasets usually would
require more epochs of training for them to learn effective representations
and are likely to suffer from gradient explosion and thus oscillate around the
local minimum. Despite ForceNet benefits much from pretraining on both
datasets, it may benefit more from pretraining and then finetuning in datasets
with more diverse configurations with enhanced training stability and faster
convergence.