SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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\mathring{\mathrm{A}})$ .
-  $h(r)$  MAE, with  $h(r)$  representing the distribution of interatomic distances during simulation. According to [8], MAE here is calculated as the  $l1$ -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
Dataset will be made available upon publication.
Molecule	Metric	DimeNet	GemNet-T	GemNet-dT	NequIP	TorchMDNet	ET-ORO	ET-OREO
Aspirin	Force (↓)	10.0	3.3	5.1	2.3	7.4	4.2	1.0
Stability (↑)	54(12)	72(50)	192(132)	300(0)	102(45)	94(42)	300(0)
h(r) (↓)	0.04(0.00)	0.04(0.02)	0.04(0.01)	0.02(0.00)	0.04(0.00)	0.04(0.00)	0.02(0.00)
Ethanol	Force	4.2	2.1	1.7	1.3	5.6	3.1	1.0
Stability	26(10)	169(98)	300(0)	300(0)	121(34)	300(0)	300(0)
h(r)	0.15(0.03)	0.10(0.02)	0.09(0.00)	0.08(0.00)	0.12(0.01)	0.10(0.00)	0.03(0.00)
Naphthalene	Force	5.7	1.5	1.9	1.1	3.3	2.0	0.9
Stability	85(68)	8(2)	25(10)	300(0)	50(20)	25(9)	300(0)
h(r)	0.10(0.01)	0.13(0.00)	0.12(0.01)	0.12(0.01)	0.12(0.00)	0.11(0.00)	0.03(0.00)
Salicylic Acid	Force	9.6	4.0	4.0	1.6	4.7	2.5	0.9
Stability	73(82)	26(24)	94(109)	300(0)	60(69)	94(58)	300(0)
h(r)	0.06(0.02)	0.08(0.04)	0.07(0.03)	0.03(0.00)	0.06(0.02)	0.05(0.01)	0.02(0.00)
Table 2: Simulation results on MD17. For all results, force MAE is reported in the unit of [meV/Å], and stability is reported in the unit of [ps]. The distribution of interatomic distances  $h(r)$  MAE is unitless. FPS stands for frames per second. For all metrics (↓) indicates the lower the better, and (↑) indicates the higher the better. The first group of methods is taken from [8]. The second group of methods is our new baselines, including TorchMDNet [8], ET-ORO, and ET-OREO. These models share the same architecture and have the same FPS.
(a) Interatomic distance distribution during simulations.
(b) Potential energy curve during simulations.
Figure 1: Illustration of simulation statistics for ET-OREO on MD17 dataset. Figure (a): distributions of interatomic distances during reference simulation and simulation with ET-OREO predictions finetuned on MD17. Figure (b): the curve of potential energy predicted by ET-OREO during simulation.
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  $\approx 100$  atoms. Therefore, ET-OREO is required to perform accurate MD simulations on unseen polymers, testing both its simulation and generalization ability.
ID	Polymer Name	# atoms	SMILES	Monomer	Polymer
CK	cyclooctane	360	[*]CCCCCC[∗]
OTH	n-alkane substituted δ-valerolactone	360	[∗]OC(CCC)CCCCCC([∗])=O
LAC1	γ-butyrolactone	180	[∗]OCCOCC([∗])=O
LAC2	3-methyl-1,4-dioxan-2-one	240	[∗]OCCOC(C)C([∗])=O
Table 3: Details of 15-loop polymers we used for simulations, including their ID, polymer name, number of atoms for the 15-loop polymer, SMILES string, and rdkit visualizations.
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\mathrm{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).
OTH
LAC1
LAC2
(a) Forces correlation of ET-OREO to DFT forces during simulations.
(b) Potential energy curve during ET-OREO simulations.
(c) Sampling of conformation spaces during ET-OREO simulations.
Figure 2: Accuracy and Robustness of ET-OREO in MD simulations for large polymers. In Figure (a) it is observed that ET-OREO forces are highly correlated with DFT forces during its own MD simulations. Figure (b) shows the RMSD and potential energy curves of ET-OREO during MD simulations. The curves suggest that ET-OREO simulations experience equilibration. Figure (c) shows the visualization of sample conformations captured at various time points during the simulation.
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  $\sim 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 off-equilibrium 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.
TorchMDNet	NoisyNode	ET-OREO
εHOMO	20.3	15.6	16.8
εLUMO	17.5	13.2	14.5
Δε	36.1	24.5	26.4
Table 4: Fine-tuning on HOMO-LUMO properties on QM9. Metrics are MAE in meV.
We presented ET-OREO, a pre-trained model for 3D molecules on  $\approx 15\mathrm{M}$  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.
[1] 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):1-11, 2022.
[2] Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-energy surfaces. Physical review letters, 98(14):146401, 2007.
[3] 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.
[4] Venkatesh Botu and Rampi Ramprasad. Learning scheme to predict atomic forces and accelerate materials simulations. Physical Review B, 92(9):094306, 2015.
[5] 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.
[6] 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.
[7] 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.
[8] 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.
[9] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d roto-translation equivariant attention networks. Advances in Neural Information Processing Systems, 33:1970–1981, 2020.
[10] 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.
[11] Johannes Gasteiger, Florian Becker, and Stephan Gunnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790-6802, 2021.
[12] Johannes Gasteiger, Janek Groß, and Stephan Günneumann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123, 2020.
[13] Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Velicković, James Kirkpatrick, and Peter Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2106.07971, 2021.