Title: May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular Conformations
Venue: NeurIPS 2023
Accepted: yes
Keywords: molecular pretraining

molecular representation learning
Abstract: Recent works have shown the promise of learning pre-trained models for 3D molecular representation.
However, existing pre-training models focus predominantly on equilibrium data and largely overlook off-equilibrium conformations.
It is challenging to extend these methods to off-equilibrium data because their training objective relies on assumptions of
conformations being the local energy minima. We address this gap by proposing a force-centric pretraining model for 3D molecular conformations covering both equilibrium and off-equilibrium data.
For off-equilibrium data, our model learns directly from their atomic forces. For equilibrium data, we introduce zero-force regularization and forced-based denoising techniques to approximate near-equilibrium forces.
We obtain a unified pre-trained model for 3D molecular representation with over 15 million diverse conformations. Experiments show that, with our pre-training objective, we increase forces accuracy by around 3 times compared to the un-pre-trained Equivariant Transformer model. By incorporating regularizations on equilibrium data, we solved the problem of unstable MD simulations in vanilla Equivariant Transformers, achieving state-of-the-art simulation performance with 2.45 times faster inference time than NequIP.  As a powerful molecular encoder, our pre-trained model achieves on-par performance with state-of-the-art property prediction tasks.

May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular Conformations

Abstract

Recent works have shown the promise of learning pre-trained models for 3D molecular representation. However, existing pre-training models focus predominantly on equilibrium data and largely overlook off-equilibrium conformations. It is challenging to extend these methods to off-equilibrium data because their training objective relies on assumptions of conformations being the local energy minima. We address this gap by proposing a force-centric pretraining model for 3D molecular conformations covering both equilibrium and off-equilibrium data. For off-equilibrium data, our model learns directly from their atomic forces. For equilibrium data, we introduce zero-force regularization and forced-based denoising techniques to approximate near-equilibrium forces. We obtain a unified pre-trained model for 3D molecular representation with over 15 million diverse conformations. Experiments show that, with our pre-training objective, we increase forces accuracy by around 3 times compared to the un-pre-trained Equivariant Transformer model. By incorporating regularizations on equilibrium data, we solved the problem of unstable MD simulations in vanilla Equivariant Transformers, achieving state-of-the-art simulation performance with 2.45 times faster inference time than NequIP. As a powerful molecular encoder, our pre-trained model achieves on-par performance with state-of-the-art property prediction tasks.

1 Introduction

Representation learning for 3D molecules is a crucial task in the field of AI for science. Recent years have seen significant advancements in Graph Neural Networks [24, 19, 12] and Equivariant Graph Neural Networks architectures [30, 9, 1, 11, 23] to capture interatomic features in 3D molecular 37th Conference on Neural Information Processing Systems (NeurIPS 2023).

conformations, thereby laying foundations for learning powerful representations. Leveraging such architectures, several existing works [33, 10, 17] have explored pre-training on 3D molecular structures on equilibrium data (PCQM4MV2 [16, 15]). Specifically, [33] developed a de-noising approach that learns meaningful information about the equilibrium conformations and their surroundings. [17, 10] leveraged the potential energy of equilibrium conformations to derive a supervised pretraining objective. Such pre-trained models have achieved promising performance on diverse molecular property prediction tasks.

Nonetheless, these existing 3D molecule pre-training models are based on the equilibrium assumption and predominantly overlook off-equilibrium conformations. In this paper, we consider off-equilibrium conformations to be those whose atomic forces significantly deviate from zero. In this case, existing work cannot be indiscriminately extended to non-zero forces, because their modeling assumption is fundamentally tied to equilibrium conformations being the local energy minima. The de-noising technique in [33], which is theoretically equivalent to learning forces around equilibrium conformations, may potentially be ill-posed for off-equilibrium conformations. Similarly, the energy-based supervised pretraining objectives in [17, 10] are based on the local minimality of these conformations and cannot be straightforwardly extended to off-equilibrium data. To date, the pre-training of a representation model for 3D molecular conformations that encompasses unifying both equilibrium and off-equilibrium molecular conformations remains largely underexplored. On the other hand, off-equilibrium conformations represent a significant portion of the chemical space, which is crucial for comprehensive modeling of the molecular representation space. For instance, applications such as molecular dynamics (MD) simulations are heavily dependent on the model's ability to accurately represent off-equilibrium data.

In this paper, we incorporate both off-equilibrium and equilibrium conformations into a unified representation learner. We propose a new pre-training model, ' E quivariant T ransformer O ptimization and R epresentations for E quilibrium and O ff-equilibrium Molecules' ( ET-OREO ) for 3D molecular conformations. This model integrates both equilibrium and off-equilibrium data from multiple large-scale datasets. Its training objective hinges on atomic forces , defined as the negative gradient of a molecule's potential energy with respect to atomic coordinates. Atomic forces exhibit several notable properties:(1) they are physically well-defined observable, i.e., the force acting on an atom is determined solely and uniquely from its local environment, defined as the real-space distribution of its neighboring atoms; (2) they are generalizable across various molecules in the sense that atoms from different molecules that have the same local environment should experience the same atomic forces; and (3) they can unify equilibrium and off-equilibrium data, as equilibrium data can be conceptualized as local minima in the latent (configuration) space with zero forces, while off-equilibrium data aid the model in more accurately characterizing the high-energy chemical space beyond equilibrium. Among these points, (2), which is the direct consequence of (1), makes atomic forces fundamentally different from potential energy which is defined for the whole system under consideration (molecules) and can only be determined up to an additional constant. Therefore, a predictive model of atomic forces is transferable, i.e., in principle, it can be used directly for molecules of any size (number of atoms). This advantage makes the 'learning atomic forces' fundamentally different from the traditional approach in which a fictitious concept of 'atomic energy' must be defined, predicted, and combined to obtain the total potential energy of the whole system [2].

Inspired by this, we develop our force-centric pretraining model that trains on both off-equilibrium and equilibrium molecular conformations. For off-equilibrium data, our model learns directly from their atomic forces, aligning the model gradient with respect to input coordinates with atomic forces. A conceived obstacle in leveraging forces lies in data acquisition: high-accuracy forces require the application of ab-initio methods such as Density Functional Theory (DFT) [14, 20]. However, we argue that DFT, despite its time complexity of O ( N 3 ) , remains tractable for moderately sized molecules [31]. Additionally, the major computational cost overlap with the potential energy, which is common in existing 3D conformation datasets.

For equilibrium data, we impose zero-force regularizations, reflecting their status as local minima of the potential energy surface. We also introduce random noise into equilibrium conformations and consider the random noise directions as approximate forces on the perturbed conformation coordinates. This allows the model to 'denoise' the perturbed conformations by gradients. Our approach draws inspiration from the recent success of denoising techniques in molecular learning [33, 13], energybased supervised pretraining [10, 17], and score-matching generative models [27, 28, 25]. We integrate our gradient-based denoising and zero-force regularization on equilibrium conformations

with off-equilibrium force optimization, thereby providing the model with a unified landscape of molecular structures and the potential energy surface.

For model pre-training, we collated more than 15 million 3D molecular conformations. Our pretraining data leverages three open-domain datasets, including PCQM4Mv2 [16, 15], MD17 [6, 5], ANI1-x [26]. Moreover, we contribute a new dataset, poly24 , consisting of simulation trajectories of a diverse family of polymers. Pretrained on these diverse sources of equilibrium and off-equilibrium data, our model attained state-of-the-art simulation performance on both MD17 and polymers, achieving efficient and accurate simulations in terms of distributional properties and accuracy relative to DFT calculations. Our model also serves as a potent representation learner for equilibrium data, attaining performance on par with state-of-the-art molecular learning methods that focus exclusively on equilibrium data.

In summary, our contributions are as follows:

We introduce a novel force-centric molecular conformation pretraining paradigm that trains a unified conformation representation encompassing both equilibrium and off-equilibrium molecules. Our paradigm enables the representation learner to portray not only the equilibrium conformation space but also the extensive off-equilibrium spaces between them.

Our model achieves highly accurate molecular dynamics (MD) simulations, by efficiently fine-tuning its parameters for use with molecules and polymers. This allows for fast and reliable MD simulation, as our model is able to accurately replicate the ab initio forces and distributional properties of conformations. Furthermore, our model has demonstrated comparable performance to state-of-the-art models that are solely focused on equilibrium data in molecular property prediction.

We provide the community with a diverse set of DFT simulation data comprising a varied set of polymers, which are valuable not only for studying polymer properties such as ring-opening enthalpies but also for the general modeling of molecular forces.

2 Related Work

Machine Learning Forcefield In recent years, there has been a surge in the development of deep learning models for molecule forcefields. Two prominent groups of models have emerged: geometric message passing models [24, 19, 12] and the Tensor Field Network (TFN) family [30, 9, 1, 11]. Geometric message-passing models [24, 19, 12] utilize traditional graph neural networks or messagepassing networks on pairwise radial features that are translational and rotational invariant. On the other hand, the TFN family models learn SE(3)-equivariant features by leveraging harmonic functions and the irreducible representations of the SE(3) group. NequIP [1] is the state-of-the-art model in this family and has achieved high stability and fidelity on the MD17 dataset.

Along another line, EGNN [23] learns equivariant features by integrating directional distance vectors between atom pairs in their implicit 3-dimensional coordinates. This approach allows for the learning of equivariant features without incurring the computational cost of the TFN family. Our method's backbone model is TorchMDNet [29], which we view as an extension of EGNN. TorchMDNet embeds the implicit 3-dimensional coordinates in a latent high-dimensional space, enhancing the model's capacity to represent 3-dimensional equivariant features. In addition, TorchMDNet leverages the concept of continuous distance filters from SchNet [24] to further increase its representation power.

Molecular Pretraining Inspired by the success of pre-training foundation models in the fields of NLP and CV, there have been several attempts to pretraining for 3D molecular structures [33, 10, 17]. Specifically, the NoisyNode [33] work adopted the denoising regularization technique of graph neural networks [13] as a pretraining method for equilibrium conformations, achieving state-of-the-art performance on molecular property predictions. [10] and [17] both base their pre-training on the supervised energy data from equilibrium conformations with forces regularization. However, as aforementioned, these existing methods predominantly rely on equilibrium data and cannot be easily extended to off-equilibrium data. While achieving high performance for property prediction tasks, they leave the vast chemical space of off-equilibrium conformations unexplored.

Force-based training over off-equilibrium molecular conformations has been explored in literature. [6, 5, 7, 4, 3] all focus on learning atomic forces and predict energy based on numeric integration, based on the claims that the potential energy and forces have different noise distribution pattern[6, 5] and forces being local and immediate atomic features [4, 3]. Their methods still focus on training on off-equilibrium simulation data from single molecules . In supplementary materials, we showed that the joint optimization of potential energy and forces is problematic for multi-molecule settings. We instead propose force-centered and energy-free objectives to learn a unified pretraining model for both off-equilibrium and equilibrium data; and we show that force-centered pre-training improves multi-molecule optimization and generalization.

3 Methodology

3.1 Problem Setup

Consider a molecule type x that can exist in various 3D conformations. Let nx represent the number of atoms in molecule x . The distribution of the molecule's conformations is represented as r x ∈ R nx × 3 , and its atom types are represented as z x ∈ Z nx . The potential energy of a 3D molecule is determined by its conformation r ∼ r x and atom types z , denoted as E = E ( r , z ) ∈ R . The forces applied to the atoms, which are defined as the negative gradient of the potential energy with respect to atomic coordinates, are given by F = -¶ E ¶ r ∈ R nx × 3 . In theory, a stable conformation is one where the potential energy achieves a local minimum with zero forces. In molecular dynamics (MD) simulations, the molecular conformations are moved according to forces and the thermostat of the simulation.

Our goal is to learn an equivariant model F q ( r , z ) parameterized by q that learns molecular representations from both equilibrium and off-equilibrium conformations. In this paper, F q ( r , z ) needs to predict energy-conservative atomic forces. To achieve this, we follow the standard machine learning forcefield paradigm. For simplicity, we omit z in E and F , and we ignore the molecule type index x when referring to general molecules. We use F : ( R n × 3 ) × Z → R to output the potential energy and take -GLYPH<209>F q as forces. To distinguish between equilibrium and off-equilibrium conformations, we define the equilibrium set of conformations as E : = { r : F ( r ) = 0 } and the off-equilibrium set S as the complementary set where the forces are non-zero.

3.2 Joint Training of Forces on Equilibrium and Off-equilibrium Data

Our methodology is centered on forces optimization. For off-equilibrium conformations x ∼ S with non-zero forces, we directly optimize on the atomic forces by minimizing ‖-GLYPH<209>F q ( r x ) -F ( r x ) ‖ 2 2 . For equilibrium molecular conformations x ∼ E , we assume their forces to be zero and impose a zero-force regularization: ‖-GLYPH<209>F q ( r x ) ‖ 2 2 . However, this objective gives the model little knowledge of the conformation structures in the neighborhood of r x . To better inform the model of the forcefield around equilibrium conformations, we further use a de-noising equilibrium regularization , where we add Gaussian noise on equilibrium conformations and use the noise direction as noisy forces:

The rationale of the above force-guided denoising objective is as follows. The perturbed conformation, denoted as r x -e , could either approximate a high-energy, off-equilibrium confirmation or, alternatively, represent an unphysical state. In the former scenario, the conformation is anticipated to relax back to the proximate local minimum, r x , thereby yielding forces that are consistent with the direction e . Conversely, in the latter situation, the learning model is still sufficiently equipped to maintain robust representations, thereby enabling a return to a stable conformation from any unphysical deviations. More formally, (1) can be interpreted as learning approximate forcefield around equilibrium conformations, as will be discussed in Section 3.3.

Combining forces optimization on off-equilibrium data and zero-force regularization and de-noising objective on equilibrium data, we have the unified force-centric pre-training objective

where the first expectation over x ∼ E samples equilibrium conformations and imposes a zero-force regularization and the denoising objectives on the atomic coordinates. The expectation over x ∼ S samples off-equilibrium conformations and optimizes the model gradient with forces.

3.3 On the De-noising Equilibrium Objective

Objective function (1) can be viewed as learning an approximate forcefield around equilibrium data. In fact, by well-known results in de-noising score-matching [32, 27, 28], the objective

is equivalent to the score-matching-like objective

which by [33] is equivalent to learning implicit forces around r , assuming r is an local minimizer of the energy. Hence, (1) is equivalent to learning forces around equilibrium conformations up to a scaling constant. This proof, however, cannot be naïvely extended to molecules whose forces are non-zero, which is the driving motivation for us to supplement the de-noising objective with forces from off-equilibrium conformations.

While [33] used the same argument for their de-noising objective, our implementation is different. They used a prediction head for predicting the noises on top of an Equivariant Transformer, while we directly predicted the noise with forces predicted by the model gradient. Our design has several unique advantages: 1) Forces principally govern atomic movements. By formulating the de-noising objective with predicted forces, our model has a physical interpretation where perturbed conformations, be it high-energy off-equilibrium conformations or unphysical conformations, could relax back to stable and physical conformations. This property is especially helpful in MD simulations, as will be shown in Section 4. 2) We can unify the prediction of the forces for both equilibrium and off-equilibrium data, providing the model with a consistent energy landscape across conformations of different molecules and states.

3.4 Model Pre-Training and Fine-Tuning

Model Architecture. We use TorchMDNet [29] (also known as Equivariant Transformer) as our molecule encoder F q , following [33]. TorchMDNet is one of the best-performing models in terms of predicting molecular properties and atomic forcefields. While [1] achieves higher accuracy in molecular prediction tasks, we find TorchMDNet more favorable for pre-training due to its expressivity and better computational efficiency.

Pre-training. Our model has 8 layers and 256 embedding sizes, 8 attention heads, and a radial basis function (RBF) dimension of 64, consistent with [29] and [33]'s best-performing model. The model parameters are optimized with the AdamW[22, 18] optimizer with a peak learning rate set as 1 e -4. The learning rate is scheduled with 10,000 warmup steps that gradually increase the learning rate from 1 e -7 to 1 e -4, and afterward will decrease with a multiplier of 0.8 after 30 validation epochs of non-decreasing validation loss. On every 500 training steps, one validation epoch is performed. The model is trained with a batch size of 32 samples for 3 epochs for 468,750 gradient steps. For the denoising objective, the variance of noise s 2 is set to be 0.04.

Model Fine-tuning. The pre-training stage provides the model with a good initial representation of molecules in general but is not necessarily optimal for each specific dataset or task. For our experiments, the model is further fine-tuned on the target dataset to optimize task-specific performance. For ET-OREO, the training/testing split of overlapping datasets in both the pre-training and finetuning stages are consistent so that during fine-tuning, the test set will not include any data seen during training.

3.5 Pre-Training Data

Our paper focuses primarily on the structures of organic molecules and polymers in a vacuum. This approach enables us to maintain consistent learning of quantum mechanical interactions between

atoms within a uniform environment. For this purpose, we use three public 3D molecular conformation datasets MD17, ANI1-x, and PCQM4Mv2; and we also create a new polymer simulation dataset, detailed as follows.

Poly24: MD Simulations for Polymers. We contribute poly24 , a DFT-based MD simulation dataset for polymers. 1 We generated DFT simulation data for 24 types of polymer families, broadly categorized into cycloalkanes, lactones, ethers, and others. Each polymer family consists of a cyclic monomer and its ring-opening polymerizations. Within ring-opening polymerization, a ring of the cyclic monomers is broken and then the 'opened' monomer is added to a long chain, forming a polymer chain. The details on data generation can be found in the Appendix. In total, we run generally 10 DFT simulations for different initialization of each L -loop ( L = 1 , 3 , 4 , 5 , 6) polymer across the 24 types of polymers. We have 1311 DFT trajectories and 6,552,624 molecular conformations. Only polymers with less than or equal to 64 atoms are used for pre-training, totaling 3,851,540 conformations. The remaining larger-polymer conformations are used for fine-tuning and testing.

MD17, ANI1-x, and PCQM4Mv2. In addition to our own poly24 , we have utilized three existing public datasets, namely MD17[6, 5], ANI1-x[26], and PCQM4Mv2[16, 15] for our model pre-training.

These datasets contain small organic molecules in a vacuum, and property prediction for such molecules is an area of great interest to the cheminformatics community. The machine learning for molecules community has also extensively studied and benchmarked these datasets. Table 1 summarizes all the dataset used in the pre-training stage. In total, we have more than 15 million samples from MD17, ANI1-x, PCQM4Mv2, and Poly24, covering equilibrium and off-equilibrium conformations for diverse organic molecules.

4 Experiments

4.1 MDSimulations for Small Molecules on MD17

Setup After pre-training our model ET-OREO, we further finetune it on the MD17 dataset to validate the performance of molecular simulations using the forces predicted by our fine-tuned model. We run simulations with our model's predicted forces with a Nose-Hoover thermostat, with 500K temperature and 0.5fs timestep for 600k steps. The simulation setting is the same as [8], which benchmarked popular deep learning forcefields for MD simulations. We use the ASE package for the implementation of molecular simulation environment [21] with a characteristic time for the Nose-Hoover thermostat set to 25fs.

Baselines and Metrics. Following [8], we report 3 metrics:

Forces MAE , which is the mean absolute error of forces on the DFT trajectories;

Stability , which measures how long the simulation can run before blowing up. According to [8], this is detected when the radial distance function (RDF) deviates from the reference simulation by a threshold (0.10Å).

h ( r ) MAE , with h ( r ) representing the distribution of interatomic distances during simulation. According to [8], MAE here is calculated as the l 1-norm between the reference distribution and the predicted distribution.

We compare ET-OREO with the best-performing models reported by [8]. Furthermore, ET-OREO is compared with TorchMDNet [29] trained on MD17 from scratch, and ET-ORE, which is an ablation version of ET-OREO without the zero-force and de-noising regularizations. Consistent with previous

Aspirin

Ethanol

Naphthalene

Salicylic Acid

0.0

2.5

5.0

7.5

10.0

Interatomic Distances r (

Å

)

0.0

0.1

0.2

0.3

0.4

0.5

Density

h(r): 0.02

Reference

Prediction

0.0

2.5

5.0

7.5

10.0

Interatomic Distances r (

Å

)

0.0

0.5

1.0

1.5

Density

h(r): 0.09

Reference

Prediction

0.0

2.5

5.0

7.5

10.0

Interatomic Distances r (

Å

)

0.0

0.2

0.4

0.6

0.8

Density

h(r): 0.08

Reference

Prediction

0.0

2.5

5.0

7.5

10.0

Interatomic Distances r (

Å

)

0.0

0.2

0.4

0.6

0.8

Density

h(r): 0.02

Reference

Prediction

(a) Interatomic distance distribution during simulations.

0

100

200

300

Time [ps]

-0.83

-0.83

-0.83

-0.83

-0.83

Energy [meV]

0

100

200

300

Time [ps]

-0.83

-0.83

-0.83

-0.83

-0.83

Energy [meV]

0

100

200

300

Time [ps]

-0.83

-0.83

-0.83

-0.83

-0.83

Energy [meV]

0

100

200

300

Time [ps]

-0.83

-0.83

-0.83

-0.83

-0.83

Energy [meV]

(b) Potential energy curve during simulations.

simulation benchmark [8], TorchMDNet, ET-ORE, and ET-OREO are fine-tuned with 9,500 conformations for training and 500 for validation. The training and validation conformations are sampled from the same training data used in the pre-training of ET-OREO, and the metrics are reported on the rest of MD17 data unseen to both the fine-tuning and the pre-training of ET-OREO and ET-ORE. Furthermore, in our implementation, we focus solely on forces optimization for TorchMDNet. We show in Appendix that this improves forces accuracy without sacrificing the model's ability to predict the potential energy.

Results. In Table 2, we compare the performance of ET-OREO with baseline models. We make the following key observations: (1) On all of the simulations, we have achieved a lower interatomic distance distribution h ( r ) than the previous best-performance model NequIP reported by [8]. Particularly, we reduced h ( r ) for aspirin and salicylic acid by half, 10% for ethanol, and 18% for naphthalene. While NequIP [1] can achieve stable and accurate MD simulations by training from scratch on MD17, it suffers from a lower FPS as it is computationally expensive. In contrast, ET-OREO can achieve both fast and high-quality MD simulations. (2) ET-OREO improves forces accuracy over TorchMDNet by over three times, achieving state-of-the-art MAE on all four tested molecules. The major difference between ET-OREO and TorchMDNet is that ET-OREO is pre-trained with our force-centric objective before fine-tuning on MD17 data. Such significant performance gains show the benefits of pre-training over diverse equilibrium and off-equilibrium conformations. (3) ET-OREO produces

stable simulations and accurate interatomic distance distributions. ET-OREO is able to stably run simulations on all four molecules with accurate interatomic distance distribution h ( r ) compared to the reference trajectory in MD17. Figure 1 visualizes the close approximation of ET-OREO predicted h ( r ) compared with reference data, and that ET-OREO can produce energy-conserving simulations that sample around the equilibrium.

Regularization on Equilibrium Conformations is Vital for Simulations. We found that training TorchMDNet itself on MD17 cannot produce stable MD simulations. Meanwhile, the performance of ET-ORE shows that pre-training with forces only helped with higher forces accuracy. However, ET-ORE still does not perform well for MD17 simulations. Except for ethanol, ET-ORE cannot produce stable MD simulations, despite consistently improved forces accuracy. This is in line with [8]'s observation that higher forces accuracy does not always guarantee simulation performance. The difference between ET-ORE and ET-OREO is that the latter incorporates stable zero-force conformations and a de-noising regularization objective. The additional regularization on the conformation space proves vital to stable and accurate simulations.

ET-OREO achieves accurate MD simulation with 3 times inference speed. From Table 2, NequIP reported similar simulation performance to ET-OREO without need for pre-training. However, NequIP has a significantly slower inference time. The NequIP model in Table 2 has an inference time of approximately 119ms per step on NVIDIA V100 with 1.05 million model parameters. [8] In comparison, ET-OREO has an inference time of 48.5ms on the same hardware with 6.87 million parameters. Hence, we have achieved state-of-the-art simulation performance with a 2.45 faster inference speed.

4.2 Simulation on Large-loop Polymers

The vast majority of our model's training data is small molecules. In this section, we investigate the model's ability to generalize to larger molecules, consisting of 15-loop polymers unseen in the training data.

Setup We finetune our model on the forces of small polymers (5-loop or less) in the poly24 dataset for one epoch with a learning rate of 0 . 0001. For testing the model's MD simulation performance on large polymers, we take the larger 15-loop polymers from the poly24 dataset and run MD simulations on them with ET-OREO-fine-tuned forcefields for a maximum of 600K steps. The number of atoms for these 15-loop polymers are, respectively: 360, 360, 180, and 240. In the training data, we have at most molecules consisting of ≈ 100 atoms. Therefore, ET-OREO is required to perform accurate MDsimulations on unseen polymers, testing both its simulation and generalization ability.

ID

Polymer Name

# atoms

SMILES

Monomer

Polymer

CK

cyclooctane

360

[*]CCCCCCCC[*]

OTH

n-alkane substituted

d

-valerolactone

360

[*]OC(CCC)CCCC([*])=O

g

-butyrolactone

180

[*]OCCOCC([*])=O

LAC2

3-methyl-1,4-dioxan-2-one

240

[*]OCCOC(C)C([*])=O

Results. Figure 2 (a) visualize the comparison results between ET-OREO predicted forces and DFT reference data. For all 15 loop polymers, ET-OREO obtains highly accurate forces to DFT calculations, with 0.01 eV/ MAE in energy and close to zero cosine distance. Hence, ET-OREO can perform MD simulations on multiple types of polymers with uniformly high correlation with DFT references. The test 15-loop polymers contain up to 360 atoms, unseen in the training data. This shows that ET-OREO can extrapolate well to large unseen polymers with known monomers. Furthermore, in practice, ET-OREO only needs training data from small polymers, which are cheap to generate ab initio data and thus greatly reduce the cost of DFT simulation.

Figure 2 (b) shows the potential energy curves during simulation. The potential energy curves indicate that the polymer in the simulation converges to a stable near-equilibrium distribution quickly after

initial fluctuations. Furthermore, ET-OREO can explore different possible conformation states of the polymers, as illustrated in Figure 2 (c).

CK

OTH

LAC2

(a) Forces correlation of ET-OREO to DFT forces during simulations.

0

200

400

Time [ps]

-45.0

-40.0

-35.0

-30.0

-25.0

Energy [eV]

0

200

400

Time [ps]

-45.0

-40.0

-35.0

Energy [eV]

0

200

400

Time [ps]

-28.0

-26.0

-24.0

-22.0

Energy [eV]

0

200

400

Time [ps]

-32.5

-30.0

-27.5

-25.0

Energy [eV]

(b) Potential energy curve during ET-OREO simulations.

Initial

150ps

300ps

Initial

150ps

300ps

Initial

150ps

300ps

Initial

150ps

300ps

Predicted Force (eV/A)

6

(eV/A)

6

Predicted Force (eV/A)

6

Predicted Force (eV/A)

6

4

4

4

4

Force

2

2

2

2

0

0.1

0.0

10

Predicted

0

0.1

0.0

0.1

0

0.1

0.0

T0

O

T'0-

0.0

0.1

2

一

-2

2

-2

-4

MAE:4.2meV/A

CosD:0.0000

-4

CosD:0.0001

MAE:8.4meV/A

-4

MAE:.10.3meV/A

CosD:0.0002

-4

MAE:11.7meV/A

CosD:0.0002

-6

-5.0-2.50.02.55.0

-6

-5.0-2.5 0.02.55.0

-6

-5.0 -2.50.02.55.0

-6

-5.0-2.50.0

2.55.0

DFTForce (eV/A)

DFT Force (eV/A)

DFT Force (eV/A)

DFT Force (eV/A)

4.3 Property Prediction on QM9

To test ET-OREO's ability to encode equilibrium conformation for property prediction tasks, we follow the same experiment setting as NoisyNode [33] on QM9. [33] exclusively pre-trained on equilibrium conformations from PCQM4Mv2 with a de-noising objective. In comparison, ET-OREO has access to

larger off-equilibrium data supervised by forces. We follow the exact same fine-tuning setup as [33] save the de-noising regularization in the fine-tuning stage. In Table 4, we compare ET-OREO with NoisyNode and TorchMDNet [29] on HOMO-LUMO properties prediction on QM9. We report the performance of NoisyNode trained on the TorchMDNet encoder reported in [33], hence all three models have the same encoders, except that TorchMDNet trains from scratch, and NoisyNode and ET-OREO fine-tune from their pre-trained parameters. ET-OREO improves the performance of property prediction by ∼ 30% compared to TorchMDNet trained from scratch, implying that our pre-training paradigm provided the model with useful information about the quantum mechanical properties of equilibrium conformations. Our performance is on par with NoisyNode, and unlike in simulation experiments, we observe no accuracy improvement by the inclusion of the rich offequilibrium conformation. We speculate that such equilibrium property prediction tasks benefit little from additional information on off-equilibrium conformation. We want to particularly point out that, despite the pre-training effort, NoisyNode and ET-OREO both require extensive fine-tuning with a large amount of data (more than 10,000) and training epochs (roughly 150 epochs for convergence) for optimal performance. An important future improvement in molecular pre-training should be in improving the data and time efficiency of fine-tuning.

5 Discussion

We presented ET-OREO, a pre-trained model for 3D molecules on ≈ 15M 3D conformations of both equilibrium and off-equilibrium states. We unified the pre-training of molecular conformations of different sources and states with a force-centric pre-training objective. With our pre-training model, we achieved state-of-the-art performance on MD simulations, with both high force accuracy and simulation efficiency. As a potent encoder for conformations, our model also attained a state-of-the-art level of performance on property prediction on equilibrium data.

Our current model is limited mostly to single and small molecules in a vacuum, leaving more complex molecular environments as future work. We also deem it a promising direction to explore more model architectures and larger model sizes. Furthermore, our de-noising objective on equilibrium data does not leverage the model's learned information on atomic forces. It will be interesting to study the possibility of a more sophisticated coupling of equilibrium and off-equilibrium optimization, to make the model progressively leverage its quantum mechanical knowledge for equilibrium data.

References

Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications , 13(1):111, 2022.

Jörg Behler and Michele Parrinello. Generalized neural-network representation of highdimensional potential-energy surfaces. Physical review letters , 98(14):146401, 2007.

Venkatesh Botu, Rohit Batra, James Chapman, and Rampi Ramprasad. Machine learning force fields: construction, validation, and outlook. The Journal of Physical Chemistry C , 121(1):511-522, 2017.

Venkatesh Botu and Rampi Ramprasad. Learning scheme to predict atomic forces and accelerate materials simulations. Physical Review B , 92(9):094306, 2015.

Stefan Chmiela, Huziel E Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learned force fields. Nature communications , 9(1):3887, 2018.

Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances , 3(5):e1603015, 2017.

Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T Unke, Adil Kabylda, Huziel E Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. arXiv preprint arXiv:2209.14865 , 2022.

Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Gomez-Bombarelli, and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. arXiv preprint arXiv:2210.07237 , 2022.

Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d rototranslation equivariant attention networks. Advances in Neural Information Processing Systems , 33:1970-1981, 2020.

Xiang Gao, Weihao Gao, Wenzhi Xiao, Zhirui Wang, Chong Wang, and Liang Xiang. Supervised pretraining for molecular force fields and properties prediction. arXiv preprint arXiv:2211.14429 , 2022.

Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems , 34:6790-6802, 2021.

Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123 , 2020.

Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veliˇ ckovi´ c, James Kirkpatrick, and Peter Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2106.07971 , 2021.

P. Hohenberg and W Kohn. Inhomogeneous electron gas. Phys. Rev. , 136:B864-B871, 1964.

Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogblsc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.09430 , 2021.

Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems , 33:22118-22133, 2020.

Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. Energy-motivated equivariant pretraining for 3d molecular graphs. arXiv preprint arXiv:2207.08824 , 2022.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123 , 2020.

W. Kohn and L Sham. Self-consistent equations including exchange and correlation effects. Phys. Rev. , 140:A1133-A1138, 1965.

Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E Castelli, Rune Christensen, Marcin Dułak, Jesper Friis, Michael N Groves, Bjørk Hammer, Cory Hargus, et al. The atomic simulation environment-a python library for working with atoms. Journal of Physics: Condensed Matter , 29(27):273002, 2017.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning , pages 9323-9332. PMLR, 2021.

Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics , 148(24):241722, 2018.

Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning , pages 9558-9568. PMLR, 2021.

Justin S Smith, Roman Zubatyuk, Benjamin Nebgen, Nicholas Lubbers, Kipton Barros, Adrian E Roitberg, Olexandr Isayev, and Sergei Tretiak. The ani-1ccx and ani-1x data sets, coupled-cluster and density functional theory properties for molecules. Scientific data , 7(1):134, 2020.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Philipp Thölke and Gianni De Fabritiis. Torchmd-net: Equivariant transformers for neural network based molecular potentials. arXiv preprint arXiv:2202.02541 , 2022.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219 , 2018.

Oliver T Unke, Stefan Chmiela, Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schütt, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine learning force fields. Chemical Reviews , 121(16):10142-10186, 2021.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.00133 , 2022.

Peer reviews:
comments: SUMMARY:
1. proposing a pretraining framework for conformation space including equilibrium and off-equilibrium.
2. this method also shows highly accurate in MD or relax, including molecules and polymers.
3. providing the community with a diverse set of DFT dataset of polymers.

STRENGTHS:
1. This paper is well-written and well-organized, effectively presenting the significance of force-centric pretraining for molecular systems. The proposed framework demonstrates comprehensive and impressive results in molecular dynamics simulations of small molecules and polymers, as well as in quantum chemistry property predictions.
2. The authors also introduce an innovative concept in their work: the de-noising equilibrium balance and off-equilibrium co-training method.

WEAKNESSES:
1. As previously mentioned, I am particularly intrigued by the efficiency of the de-noising equilibrium tasks. It would be better to conduct an detail ablation study of de-noise equillibrium and force regularization.

QUESTIONS:
1. Would it be possible to incorporate the PCQM4Mv2 homo-lumo prediction results in the downstream analysis? This benchmark is well-defined and provides a competitive comparison for evaluating the performance of various models.

2. About code and data access, I believe that the Poly24 dataset is of significant value to the research community for future investigations.

LIMITATIONS:
In the discussion section, it is mentioned that both the backbone and time efficiency play crucial roles in the performance of the proposed AI model. 
Besides I recommend incorporating partial charge prediction during the pretraining or downstream stages of your framework, if feasible. This addition could potentially enhance the overall effectiveness and applicability of your model.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper incorporates off-equilibrium molecular conformations into consideration, in which the energy is not minimized. This paper defines pre-training objectives for force field learning on sampled off-equilibrium molecular conformations. Experimental results validate the effectiveness of pre-training on additional off-equilibrium data.\n\nSTRENGTHS:\n(1) The presentation of the paper is quite clear and easy to follow. The writing is well-organized;\n\n(2) The proposed idea is quite interesting and necessary. How to utilize off-equilibrium data is a useful and important research question for pre-training 3D molecular property prediction. This reviewer agrees with the significance of the proposed idea.\n\nWEAKNESSES:\n(1) The novelty of the proposed method is not very high. Considering there are already published papers researching denoised pretraining strategies on 3d molecules, which lowers the novelty of the proposed method a bit. Then the major contribution of the proposed method is mainly about introducing additional data. This reviewer suggests that this paper might be more suitable for a domain journal; \n\n(2) The contribution of this paper might not be enough. Currently, the major contribution includes introducing additional data, and additional pre-training objectives. There are some important works like SCN (Spherical channel networks for modeling atomic interactions) utilizing separate force predictions. So separately predicting forces cannot be regarded as an important contribution;\n\n(3) The experimental results do not demonstrate the effectiveness of the proposed method. It seems that the proposed method seldom outperforms previous baseline methods, which makes this reviewer feel a little bit confused.\n\nQUESTIONS:\nIt seems the reported results on MD17 are not that superior. Probably this reviewer misunderstands the experimental tables but it seems the sota method is usually not the proposed method. Does the proposed method outperform the previous baseline methods?\n\nLIMITATIONS:\nThe major limitation of the proposed method is that the force prediction might only play a role as a regularizer since learning force field is also a very challenging task since force is a vector feature instead of a scalar feature like energy.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper propose a set of additional objectives for training machine learning interatomic potentials, in particular for equilibrum conformations as well as a denoising term. They show that this improves stability, when pretrained on a large dataset. They also introduce a new dataset of polymers.\n\nSTRENGTHS:\nThe improvement from ET-ORE to ET-OREO is interesting and surprising. It doesn't really make sense that equilibrum forces should help The polymer data set is great!\n\nWEAKNESSES:\nProbably the biggest weakness is that the authors chose a weak baseline model with TorchMDNet. The entire paper hinges on work from Fu et al. that showed that in their experiments all methods except the NequIP potential were unstable in MD simulation. With many additional tricks and an enormous amount of data they now match NequIP. But will this entire improvement just go away when you apply the ideas to a stronger baseline model like NequIP? I would strongly suggest the authors to repeat the experiments with a better baseline and see if there are still improvements. Otherwise why wouldn't anyone just use NequIP to begin with? The speedup is small and as explained below mostly an artefact of hyperparameters.\n\nFinally, a central weakness it that ET-OREO is pretrained on a massive dataset while the baseline methods that are stable can do this from a much much much smaller dataset. \n\nAnother big weakness is that this unexplained improvement from ET-ORE to ET-OREO simultaneously adds the equilibrum and the denoising objective. Why no ablation study to understand which of them is the important one?\n\nThe inference speed comparison to NequIP is not thorough. NequIP has shallow pareto curves on accuracy vs efficiency which is well known, i.e. a L=2 models gets almost exactly the same accuracy as L=3, but is much faster (see original NequIP paper). The paper compares to a NequIP model that is solely optimized for accuracy, not for efficienc. This would be helpful to compare to. In addition, NequIP modern versions of NequIP now exists that are faster, e.g. Allegro or MACE. \n\nThe presentation in addition has several obvious, outlined item-by-item below: \n\n- \"they are locally well-defined for atoms and generalizable across various molecules\" --> this statement is clearly not true in this generality. In what way is a force generalizable across molecules? Needs clarification what is meant here? \n\n- following immediately: \"they reflect quantum mechanical interactions among atoms and expose the underlying quantum structures of molecules;\" --> extremely generic, if even true, atoms should be electrons here... \n\n- \"By harnessing these factors, we posit that learning from atomic forces can\n57 guide the model to learn a comprehensive and unified landscape of molecular conformations across\n58 different molecules, datasets, and equilibrium states\" --> this has been explored for decades in every single paper training interatomic potentials, leveraging atomic forces is in now way a novel idea, this needs a strong rephrasing. \n\n- \"We introduce a novel force-centric molecular conformation pretraining paradigm that trains\n87 a unified conformation representation encompassing both equilibrium and off-equilibrium\n88 molecules\" this is a vast overstatement, see above\n\n- 'The above-mentioned papers followed the same learning paradigm, where different models are\n116 trained for different molecules. Consequently, these methods cannot be easily generalized to unseen\n117 molecules\" --> there have been many efforts at building general-purpose force-fields, see e.g. the ANI series. \n\n- \". Despite their impressive accuracy in static prediction, their ability\n119 of running simulations were not reported until [26] showed that the increased static accuracy does not\n120 translate to better simulation performance\" again, horrendously false statement, many of the MLIP papers you cite ran various simulations, across diverse and difficult phenoma from ionic diffusion to heterogeneous catalysis to small molecules. \n\n- \"In molecular dynamics (MD) simulations, the molecular\n149 conformations are moved according to forces and the thermostat of the simulation.\" what about NVE simulations that don't have a thermostat? \n\n- \"For equilibrium molecular conformations x ∼ E , we assume their forces to be zero and impose a\nzero-force regularization:\" --> what is the reasoning behind this, if you *have* the forces anyway, they will be zero at equilibrium, so you're just again training on forces, again what this community has been doing since the 60s. \n\n- the denoising objective seems like it has obvious failure cases: perturbing a relaxed structure and argueing the force must go into the direction of the GS completely ignores the many-bodyness of the PES, discussing these would be helpful\n\n- a lot of the work seems to be based on improving one or two specific previously introduce pretraining objectives for equilibrium data. In the opening of the manuscript, these are cited several times, but never explained. This makes the manuscript difficult to read.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-08"}