SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rui Feng
rfeng@gatech.edu
Huan Tran
huan.tran@mse.gatech.edu
Aubrey Toland
artoland@gatech.edu
Qi Zhu
qi.zhu. ck@gmail.com
Binghong Chen
binghong@gatech.edu
Rampi Ramprasad
rampi. ramprasad@mse.gatech.edu
Chao Zhang
chaozhang@gatech.edu
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.
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, "Equivariant Transformer Optimization and Representations for Equilibrium and Off-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], energy-based 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.
- 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.
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 message-passing 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.
Consider a molecule type  $x$  that can exist in various 3D conformations. Let  $n_x$  represent the number of atoms in molecule  $x$ . The distribution of the molecule's conformations is represented as  $\mathbf{r}_x \in \mathbb{R}^{n_x \times 3}$ , and its atom types are represented as  $\mathbf{z}_x \in \mathbb{Z}^{n_x}$ . The potential energy of a 3D molecule is determined by its conformation  $\mathbf{r} \sim \mathbf{r}_x$  and atom types  $\mathbf{z}$ , denoted as  $E = E(\mathbf{r}, \mathbf{z}) \in \mathbb{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 = -\frac{\partial E}{\partial \mathbf{r}} \in \mathbb{R}^{n_x \times 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  $\Phi_{\theta}(\mathbf{r},\mathbf{z})$  parameterized by  $\theta$  that learns molecular representations from both equilibrium and off-equilibrium conformations. In this paper,  $\Phi_{\theta}(\mathbf{r},\mathbf{z})$  needs to predict energy-conservative atomic forces. To achieve this, we follow the standard machine learning forcefield paradigm. For simplicity, we omit  $\mathbf{z}$  in  $E$  and  $F$ , and we ignore the molecule type index  $x$  when referring to general molecules. We use  $\Phi : (\mathbb{R}^{n\times 3}) \times \mathbb{Z} \to \mathbb{R}$  to output the potential energy and take  $-\nabla \Phi_{\theta}$  as forces. To distinguish between equilibrium and off-equilibrium conformations, we define the equilibrium set of conformations as  $\mathcal{E} := \{\mathbf{r} : F(\mathbf{r}) = 0\}$  and the off-equilibrium set  $\mathcal{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 \sim \mathcal{S}$  with non-zero forces, we directly optimize on the atomic forces by minimizing  $\| -\nabla \Phi_{\theta}(\mathbf{r}_x) - F(\mathbf{r}_x)\|_2^2$ . For equilibrium molecular conformations  $x \sim \mathcal{E}$ , we assume their forces to be zero and impose a zero-force regularization:  $\| -\nabla \Phi_{\theta}(\mathbf{r}_x)\|_2^2$ . However, this objective gives the model little knowledge of the conformation structures in the neighborhood of  $\mathbf{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:
$$
\mathbb {E} _ {\varepsilon \sim \mathcal {N} (0, \sigma^ {2})} \left[ \| \nabla_ {\mathbf {r} _ {x}} \Phi (\mathbf {r} _ {x} - \varepsilon) - \varepsilon \| _ {2} ^ {2} \right], \quad x \sim \mathscr {E}. \tag {1}
$$
The rationale of the above force-guided denoising objective is as follows. The perturbed conformation, denoted as  $\mathbf{r}_x - \varepsilon$ , 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,  $\mathbf{r}_x$ , thereby yielding forces that are consistent with the direction  $\varepsilon$ . 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
$$
\mathbb {E} _ {x \sim \mathcal {E}} \left[ \underbrace {\| \nabla_ {\mathbf {r} _ {x}} \Phi (\mathbf {r} _ {x}) \| _ {2} ^ {2}} _ {\text {z e r o - f o r c e r e g u l a r i z a t i o n}} + \underbrace {\mathbb {E} _ {\varepsilon \sim \mathcal {N} (0 , \sigma^ {2})} \left[ \| \nabla_ {\mathbf {r} _ {x}} \Phi (\mathbf {r} _ {x} - \varepsilon) - \varepsilon \| _ {2} ^ {2} \right]} _ {\text {d e - n o i s i n g e q u i l i b r i u m}} \right] + \mathbb {E} _ {x \in \mathscr {S}} \left[ \underbrace {\| F (\mathbf {r} _ {x}) - \nabla_ {\mathbf {r} _ {x}} \Phi (\mathbf {r} _ {x}) \| _ {2} ^ {2}} _ {\text {f o r c e s o p t i m i z a t i o n}} \right], \tag {2}
$$
where the first expectation over  $x \sim \mathcal{E}$  samples equilibrium conformations and imposes a zero-force regularization and the denoising objectives on the atomic coordinates. The expectation over  $x \sim \mathcal{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
$$
\mathbb {E} _ {\varepsilon \sim \mathcal {N} (0, \sigma^ {2})} \left[ \left\| \nabla_ {\mathbf {r}} \Phi (\mathbf {r} - \varepsilon) - \varepsilon / \sigma^ {2} \right\| _ {2} ^ {2} \right] \tag {3}
$$
is equivalent to the score-matching-like objective
$$
\mathbb {E} _ {q _ {\sigma} (\tilde {\mathbf {r}})} \left[ \| \nabla_ {\tilde {\mathbf {r}}} \Phi (\tilde {\mathbf {r}}) - \nabla_ {\tilde {\mathbf {r}}} \log q _ {\sigma} (\tilde {\mathbf {r}}) \| _ {2} ^ {2} \right], \quad q _ {\sigma} (\tilde {\mathbf {r}} | \mathbf {r}) := \mathbf {r} + \varepsilon , \tag {4}
$$
which by [33] is equivalent to learning implicit forces around  $\mathbf{r}$ , assuming  $\mathbf{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 naively 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  $\Phi_{\theta}$ , 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  $1e - 4$ . The learning rate is scheduled with 10,000 warmup steps that gradually increase the learning rate from  $1e - 7$  to  $1e - 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  $\sigma^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 fine-tuning 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. 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.
Dataset	# Conformations	Equilibrium	Off-equilibrium
PCQM4Mv2	3,378,606	✓	✗
ANI1x	4,956,005	✓	✓
MD17	3,611,115	✗	✓
poly24	3,851,540	✗	✓
Total	15,718,279	✓	✓
Table 1: Datasets used in our model pre-training process.
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.1 MD Simulations for Small Molecules on MD17